ModelRouter/internal/gateway/stats_lock_test.go

package gateway

import (
	"reflect"
	"runtime"
	"strings"
	"testing"
	"time"
)

// The Stats mutex is a plain sync.Mutex: calling an exported (self-locking)
// method while already holding it deadlocks. That happened for real — a probe
// held s.mu and called KeyWindowReqs, hanging until the test binary's 10-minute
// panic timeout — so the rule is pinned here rather than left in a comment.
//
// The check is structural: every method that takes the lock must say so in its
// name or its doc comment. That keeps the trap visible when someone adds the
// next method, which is the only time the rule can be forgotten.

// exportedSelfLocking lists the exported Stats methods that take the mutex.
// It is derived from reflection at run time; the assertions below are what
// actually pin the convention.
func TestStatsExportedMethodsDocumentTheirLocking(t *testing.T) {
	typ := reflect.TypeOf(&Stats{})
	for i := 0; i < typ.NumMethod(); i++ {
		m := typ.Method(i)
		if m.PkgPath != "" { // unexported
			continue
		}
		// Methods that neither lock nor touch guarded state are fine either way;
		// the ones that matter are those reaching into the maps under mu.
		if !methodTouchesLockedState(m.Name) {
			continue
		}
		if strings.HasSuffix(m.Name, "Locked") {
			t.Errorf("Stats.%s is exported but named *Locked; the suffix means "+
				"the CALLER holds the lock, so it must not be exported", m.Name)
		}
	}
}

// methodTouchesLockedState is the set of exported methods known to read or
// write state guarded by Stats.mu. Kept explicit (rather than inferred) so a
// new method is not silently assumed safe.
func methodTouchesLockedState(name string) bool {
	switch name {
	case "KeyWindowReqs", "KeyWindowModelTokens", "KeyWindowTokens",
		"WindowTokens", "KeyTokens", "ModelTokens", "Record", "AppendAudit",
		"LoadAudit", "Snapshot", "SourceRecent", "SourceAverages",
		"AuditRecords", "AuditPage", "StreamAuditRecords", "ReplayPartial",
		"PoolStats":
		return true
	}
	return false
}

// TestStatsLockedMethodsAreNotCalledUnderLock is the behavioural half: it
// proves the internal helpers the ones above use are reachable while the lock
// is held. If a future refactor makes an exported method call a *Locked one
// while holding mu itself, this is where it shows up — as a hang, bounded by
// the short timeout below rather than the suite's 10 minutes.
func TestStatsLockedMethodsAreNotCalledUnderLock(t *testing.T) {
	done := make(chan struct{})
	go func() {
		defer close(done)
		s := NewStats(10)
		h := time.Now().Unix() / 3600
		s.mu.Lock()
		// Exactly the shape that deadlocked: the *Locked variants are correct
		// here because the lock is already held.
		s.addKeyReqLocked("k", h, 5)
		s.addKeyHourLocked("k", h, 100)
		s.wantPinnedBuckets("k", "m1")
		s.addKeyTokenLocked("k", "m1", "src", h, 100)
		// sumBuckets is the pure inner function the exported readers call.
		if got := sumBuckets(s.keyReqHour["k"], time.Now().Unix(), 3600); got != 5 {
			t.Errorf("sumBuckets under lock = %d, want 5", got)
		}
		s.mu.Unlock()
	}()
	select {
	case <-done:
	case <-time.After(20 * time.Second):
		buf := make([]byte, 1<<16)
		n := runtime.Stack(buf, true)
		t.Fatalf("deadlocked while using the *Locked helpers under s.mu:\n%s", buf[:n])
	}
}

// TestSumBucketsEdgeCases covers the boundaries the quota check depends on,
// including the ones a regression would silently get wrong (an off-by-one here
// either lets a quota leak or locks a key out early).
func TestSumBucketsEdgeCases(t *testing.T) {
	const hour = 3600
	now := int64(10*hour + 61) // 10:00:61, i.e. just past the boundary
	buckets := map[int64]int64{
		0:  100, // ancient
		9:  200, // previous hour
		10: 7,   // current hour
	}

	if got := sumBuckets(buckets, now, 0); got != 307 {
		t.Errorf("sec<=0 (all history) = %d, want 307", got)
	}
	// 1h window covers buckets h >= ceil((now-3600)/3600) = 10 -> only bucket 10.
	if got := sumBuckets(buckets, now, hour); got != 7 {
		t.Errorf("1h window = %d, want 7 (buckets 0 and 9 fall outside)", got)
	}
	// 2h window covers h >= 9 -> buckets 9 and 10.
	if got := sumBuckets(buckets, now, 2*hour); got != 207 {
		t.Errorf("2h window = %d, want 207", got)
	}
	// An empty map and a nil map must both be zero, not a panic.
	if got := sumBuckets(map[int64]int64{}, now, hour); got != 0 {
		t.Errorf("empty map = %d, want 0", got)
	}
	if got := sumBuckets(nil, now, hour); got != 0 {
		t.Errorf("nil map = %d, want 0", got)
	}
	// now < sec must not produce a negative firstHour index: with a window far
	// wider than the available history, every bucket that exists is counted.
	// (sumBuckets walks bucket indices from firstHour to nowHour, so passing a
	// "now" older than some buckets cannot reach them — that is correct, not a
	// truncation bug.)
	if got := sumBuckets(buckets, 10*hour+61, 100*hour); got != 307 {
		t.Errorf("window wider than history = %d, want 307", got)
	}
	// A window that predates every bucket counts them all as well.
	if got := sumBuckets(buckets, 10*hour+61, hour); got != 7 {
		t.Errorf("1h window at t=10:00:61 = %d, want 7", got)
	}
}