Vihren Atlas

Atlas › Test

TestMapCounter_HeapIndexTracking

go.temporal.io/server/service/matching/counter/TestMapCounter_HeapIndexTracking

Package
go.temporal.io/server/service/matching/counter
Suite / test hierarchy
TestMapCounter_HeapIndexTracking
Test
TestMapCounter_HeapIndexTracking
Introduced at
TestMapCounter_HeapIndexTracking Test frontier
Covered ranges
13
Covered lines
56
Covered files
1

Covered source

Expand a file to inspect source; the > gutter marks covered lines.

go.temporal.io/server/service/matching/counter/map.go 56 covered LOC · 13 ranges

Open complete file

18
19 // NewMapCounter creates a mapCounter that also tracks the top K entries.
20 > func NewMapCounter(limit int) *mapCounter { map.go
21 > return &mapCounter{
22 > m: make(map[string]int),
23 > limit: limit,
24 > }
25 > }
26
27 > func (m *mapCounter) GetPass(key string, base, inc int64) int64 { map.go
28 > c, _ := m.getPassWithOverflow(key, base, inc)
29 > return c
30 > }
31
32 > func (m *mapCounter) getPassWithOverflow(key string, base, inc int64) (int64, bool) { map.go
33 > if idx, ok := m.m[key]; ok {
34 > prev := m.heap[idx].Count map.go
35 > count := max(base, prev+inc)
36 > // inline simple case of updateHeap
37 > m.heap[idx].Count = count
38 > heap.Fix(m, idx)
39 > return count, false
40 > }
41 // not present, fall back to full updateHeap
42 > count := max(base, inc) map.go
43 > return count, m.updateHeap(key, count)
44 }
45
53 }
54
55 > func (m *mapCounter) updateHeap(key string, count int64) bool { map.go
56 > if idx, ok := m.m[key]; ok {
57 // already in heap - update count and fix
58 m.heap[idx].Count = count
61 }
62
63 > if len(m.heap) < m.limit { map.go
64 > // heap not full - add
65 > m.m[key] = len(m.heap)
66 > heap.Push(m, TopKEntry{Key: key, Count: count})
67 > return false
68 > }
69
70 // heap is full - only add if count > min
71 > if count > m.heap[0].Count { map.go
72 > // evict min
73 > evicted := heap.Pop(m).(TopKEntry) // nolint:revive // unchecked-type-assertion
74 > delete(m.m, evicted.Key)
75 > // add new
76 > m.m[key] = len(m.heap)
77 > heap.Push(m, TopKEntry{Key: key, Count: count})
78 > }
79 > return true
80 }
81
82 // implements heap.Interface using m.heap
83 > func (m *mapCounter) Len() int { return len(m.heap) } map.go
84 > func (m *mapCounter) Less(i, j int) bool { return m.heap[i].Count < m.heap[j].Count } map.go
85 > func (m *mapCounter) Swap(i, j int) { map.go
86 > m.heap[i], m.heap[j] = m.heap[j], m.heap[i]
87 > // don't forget to fix the map:
88 > m.m[m.heap[i].Key] = i
89 > m.m[m.heap[j].Key] = j
90 > }
91
92 > func (m *mapCounter) Push(x any) { map.go
93 > m.heap = append(m.heap, x.(TopKEntry))
94 > }
95
96 > func (m *mapCounter) Pop() any { map.go
97 > n := len(m.heap)
98 > entry := m.heap[n-1]
99 > m.heap[n-1] = TopKEntry{}
100 > m.heap = m.heap[0 : n-1]
101 > return entry
102 > }