Add QueryTracker

Signed-off-by: Justin Jung <[email protected]>

Author: justinjung04

pkg/util/tracker/query.go

@@ -0,0 +1,200 @@
+package tracker
+
+import (
+	"container/heap"
+	"sync"
+	"time"
+)
+
+const (
+	ttl               = 2 * time.Second
+	slidingWindowSize = 3 * time.Second
+	maxTrackedQueries = 100
+)
+
+type QueryTracker struct {
+	heap   *queryHeap
+	lookup map[string]*queryItem
+	mu     sync.Mutex
+}
+
+type queryItem struct {
+	requestID  string
+	bytesRate  *slidingWindow
+	lastUpdate time.Time
+	index      int
+}
+
+type queryHeap []*queryItem
+
+func (h queryHeap) Len() int           { return len(h) }
+func (h queryHeap) Less(i, j int) bool { return h[i].bytesRate.rate() < h[j].bytesRate.rate() }
+func (h queryHeap) Swap(i, j int) {
+	h[i], h[j] = h[j], h[i]
+	h[i].index = i
+	h[j].index = j
+}
+
+func (h *queryHeap) Push(x interface{}) {
+	item := x.(*queryItem)
+	item.index = len(*h)
+	*h = append(*h, item)
+}
+
+func (h *queryHeap) Pop() interface{} {
+	old := *h
+	n := len(old)
+	item := old[n-1]
+	item.index = -1
+	*h = old[0 : n-1]
+	return item
+}
+
+func NewQueryTracker() *QueryTracker {
+	h := &queryHeap{}
+	heap.Init(h)
+	tracker := &QueryTracker{
+		heap:   h,
+		lookup: make(map[string]*queryItem),
+	}
+
+	go tracker.loop()
+	return tracker
+}
+
+func (q *QueryTracker) loop() {
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+
+	for range ticker.C {
+		q.cleanup()
+	}
+}
+
+func (q *QueryTracker) cleanup() {
+	now := time.Now()
+	stale := now.Add(-ttl)
+
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	var toRemove []*queryItem
+	for _, item := range *q.heap {
+		if item.lastUpdate.Before(stale) {
+			toRemove = append(toRemove, item)
+		}
+	}
+
+	for _, item := range toRemove {
+		heap.Remove(q.heap, item.index)
+		delete(q.lookup, item.requestID)
+	}
+}
+
+func (q *QueryTracker) Add(requestID string, bytes uint64) {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	now := time.Now()
+	item, exists := q.lookup[requestID]
+
+	if !exists {
+		item = &queryItem{
+			requestID: requestID,
+			bytesRate: newSlidingWindow(slidingWindowSize),
+		}
+		item.bytesRate.add(bytes)
+		item.lastUpdate = now
+
+		if q.heap.Len() < maxTrackedQueries {
+			heap.Push(q.heap, item)
+			q.lookup[requestID] = item
+		} else {
+			minItem := (*q.heap)[0]
+			if item.bytesRate.rate() > minItem.bytesRate.rate() {
+				delete(q.lookup, minItem.requestID)
+				heap.Pop(q.heap)
+				heap.Push(q.heap, item)
+				q.lookup[requestID] = item
+			}
+		}
+	} else {
+		item.bytesRate.add(bytes)
+		item.lastUpdate = now
+		heap.Fix(q.heap, item.index)
+	}
+}
+
+func (q *QueryTracker) GetWorstQuery() (string, float64) {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	if q.heap.Len() == 0 {
+		return "", 0
+	}
+
+	var worstQueryID string
+	var worstRate float64
+
+	for _, item := range *q.heap {
+		rate := item.bytesRate.rate()
+		if rate > worstRate {
+			worstRate = rate
+			worstQueryID = item.requestID
+		}
+	}
+
+	return worstQueryID, worstRate
+}
+
+type slidingWindow struct {
+	buckets    []uint64
+	windowSize time.Duration
+	lastUpdate time.Time
+	currentIdx int
+	mu         sync.Mutex
+}
+
+func newSlidingWindow(windowSize time.Duration) *slidingWindow {
+	seconds := int(windowSize.Seconds())
+	return &slidingWindow{
+		buckets:    make([]uint64, seconds),
+		windowSize: windowSize,
+		lastUpdate: time.Now().Truncate(time.Second),
+	}
+}
+
+func (swr *slidingWindow) add(bytes uint64) {
+	swr.mu.Lock()
+	defer swr.mu.Unlock()
+
+	now := time.Now().Truncate(time.Second)
+
+	// Calculate how many seconds have passed since last update
+	secondsDrift := int(now.Sub(swr.lastUpdate).Seconds())
+	if secondsDrift > 0 {
+		// Clear old buckets
+		for i := 0; i < min(secondsDrift, len(swr.buckets)); i++ {
+			nextIdx := (swr.currentIdx + i) % len(swr.buckets)
+			swr.buckets[nextIdx] = 0
+		}
+		// Update current index
+		swr.currentIdx = (swr.currentIdx + secondsDrift) % len(swr.buckets)
+		swr.lastUpdate = now
+	}
+
+	// Add bytes to current bucket
+	swr.buckets[swr.currentIdx] += bytes
+}
+
+func (swr *slidingWindow) rate() float64 {
+	swr.mu.Lock()
+	defer swr.mu.Unlock()
+
+	var totalBytes uint64
+	for _, bytes := range swr.buckets {
+		totalBytes += bytes
+	}
+
+	return float64(totalBytes) / swr.windowSize.Seconds()
+}

pkg/util/tracker/query_test.go

@@ -0,0 +1,125 @@
+package tracker
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestQueryTracker(t *testing.T) {
+	queryTracker := NewQueryTracker()
+	queryTracker.Add("r-1", 3000)
+	queryTracker.Add("r-2", 300)
+	requestID, rate := queryTracker.GetWorstQuery()
+
+	assert.Equal(t, "r-1", requestID)
+	assert.Equal(t, float64(1000), rate)
+}
+
+func TestQueryTracker_MaxQueries(t *testing.T) {
+	queryTracker := NewQueryTracker()
+
+	// Add more than maxTrackedQueries with low rates
+	for i := 0; i < 200; i++ {
+		queryTracker.Add(fmt.Sprintf("low-%d", i), 100)
+	}
+
+	// Add high-rate query
+	queryTracker.Add("high-rate", 5000)
+
+	// Should track the high-rate query and evict low-rate ones
+	requestID, rate := queryTracker.GetWorstQuery()
+	assert.Equal(t, "high-rate", requestID)
+	assert.True(t, rate > 1000)
+
+	// Verify heap size is bounded
+	queryTracker.mu.Lock()
+	heapSize := queryTracker.heap.Len()
+	queryTracker.mu.Unlock()
+	assert.LessOrEqual(t, heapSize, maxTrackedQueries)
+}
+
+func TestQueryTracker_TTL(t *testing.T) {
+	queryTracker := NewQueryTracker()
+	for i := 0; i < 200; i++ {
+		queryTracker.Add(fmt.Sprintf("low-%d", i), 300)
+	}
+	requestID, rate := queryTracker.GetWorstQuery()
+	assert.Equal(t, float64(100), rate)
+
+	time.Sleep(4 * time.Second) // expire all items
+	requestID, rate = queryTracker.GetWorstQuery()
+
+	assert.Equal(t, "", requestID)
+	assert.Equal(t, float64(0), rate)
+}
+
+func TestQueryTracker_EmptyHeap(t *testing.T) {
+	queryTracker := NewQueryTracker()
+	requestID, rate := queryTracker.GetWorstQuery()
+	assert.Equal(t, "", requestID)
+	assert.Equal(t, float64(0), rate)
+}
+
+func TestQueryTracker_UpdateExistingQuery(t *testing.T) {
+	queryTracker := NewQueryTracker()
+	queryTracker.Add("r-1", 1000)
+	initialRate := queryTracker.lookup["r-1"].bytesRate.rate()
+	
+	// Add more bytes to same query
+	queryTracker.Add("r-1", 2000)
+	updatedRate := queryTracker.lookup["r-1"].bytesRate.rate()
+	
+	assert.True(t, updatedRate > initialRate)
+	requestID, rate := queryTracker.GetWorstQuery()
+	assert.Equal(t, "r-1", requestID)
+	assert.Equal(t, float64(1000), rate) // 3000 bytes / 3 seconds
+}
+
+func TestQueryTracker_HeapEviction(t *testing.T) {
+	queryTracker := NewQueryTracker()
+	
+	// Fill heap to capacity with low-rate queries
+	for i := 0; i < maxTrackedQueries; i++ {
+		queryTracker.Add(fmt.Sprintf("low-%d", i), 100)
+	}
+	
+	// Verify heap is at capacity
+	assert.Equal(t, maxTrackedQueries, queryTracker.heap.Len())
+	
+	// Add a high-rate query that should evict the lowest
+	queryTracker.Add("high-rate", 10000)
+	
+	// Heap should still be at capacity
+	assert.Equal(t, maxTrackedQueries, queryTracker.heap.Len())
+	
+	// High-rate query should be tracked
+	_, exists := queryTracker.lookup["high-rate"]
+	assert.True(t, exists)
+	
+	// Should be the worst query
+	requestID, rate := queryTracker.GetWorstQuery()
+	assert.Equal(t, "high-rate", requestID)
+	assert.True(t, rate > 3000)
+}
+
+func TestQueryTracker_NoEvictionForLowRate(t *testing.T) {
+	queryTracker := NewQueryTracker()
+	
+	// Fill heap with medium-rate queries
+	for i := 0; i < maxTrackedQueries; i++ {
+		queryTracker.Add(fmt.Sprintf("med-%d", i), 1000)
+	}
+	
+	// Try to add a lower-rate query
+	queryTracker.Add("low-rate", 50)
+	
+	// Low-rate query should not be tracked
+	_, exists := queryTracker.lookup["low-rate"]
+	assert.False(t, exists)
+	
+	// Heap should still be at capacity
+	assert.Equal(t, maxTrackedQueries, queryTracker.heap.Len())
+}