add Cardinality Counter to test metrics and provider

go-kit/metrics/testmetrics/metrics.go

@@ -3,7 +3,9 @@ package testmetrics
 import (
 	"sync"
 
+	hll "github.com/axiomhq/hyperloglog"
 	"github.com/go-kit/kit/metrics"
+	xmetrics "github.com/heroku/x/go-kit/metrics"
 )
 
 // Counter accumulates a value based on Add calls.
@@ -99,3 +101,34 @@ func (h *Histogram) With(labelValues ...string) metrics.Histogram {
 	lvs := append(append([]string(nil), h.labelValues...), labelValues...)
 	return h.p.newHistogram(h.name, lvs...)
 }
+
+// Counter accumulates a value based on Add calls.
+type CardinalityCounter struct {
+	Name    string
+	lvs     []string
+	mu      sync.Mutex
+	counter *hll.Sketch
+	p       *Provider
+	sync.RWMutex
+}
+
+// With returns a new UniqueCounter with the passed in label values merged
+// with the previous label values. The counter's values are copied.
+func (c *CardinalityCounter) With(labelValues ...string) xmetrics.CardinalityCounter {
+	lvs := append(append([]string(nil), c.lvs...), labelValues...)
+	return c.p.newCardinalityCounter(c.Name, lvs...)
+}
+
+// Insert adds the item to the set to be counted.
+func (c *CardinalityCounter) Insert(i []byte) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.counter.Insert(i)
+}
+
+func (c *CardinalityCounter) Estimate() uint64 {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	return c.counter.Estimate()
+}

go-kit/metrics/testmetrics/provider.go

@@ -10,6 +10,7 @@ import (
 	"sync"
 	"testing"
 
+	hll "github.com/axiomhq/hyperloglog"
 	"github.com/go-kit/kit/metrics"
 
 	xmetrics "github.com/heroku/x/go-kit/metrics"
@@ -23,7 +24,7 @@ type Provider struct {
 	counters     map[string]*Counter
 	gauges       map[string]*Gauge
 	histograms   map[string]*Histogram
-	cardCounters map[string]*xmetrics.HLLCounter
+	cardCounters map[string]*CardinalityCounter
 	stopped      bool
 }
 
@@ -34,7 +35,7 @@ func NewProvider(t *testing.T) *Provider {
 		counters:     make(map[string]*Counter),
 		histograms:   make(map[string]*Histogram),
 		gauges:       make(map[string]*Gauge),
-		cardCounters: make(map[string]*xmetrics.HLLCounter),
+		cardCounters: make(map[string]*CardinalityCounter),
 	}
 }
 
@@ -101,13 +102,18 @@ func (p *Provider) newHistogram(name string, labelValues ...string) metrics.Hist
 
 // NewCardinalityCounter implements metrics.Provider.
 func (p *Provider) NewCardinalityCounter(name string) xmetrics.CardinalityCounter {
+	return p.newCardinalityCounter(name)
+}
+
+func (p *Provider) newCardinalityCounter(name string, labelValues ...string) xmetrics.CardinalityCounter {
 	p.Lock()
 	defer p.Unlock()
 
-	if _, ok := p.cardCounters[name]; !ok {
-		p.cardCounters[name] = xmetrics.NewHLLCounter(name)
+	k := p.keyFor(name, labelValues...)
+	if _, ok := p.cardCounters[k]; !ok {
+		p.cardCounters[k] = &CardinalityCounter{Name: name, p: p, lvs: labelValues, counter: hll.New()}
 	}
-	return p.cardCounters[name]
+	return p.cardCounters[k]
 }
 
 // CheckCounter checks that there is a registered counter
@@ -308,25 +314,30 @@ func (p *Provider) CheckStopped() {
 
 // CheckCardinalityCounter checks that there is a registered cardinality
 // counter with the name and estimate provided.
-func (p *Provider) CheckCardinalityCounter(name string, estimate uint64) {
+func (p *Provider) CheckCardinalityCounter(name string, estimate uint64, labelValues ...string) {
 	p.t.Helper()
 
 	p.Lock()
 	defer p.Unlock()
 
-	cc, ok := p.cardCounters[name]
+	k := p.keyFor(name, labelValues...)
+	cc, ok := p.cardCounters[k]
 	if !ok {
 		keys := make([]string, 0, len(p.cardCounters))
 		for k := range p.cardCounters {
 			keys = append(keys, k)
 		}
 		available := strings.Join(keys, "\n")
-		p.t.Fatalf("no cardinality counter named %s out of available cardinality counter: \n%s", name, available)
+		p.t.Fatalf("no counter named %s out of available cardinality counters: \n%s", k, available)
 	}
 	actualEstimate := cc.Estimate()
-	if actualEstimate != estimate {
+	if cc.Estimate() != actualEstimate {
 		p.t.Fatalf("%v = %v, want %v", name, actualEstimate, estimate)
 	}
+
+	if len(labelValues) > 0 && !reflect.DeepEqual(labelValues, cc.lvs) {
+		p.t.Fatalf("want counter label values: %#v, got %#v", labelValues, cc.lvs)
+	}
 }
 
 func (p *Provider) keyFor(name string, labelValues ...string) string {