fix(router): evict stale canonical entry on double-check TTL miss

Router.Get's canonical double-check branch trusted any existing entry
under the canonical key without re-checking TTL. If an earlier
agent_key caller wrote the entry and the TTL expired, a later UUID-form
caller would resolve fresh, hit the canonical branch, find the stale
entry, and return it — indefinitely — because the raw-UUID key was
never the map key and the initial-miss eviction branch did nothing.

Re-check TTL inside the double-check branch and evict+rewrite when
stale. Regression test primes a canonical entry with cachedAt set
2×TTL in the past, then asserts a UUID-form Get re-invokes the
resolver and the returned agent reflects the fresh resolver output.
This commit is contained in:
viettranx committed 2026-04-11 21:22:23 +07:00
1 parent 6e86ba9a39
commit 7799a9c88a
2 files changed
+79 -3

No files matched your search

+9 -3
View File
@@ -106,10 +106,16 @@ func (r *Router) Get(ctx context.Context, agentID string) (Agent, error) {
// entries exist for the same logical agent.
canonicalKey := agentCacheKey(ctx, ag.ID())
r.mu.Lock()
// Double-check: another goroutine might have created it under the canonical key
// Double-check: another goroutine might have created it under the canonical key.
// Re-check TTL so a UUID-form caller cannot receive a stale canonical entry
// indefinitely — this branch is the only eviction path for entries the caller
// never wrote under the raw input key.
if existing, ok := r.agents[canonicalKey]; ok {
r.mu.Unlock()
return existing.agent, nil
if r.ttl == 0 || time.Since(existing.cachedAt) < r.ttl {
r.mu.Unlock()
return existing.agent, nil
}
delete(r.agents, canonicalKey)
}
r.agents[canonicalKey] = &agentEntry{agent: ag, cachedAt: time.Now()}
r.mu.Unlock()
@@ -5,6 +5,7 @@ import (
"strings"
"sync/atomic"
"testing"
"time"
"github.com/google/uuid"
@@ -171,3 +172,72 @@ func TestRouterGet_DualTenantSameAgentKey(t *testing.T) {
t.Errorf("tenantB cache entry %q missing", keyB)
}
}
// TestRouterGet_CanonicalDoubleCheckEvictsStaleEntry pins the TTL check inside
// the canonical double-check branch. Scenario: an earlier caller wrote the
// canonical entry, TTL expires, a UUID-form caller arrives. The raw UUID key
// is not in the map so the initial-miss eviction branch does nothing; the
// resolver runs and the canonical double-check finds the stale entry. Without
// the TTL re-check this would return the stale agent indefinitely — only
// InvalidateAgent could rescue it. The fix evicts + rewrites.
func TestRouterGet_CanonicalDoubleCheckEvictsStaleEntry(t *testing.T) {
r := NewRouter()
var resolveCount atomic.Int32
calls := 0
r.SetResolver(func(_ context.Context, _ string) (Agent, error) {
resolveCount.Add(1)
calls++
return &stubAgent{id: "goctech-leader", running: calls == 1}, nil
})
tenantID := uuid.New()
ctx := store.WithTenantID(context.Background(), tenantID)
canonicalKey := tenantID.String() + ":goctech-leader"
// Prime canonical entry with agent_key caller.
if _, err := r.Get(ctx, "goctech-leader"); err != nil {
t.Fatalf("initial Get(agent_key): %v", err)
}
if got := resolveCount.Load(); got != 1 {
t.Fatalf("resolver calls after prime = %d, want 1", got)
}
// Force the canonical entry to look stale (cachedAt far past TTL).
r.mu.Lock()
entry, ok := r.agents[canonicalKey]
if !ok {
r.mu.Unlock()
t.Fatalf("canonical entry %q missing after prime", canonicalKey)
}
entry.cachedAt = time.Now().Add(-2 * defaultRouterTTL)
r.mu.Unlock()
// UUID-form caller arrives. Initial miss branch does nothing (raw UUID key
// was never written). Resolver runs. Canonical double-check must detect the
// stale entry and evict+rewrite, not return it.
uuidInput := uuid.New().String()
got, err := r.Get(ctx, uuidInput)
if err != nil {
t.Fatalf("Get(uuid) after stale prime: %v", err)
}
if resolveCount.Load() != 2 {
t.Errorf("resolver calls after stale Get = %d, want 2", resolveCount.Load())
}
if got.IsRunning() {
// The stale entry had running=true (calls==1). The fresh resolver return
// has running=false (calls==2). If IsRunning is true we got the stale one.
t.Errorf("Get returned stale entry (running=true) instead of fresh resolver result")
}
// Canonical entry must now hold a fresh cachedAt.
r.mu.RLock()
fresh, ok := r.agents[canonicalKey]
r.mu.RUnlock()
if !ok {
t.Fatalf("canonical entry %q missing after stale refresh", canonicalKey)
}
if time.Since(fresh.cachedAt) > time.Second {
t.Errorf("canonical entry cachedAt not refreshed: age = %v", time.Since(fresh.cachedAt))
}
}