Fairness counter: keep top K on resize (#9187)

## What changed?
Fairness counters can use a count-min sketch that can automatically grow
on contention. But when it grows it drops all values. This adds a
mechanism to preserve counts of the top 100 keys.

## Why?
More accurate fairness

## How did you test it?
- [x] added new unit test(s)

service/matching/backlog_manager_test.go

@@ -94,7 +94,7 @@ func (s *BacklogManagerTestSuite) SetupTest() {
 			s.logger,
 			nil,
 			metrics.NoopMetricsHandler,
-			func() counter.Counter { return counter.NewMapCounter() },
+			func() counter.Counter { return counter.NewMapCounter(1000) },
 			false,
 		)
 	} else if s.newMatcher {

service/matching/counter/cmsketch.go

@@ -21,6 +21,9 @@ type (
 		MaxW          int     // cap for W
 	}
 
+	// topKFunc is a callback that returns the top-K entries to preserve during resize.
+	topKFunc func() []TopKEntry
+
 	// cmSketch is a Counter that uses a count-min sketch.
 	// cmSketch is not safe for concurrent use.
 	cmSketch struct {
@@ -31,20 +34,23 @@ type (
 		cells []int64
 
 		skips, incs int // used to calculate skip rate
+
+		topKProvider topKFunc // callback to get top-K entries on resize
 	}
 )
 
 var _ Counter = (*cmSketch)(nil)
 
-func NewCMSketchCounter(params CMSketchParams, src rand.Source) *cmSketch {
-	params.D = max(2, params.D)
-	params.W = max(2, params.W)
+func NewCMSketchCounter(params CMSketchParams, src rand.Source, topKProvider topKFunc) *cmSketch {
+	params.D = max(1, params.D)
+	params.W = max(1, params.W)
 	params.Grow.SkipRateDecay = max(1_000, params.Grow.SkipRateDecay)
 	return &cmSketch{
-		params: params,
-		seed0:  maphash.MakeSeed(),
-		seeds:  makeSeeds(params.D, src),
-		cells:  make([]int64, params.W*params.D),
+		params:       params,
+		seed0:        maphash.MakeSeed(),
+		seeds:        makeSeeds(params.D, src),
+		cells:        make([]int64, params.W*params.D),
+		topKProvider: topKProvider,
 	}
 }
 
@@ -84,6 +90,10 @@ func (s *cmSketch) EstimateDistinctKeys() int {
 	return count / s.params.D
 }
 
+func (s *cmSketch) TopK() []TopKEntry {
+	return s.topKProvider()
+}
+
 func (s *cmSketch) fillIndexes(k string, indexes []int) {
 	w := s.params.W
 	// get 64 bits of hash
@@ -106,11 +116,28 @@ func (s *cmSketch) maybeGrow() {
 		s.SkipRate() < s.params.Grow.Threshold {
 		return
 	}
-	// TODO: instead of throwing away the whole sketch, try to preserve a sample of large values
+	// get top entries before resetting (if provider available)
+	var topK []TopKEntry
+	if s.topKProvider != nil {
+		topK = s.topKProvider()
+	}
+
 	s.params.W = min(int(float64(s.params.W)*s.params.Grow.Ratio), s.params.Grow.MaxW)
 	s.seed0 = maphash.MakeSeed()
 	s.cells = make([]int64, s.params.W*s.params.D)
 	s.skips, s.incs = 0, 0
+
+	if len(topK) > 0 {
+		// restore top entries after resize. GetPass can in theory call back into maybeGrow,
+		// but only if called more than SkipRateDecay times, so just limit to that.
+		topK = topK[:min(len(topK), s.params.Grow.SkipRateDecay)]
+		for _, entry := range topK {
+			_ = s.GetPass(entry.Key, entry.Count, 0)
+		}
+	}
+
+	// reset these again
+	s.skips, s.incs = 0, 0
 }
 
 func (s *cmSketch) getByIndexes(indexes []int) int64 {

service/matching/counter/cmsketch_test.go

@@ -11,7 +11,7 @@ import (
 
 func TestCMSketch_Basic(t *testing.T) {
 	src := rand.NewPCG(rand.Uint64(), rand.Uint64())
-	cms := NewCMSketchCounter(CMSketchParams{W: 10, D: 3}, src)
+	cms := NewCMSketchCounter(CMSketchParams{W: 10, D: 3}, src, nil)
 
 	// one key
 	assert.Equal(t, int64(1), cms.GetPass("one", 0, 1))
@@ -27,7 +27,7 @@ func TestCMSketch_Basic(t *testing.T) {
 
 func TestCMSketch_CrossMaxInt32(t *testing.T) {
 	src := rand.NewPCG(rand.Uint64(), rand.Uint64())
-	cms := NewCMSketchCounter(CMSketchParams{W: 10, D: 3}, src)
+	cms := NewCMSketchCounter(CMSketchParams{W: 10, D: 3}, src, nil)
 
 	for _, base := range []int64{
 		math.MaxInt32 - 123,
@@ -53,7 +53,7 @@ func TestCMSketch_Grow(t *testing.T) {
 			Ratio:         2,
 			MaxW:          10_000,
 		},
-	}, src)
+	}, src, nil)
 
 	for i := range 1000 {
 		cms.GetPass(fmt.Sprintf("key%d", i), 0, 1)
@@ -62,3 +62,38 @@ func TestCMSketch_Grow(t *testing.T) {
 	cms.GetPass("onemore", 0, 1)
 	assert.Greater(t, cms.params.W, 10)
 }
+
+func TestCMSketch_Grow_PreservedOnResize(t *testing.T) {
+	var topKCalls int
+	topK := func() []TopKEntry {
+		topKCalls++
+		return []TopKEntry{
+			TopKEntry{Key: "topkey1", Count: 9999},
+			TopKEntry{Key: "topkey2", Count: 99999},
+		}
+	}
+
+	src := rand.NewPCG(rand.Uint64(), rand.Uint64())
+	cms := NewCMSketchCounter(CMSketchParams{
+		W: 10,
+		D: 3,
+		Grow: CMSGrowParams{
+			SkipRateDecay: 1_000,
+			Threshold:     0.1,
+			Ratio:         2,
+			MaxW:          10_000,
+		},
+	}, src, topK)
+
+	for i := range 1000 {
+		cms.GetPass(fmt.Sprintf("key%d", i), 0, 1)
+	}
+	assert.Equal(t, 10, cms.params.W)
+	assert.Zero(t, topKCalls)
+	cms.GetPass("onemore", 0, 1)
+	assert.Greater(t, cms.params.W, 10)
+	assert.Equal(t, 1, topKCalls)
+
+	assert.GreaterOrEqual(t, cms.GetPass("topkey1", 0, 1), int64(9999))
+	assert.GreaterOrEqual(t, cms.GetPass("topkey2", 0, 1), int64(99999))
+}

service/matching/counter/counter.go

@@ -6,4 +6,13 @@ type Counter interface {
 	GetPass(key string, base, inc int64) int64
 	// EstimateDistinctKeys returns an estimate of the number of distinct keys in the counter.
 	EstimateDistinctKeys() int
+	// Get the counts for a small number of top keys (the maximum keys returned will be limited
+	// by the Counter's configuration). Not guaranteed to return the exact top.
+	TopK() []TopKEntry
+}
+
+// TopKEntry is a count for a value with high count.
+type TopKEntry struct {
+	Key   string
+	Count int64
 }

service/matching/counter/hybrid.go

@@ -11,11 +11,14 @@ type (
 	}
 
 	// hybridCounter is a Counter that uses a mapCounter until it has params.MapLimit entries,
-	// then switches to a cmSketch. hybridCounter is not safe for concurrent use.
+	// then switches to a cmSketch. After switching, it continues to track the top MapLimit
+	// entries in the mapCounter, so they can be preserved when cmSketch resizes and persisted.
+	// hybridCounter is not safe for concurrent use.
 	hybridCounter struct {
-		Counter
-		params CounterParams
-		src    rand.Source
+		mapCounter mapCounter
+		cmSketch   *cmSketch
+		params     CounterParams
+		src        rand.Source
 	}
 )
 
@@ -37,25 +40,43 @@ var DefaultCounterParams = CounterParams{
 
 func NewHybridCounter(params CounterParams, src rand.Source) *hybridCounter {
 	return &hybridCounter{
-		Counter: NewMapCounter(),
-		params:  params,
-		src:     src,
+		mapCounter: *NewMapCounter(params.MapLimit),
+		params:     params,
+		src:        src,
 	}
 }
 
 func (h *hybridCounter) GetPass(key string, base int64, inc int64) int64 {
-	p := h.Counter.GetPass(key, base, inc)
-	if m, ok := h.Counter.(*mapCounter); ok && len(m.m) > h.params.MapLimit {
-		h.migrateToCMS(m.m)
+	if h.cmSketch != nil {
+		p := h.cmSketch.GetPass(key, base, inc)
+		// after migration, continue updating top-K tracker
+		_ = h.mapCounter.updateHeap(key, p)
+		return p
+	}
+
+	p, overflow := h.mapCounter.getPassWithOverflow(key, base, inc)
+	if overflow {
+		h.migrateToCMS()
 	}
 	return p
 }
 
-func (h *hybridCounter) migrateToCMS(m map[string]int64) {
-	cms := NewCMSketchCounter(h.params.CMS, h.src)
+func (h *hybridCounter) migrateToCMS() {
+	h.cmSketch = NewCMSketchCounter(h.params.CMS, h.src, h.mapCounter.TopK)
 	// move existing counts into CMS
-	for key, count := range m {
-		_ = cms.GetPass(key, count, 0)
+	for _, entry := range h.mapCounter.heap {
+		_ = h.cmSketch.GetPass(entry.Key, entry.Count, 0)
 	}
-	h.Counter = cms
+}
+
+func (h *hybridCounter) EstimateDistinctKeys() int {
+	if h.cmSketch != nil {
+		return h.cmSketch.EstimateDistinctKeys()
+	}
+	return h.mapCounter.EstimateDistinctKeys()
+}
+
+// TopK returns the current top-K entries being tracked.
+func (h *hybridCounter) TopK() []TopKEntry {
+	return h.mapCounter.TopK()
 }

service/matching/counter/hybrid_test.go

@@ -0,0 +1,126 @@
+package counter
+
+import (
+	"fmt"
+	"math/rand/v2"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestHybridCounter_StartsWithMap(t *testing.T) {
+	src := rand.NewPCG(rand.Uint64(), rand.Uint64())
+	h := NewHybridCounter(CounterParams{
+		MapLimit: 10,
+		CMS:      CMSketchParams{W: 10, D: 3},
+	}, src)
+
+	// using mapCounter initially
+	assert.Nil(t, h.cmSketch)
+
+	for i := range 5 {
+		h.GetPass(fmt.Sprintf("key%d", i), 0, int64(i+1))
+	}
+
+	// should still be using mapCounter
+	assert.Nil(t, h.cmSketch)
+}
+
+func TestHybridCounter_Migrates(t *testing.T) {
+	src := rand.NewPCG(rand.Uint64(), rand.Uint64())
+	h := NewHybridCounter(CounterParams{
+		MapLimit: 5,
+		CMS:      CMSketchParams{W: 100, D: 3},
+	}, src)
+
+	// add entries up to limit
+	for i := range 5 {
+		h.GetPass(fmt.Sprintf("key%d", i), 0, int64(i+1))
+	}
+	assert.Nil(t, h.cmSketch, "should still be map at limit")
+
+	// one more should trigger migration
+	h.GetPass("trigger", 0, 100)
+	assert.NotNil(t, h.cmSketch, "should have migrated")
+}
+
+func TestHybridCounter_TopKTracking(t *testing.T) {
+	src := rand.NewPCG(rand.Uint64(), rand.Uint64())
+	h := NewHybridCounter(CounterParams{
+		MapLimit: 3,
+		CMS:      CMSketchParams{W: 100, D: 3},
+	}, src)
+
+	h.GetPass("low", 0, 1)
+	h.GetPass("mid", 0, 5)
+	h.GetPass("high", 0, 10)
+	h.GetPass("trigger", 0, 3) // triggers migration
+
+	require.NotNil(t, h.cmSketch, "should have migrated")
+
+	assert.ElementsMatch(t, []TopKEntry{
+		TopKEntry{Key: "mid", Count: 5},
+		TopKEntry{Key: "high", Count: 10},
+		TopKEntry{Key: "trigger", Count: 3},
+	}, h.TopK())
+}
+
+func TestHybridCounter_TopKUpdatesAfterMigration(t *testing.T) {
+	src := rand.NewPCG(rand.Uint64(), rand.Uint64())
+	h := NewHybridCounter(CounterParams{
+		MapLimit: 2,
+		CMS:      CMSketchParams{W: 100, D: 3},
+	}, src)
+
+	h.GetPass("a", 0, 10)
+	h.GetPass("b", 0, 20)
+	h.GetPass("c", 0, 5) // triggers migration
+
+	require.NotNil(t, h.cmSketch)
+
+	// TopK should still update
+	h.GetPass("d", 0, 100)
+	assert.ElementsMatch(t, []TopKEntry{
+		TopKEntry{Key: "d", Count: 100},
+		TopKEntry{Key: "b", Count: 20},
+	}, h.TopK())
+}
+
+func TestHybridCounter_TopKPreservedOnCMSResize(t *testing.T) {
+	src := rand.NewPCG(rand.Uint64(), rand.Uint64())
+	h := NewHybridCounter(CounterParams{
+		MapLimit: 5,
+		CMS: CMSketchParams{
+			W: 10,
+			D: 5,
+			Grow: CMSGrowParams{
+				SkipRateDecay: 1_000,
+				Threshold:     0.1,
+				Ratio:         2,
+				MaxW:          1000,
+			},
+		},
+	}, src)
+
+	// add keys until we see a resize
+	for i := 0; true; i++ {
+		if h.cmSketch == nil {
+			h.GetPass(fmt.Sprintf("key%d", i), 0, 1)
+			continue
+		}
+
+		prevTop := h.TopK()
+		prevW := h.cmSketch.params.W
+
+		h.GetPass(fmt.Sprintf("key%d", i), 0, 1)
+
+		if h.cmSketch.params.W > prevW {
+			// we resized, now check
+			for _, entry := range prevTop {
+				assert.GreaterOrEqual(t, h.cmSketch.GetPass(entry.Key, 0, 1), entry.Count)
+			}
+			break
+		}
+	}
+}

service/matching/counter/map.go

@@ -1,23 +1,99 @@
 package counter
 
+import (
+	"container/heap"
+	"slices"
+)
+
 // mapCounter is a Counter that stores counts in a map.
+// It also maintains a min-heap to efficiently track the top-K entries.
 // mapCounter is not safe for concurrent use.
 type mapCounter struct {
-	m map[string]int64
+	m     map[string]int
+	heap  topKHeap
+	limit int
 }
 
 var _ Counter = (*mapCounter)(nil)
 
-func NewMapCounter() *mapCounter {
-	return &mapCounter{m: make(map[string]int64)}
+// NewMapCounter creates a mapCounter that also tracks the top K entries.
+func NewMapCounter(limit int) *mapCounter {
+	return &mapCounter{
+		m:     make(map[string]int),
+		limit: limit,
+	}
 }
 
 func (m *mapCounter) GetPass(key string, base, inc int64) int64 {
-	c := max(base, m.m[key]+inc)
-	m.m[key] = c
+	c, _ := m.getPassWithOverflow(key, base, inc)
 	return c
 }
 
+func (m *mapCounter) getPassWithOverflow(key string, base, inc int64) (int64, bool) {
+	if idx, ok := m.m[key]; ok {
+		prev := m.heap[idx].Count
+		count := max(base, prev+inc)
+		// inline simple case of updateHeap
+		m.heap[idx].Count = count
+		heap.Fix(&m.heap, idx)
+		return count, false
+	}
+	// not present, fall back to full updateHeap
+	count := max(base, inc)
+	return count, m.updateHeap(key, count)
+}
+
 func (m *mapCounter) EstimateDistinctKeys() int {
-	return len(m.m)
+	return len(m.heap)
+}
+
+// TopK returns the top-K entries by count.
+func (m *mapCounter) TopK() []TopKEntry {
+	return slices.Clone(m.heap)
+}
+
+func (m *mapCounter) updateHeap(key string, count int64) bool {
+	if idx, ok := m.m[key]; ok {
+		// already in heap - update count and fix
+		m.heap[idx].Count = count
+		heap.Fix(&m.heap, idx)
+		return false
+	}
+
+	if len(m.heap) < m.limit {
+		// heap not full - add
+		m.m[key] = len(m.heap)
+		heap.Push(&m.heap, TopKEntry{Key: key, Count: count})
+		return false
+	}
+
+	// heap is full - only add if count > min
+	if count > m.heap[0].Count {
+		// evict min
+		evicted := heap.Pop(&m.heap).(TopKEntry) // nolint:revive // unchecked-type-assertion
+		delete(m.m, evicted.Key)
+		// add new
+		m.m[key] = len(m.heap)
+		heap.Push(&m.heap, TopKEntry{Key: key, Count: count})
+	}
+	return true
+}
+
+// topKHeap implements heap.Interface as a min-heap (smallest count at root)
+type topKHeap []TopKEntry
+
+func (h topKHeap) Len() int           { return len(h) }
+func (h topKHeap) Less(i, j int) bool { return h[i].Count < h[j].Count }
+func (h topKHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
+
+func (h *topKHeap) Push(x any) {
+	*h = append(*h, x.(TopKEntry))
+}
+
+func (h *topKHeap) Pop() any {
+	n := len(*h)
+	entry := (*h)[n-1]
+	(*h)[n-1] = TopKEntry{}
+	*h = (*h)[0 : n-1]
+	return entry
 }

service/matching/counter/map_test.go

@@ -0,0 +1,89 @@
+package counter
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMapCounter_Basic(t *testing.T) {
+	m := NewMapCounter(100)
+
+	assert.Equal(t, int64(1), m.GetPass("a", 0, 1))
+	assert.Equal(t, int64(3), m.GetPass("a", 0, 2))
+	assert.Equal(t, int64(10), m.GetPass("b", 0, 10))
+	assert.Equal(t, 2, m.EstimateDistinctKeys())
+}
+
+func TestMapCounter_TopK(t *testing.T) {
+	m := NewMapCounter(3)
+
+	m.GetPass("low1", 0, 1)
+	m.GetPass("low2", 0, 2)
+	m.GetPass("mid", 0, 5)
+	m.GetPass("high1", 0, 10)
+	m.GetPass("high2", 0, 8)
+
+	assert.ElementsMatch(t, []TopKEntry{
+		TopKEntry{Key: "high1", Count: 10},
+		TopKEntry{Key: "high2", Count: 8},
+		TopKEntry{Key: "mid", Count: 5},
+	}, m.TopK())
+}
+
+func TestMapCounter_TopK_Update(t *testing.T) {
+	m := NewMapCounter(2)
+
+	// Start with two entries
+	m.GetPass("a", 0, 1)
+	m.GetPass("b", 0, 2)
+
+	topK := m.TopK()
+	assert.Len(t, topK, 2)
+
+	// Update "a" to have the highest count
+	m.GetPass("a", 0, 100)
+
+	assert.ElementsMatch(t, []TopKEntry{
+		TopKEntry{Key: "a", Count: 101},
+		TopKEntry{Key: "b", Count: 2},
+	}, m.TopK())
+}
+
+func TestMapCounter_TopK_Eviction(t *testing.T) {
+	m := NewMapCounter(2)
+
+	m.GetPass("a", 0, 10)
+	m.GetPass("b", 0, 20)
+
+	// "c" with count 5 should not evict anything
+	m.GetPass("c", 0, 5)
+	assert.ElementsMatch(t, []TopKEntry{
+		TopKEntry{Key: "a", Count: 10},
+		TopKEntry{Key: "b", Count: 20},
+	}, m.TopK())
+
+	// "d" with count 15 should evict "a"
+	m.GetPass("d", 0, 15)
+	assert.ElementsMatch(t, []TopKEntry{
+		TopKEntry{Key: "d", Count: 15},
+		TopKEntry{Key: "b", Count: 20},
+	}, m.TopK())
+}
+
+func TestMapCounter_TopK_ManyEntries(t *testing.T) {
+	m := NewMapCounter(5)
+
+	for i := range 100 {
+		m.GetPass(fmt.Sprintf("key%d", i), 0, int64(i))
+	}
+
+	assert.ElementsMatch(t, []TopKEntry{
+		TopKEntry{Key: "key99", Count: 99},
+		TopKEntry{Key: "key98", Count: 98},
+		TopKEntry{Key: "key97", Count: 97},
+		TopKEntry{Key: "key96", Count: 96},
+		TopKEntry{Key: "key95", Count: 95},
+	}, m.TopK())
+}