weave-scope/report/counters.go

package report

import (
	"bytes"
	"fmt"
	"reflect"

	"github.com/ugorji/go/codec"
	"github.com/weaveworks/ps"
)

// Counters is a string->int map.
type Counters struct {
	psMap ps.Map
}

var emptyCounters = Counters{ps.NewMap()}

// MakeCounters returns EmptyCounters
func MakeCounters() Counters {
	return emptyCounters
}

// Copy is a noop
func (c Counters) Copy() Counters {
	return c
}

// Add value to the counter 'key'
func (c Counters) Add(key string, value int) Counters {
	if c.psMap == nil {
		c = emptyCounters
	}
	if existingValue, ok := c.psMap.Lookup(key); ok {
		value += existingValue.(int)
	}
	return Counters{
		c.psMap.Set(key, value),
	}
}

// Lookup the counter 'key'
func (c Counters) Lookup(key string) (int, bool) {
	if c.psMap != nil {
		existingValue, ok := c.psMap.Lookup(key)
		if ok {
			return existingValue.(int), true
		}
	}
	return 0, false
}

// Size returns the number of counters
func (c Counters) Size() int {
	if c.psMap == nil {
		return 0
	}
	return c.psMap.Size()
}

// Merge produces a fresh Counters, container the keys from both inputs. When
// both inputs container the same key, the latter value is used.
func (c Counters) Merge(other Counters) Counters {
	var (
		cSize     = c.Size()
		otherSize = other.Size()
		output    = c.psMap
		iter      = other.psMap
	)
	switch {
	case cSize == 0:
		return other
	case otherSize == 0:
		return c
	case cSize < otherSize:
		output, iter = iter, output
	}
	iter.ForEach(func(key string, otherVal interface{}) {
		if val, ok := output.Lookup(key); ok {
			output = output.Set(key, otherVal.(int)+val.(int))
		} else {
			output = output.Set(key, otherVal)
		}
	})
	return Counters{output}
}

// String serializes Counters into a string.
func (c Counters) String() string {
	buf := bytes.NewBufferString("{")
	prefix := ""
	for _, key := range mapKeys(c.psMap) {
		val, _ := c.psMap.Lookup(key)
		fmt.Fprintf(buf, "%s%s: %d", prefix, key, val.(int))
		prefix = ", "
	}
	fmt.Fprintf(buf, "}")
	return buf.String()
}

// DeepEqual tests equality with other Counters
func (c Counters) DeepEqual(d Counters) bool {
	return mapEqual(c.psMap, d.psMap, reflect.DeepEqual)
}

func (c Counters) fromIntermediate(in map[string]int) Counters {
	out := ps.NewMap()
	for k, v := range in {
		out = out.Set(k, v)
	}
	return Counters{out}
}

// CodecEncodeSelf implements codec.Selfer
func (c *Counters) CodecEncodeSelf(encoder *codec.Encoder) {
	mapWrite(c.psMap, encoder, func(encoder *codec.Encoder, val interface{}) {
		i := val.(int)
		encoder.Encode(i)
	})
}

// CodecDecodeSelf implements codec.Selfer
func (c *Counters) CodecDecodeSelf(decoder *codec.Decoder) {
	out := mapRead(decoder, func(isNil bool) interface{} {
		var value int
		if !isNil {
			decoder.Decode(&value)
		}
		return value
	})
	*c = Counters{out}
}

// MarshalJSON shouldn't be used, use CodecEncodeSelf instead
func (Counters) MarshalJSON() ([]byte, error) {
	panic("MarshalJSON shouldn't be used, use CodecEncodeSelf instead")
}

// UnmarshalJSON shouldn't be used, use CodecDecodeSelf instead
func (*Counters) UnmarshalJSON(b []byte) error {
	panic("UnmarshalJSON shouldn't be used, use CodecDecodeSelf instead")
}