diff --git a/devshard/cmd/devshardctl/response_cache.go b/devshard/cmd/devshardctl/response_cache.go index d915029229..0e3ab01156 100644 --- a/devshard/cmd/devshardctl/response_cache.go +++ b/devshard/cmd/devshardctl/response_cache.go @@ -6,6 +6,7 @@ import ( "encoding/hex" "io" "net/http" + "sort" "strings" "sync" "time" @@ -13,10 +14,21 @@ import ( const chatResponseCacheTTL = time.Hour +// defaultChatResponseCacheMaxEntries bounds how many cached chat responses the +// gateway keeps in memory at once. The cache is a single process-wide dedup +// layer keyed by the full request body, so under a normal chat workload +// (mostly unique bodies) it would otherwise grow by one full response body per +// distinct request and never shrink: an expired entry is only reclaimed when +// the *same* key is looked up again, which for unique bodies never happens. +// Capping the entry count keeps the footprint bounded regardless of traffic or +// uptime, matching the bounded-map idiom used by rateLimiter and gossip dedup. +const defaultChatResponseCacheMaxEntries = 2048 + type chatResponseCache struct { - mu sync.Mutex - ttl time.Duration - entries map[string]cachedChatResponse + mu sync.Mutex + ttl time.Duration + maxEntries int + entries map[string]cachedChatResponse } type cachedChatResponse struct { @@ -34,8 +46,9 @@ func newChatResponseCache(ttl time.Duration) *chatResponseCache { ttl = chatResponseCacheTTL } return &chatResponseCache{ - ttl: ttl, - entries: make(map[string]cachedChatResponse), + ttl: ttl, + maxEntries: defaultChatResponseCacheMaxEntries, + entries: make(map[string]cachedChatResponse), } } @@ -83,9 +96,43 @@ func (c *chatResponseCache) Set(key string, entry cachedChatResponse, now time.T c.mu.Lock() defer c.mu.Unlock() + if _, exists := c.entries[key]; !exists && c.maxEntries > 0 && len(c.entries) >= c.maxEntries { + c.evictLocked(now) + } c.entries[key] = entry } +// evictLocked bounds the cache. It first reclaims every expired entry (the +// common case for a unique-body workload) and, only if the cache is still at +// capacity with live entries, drops the oldest half in a single pass. Batching +// amortizes the O(n) sweep across ~maxEntries/2 inserts instead of running it +// on every Set once the cache is hot. Callers must hold c.mu. On return the +// cache holds fewer than maxEntries entries, leaving room for the insert. +func (c *chatResponseCache) evictLocked(now time.Time) { + for key, entry := range c.entries { + if !entry.ExpiresAt.After(now) { + delete(c.entries, key) + } + } + if len(c.entries) < c.maxEntries { + return + } + // ExpiresAt is insertion time plus a constant TTL, so ordering by ExpiresAt + // is insertion order. Keep the newest maxEntries/2 entries, drop the rest. + type keyExpiry struct { + key string + expires time.Time + } + live := make([]keyExpiry, 0, len(c.entries)) + for key, entry := range c.entries { + live = append(live, keyExpiry{key: key, expires: entry.ExpiresAt}) + } + sort.Slice(live, func(i, j int) bool { return live[i].expires.Before(live[j].expires) }) + for i := 0; i < len(live)-c.maxEntries/2; i++ { + delete(c.entries, live[i].key) + } +} + func serveCachedChatResponse(w http.ResponseWriter, r *http.Request, entry cachedChatResponse) { if rid, ok := requestLogFromContext(r.Context()); ok { w.Header().Set("X-Request-Id", rid) diff --git a/devshard/cmd/devshardctl/response_cache_test.go b/devshard/cmd/devshardctl/response_cache_test.go new file mode 100644 index 0000000000..8326347882 --- /dev/null +++ b/devshard/cmd/devshardctl/response_cache_test.go @@ -0,0 +1,135 @@ +package main + +import ( + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func cacheTestEntry(escrow, body string, expires time.Time) cachedChatResponse { + return cachedChatResponse{ + EscrowID: escrow, + Body: []byte(body), + ExpiresAt: expires, + } +} + +// cacheLen reports the entry count without depending on any exported accessor, +// so this bug-fix test stays self-contained. +func cacheLen(c *chatResponseCache) int { + c.mu.Lock() + defer c.mu.Unlock() + return len(c.entries) +} + +func TestChatResponseCache_SetGetRoundTrip(t *testing.T) { + c := newChatResponseCache(time.Hour) + now := time.Unix(1000, 0) + key := chatCacheKey("model-a", []byte("body-a")) + + c.Set(key, cacheTestEntry("escrow-1", "hello", time.Time{}), now) + + got, ok := c.Get(key, now) + require.True(t, ok) + require.Equal(t, "escrow-1", got.EscrowID) + require.Equal(t, "hello", string(got.Body)) + + // The returned body must be a copy: mutating it cannot corrupt the cache. + got.Body[0] = 'H' + again, ok := c.Get(key, now) + require.True(t, ok) + require.Equal(t, "hello", string(again.Body)) +} + +func TestChatResponseCache_ExpiredEntryNotServed(t *testing.T) { + c := newChatResponseCache(time.Hour) + now := time.Unix(1000, 0) + key := chatCacheKey("m", []byte("b")) + + c.Set(key, cacheTestEntry("e", "x", now.Add(time.Minute)), now) + + _, ok := c.Get(key, now.Add(2*time.Minute)) + require.False(t, ok) + require.Equal(t, 0, cacheLen(c), "expired entry should be dropped on access") +} + +// TestChatResponseCache_BoundsEntryCount is the regression test for the +// unbounded-growth bug: streaming unique request bodies (the normal chat +// workload) must never let the cache exceed its configured maximum. +func TestChatResponseCache_BoundsEntryCount(t *testing.T) { + c := newChatResponseCache(time.Hour) + c.maxEntries = 8 + base := time.Unix(1000, 0) + + const inserted = 500 + for i := 0; i < inserted; i++ { + now := base.Add(time.Duration(i) * time.Millisecond) + key := chatCacheKey("m", []byte(fmt.Sprintf("body-%d", i))) + c.Set(key, cacheTestEntry("e", fmt.Sprintf("resp-%d", i), time.Time{}), now) + require.LessOrEqualf(t, cacheLen(c), c.maxEntries, + "cache exceeded max entries (%d) after %d inserts", c.maxEntries, i+1) + } + + // The most recent insert must survive: eviction drops the oldest, not the + // newest, so a just-cached response is still available for a fast retry. + lastNow := base.Add(time.Duration(inserted-1) * time.Millisecond) + lastKey := chatCacheKey("m", []byte(fmt.Sprintf("body-%d", inserted-1))) + got, ok := c.Get(lastKey, lastNow) + require.True(t, ok) + require.Equal(t, fmt.Sprintf("resp-%d", inserted-1), string(got.Body)) +} + +// TestChatResponseCache_EvictionReclaimsExpiredBeforeLive verifies that +// eviction frees expired entries before touching live ones, so a burst of +// stale entries does not evict a fresh, still-useful response. +func TestChatResponseCache_EvictionReclaimsExpiredBeforeLive(t *testing.T) { + c := newChatResponseCache(time.Hour) + c.maxEntries = 4 + base := time.Unix(1000, 0) + + for i := 0; i < 4; i++ { + key := chatCacheKey("m", []byte(fmt.Sprintf("old-%d", i))) + c.Set(key, cacheTestEntry("e", "old", base.Add(10*time.Second)), base) + } + require.Equal(t, 4, cacheLen(c)) + + // Insert a fresh entry after the old ones have expired. Eviction should + // reclaim the expired entries instead of dropping the newcomer. + later := base.Add(20 * time.Second) + freshKey := chatCacheKey("m", []byte("fresh")) + c.Set(freshKey, cacheTestEntry("e", "fresh", time.Time{}), later) + + require.LessOrEqual(t, cacheLen(c), c.maxEntries) + got, ok := c.Get(freshKey, later) + require.True(t, ok) + require.Equal(t, "fresh", string(got.Body)) +} + +// TestChatResponseCache_OverwriteDoesNotEvict confirms that re-Setting an +// existing key updates in place without triggering eviction (the entry count +// does not grow, so a full cache is not needlessly churned). +func TestChatResponseCache_OverwriteDoesNotEvict(t *testing.T) { + c := newChatResponseCache(time.Hour) + c.maxEntries = 2 + now := time.Unix(1000, 0) + + k1 := chatCacheKey("m", []byte("one")) + k2 := chatCacheKey("m", []byte("two")) + c.Set(k1, cacheTestEntry("e", "one", time.Time{}), now) + c.Set(k2, cacheTestEntry("e", "two", time.Time{}), now) + require.Equal(t, 2, cacheLen(c)) + + // Overwriting k1 keeps the cache at capacity and must not evict k2. + c.Set(k1, cacheTestEntry("e", "one-v2", time.Time{}), now.Add(time.Second)) + require.Equal(t, 2, cacheLen(c)) + + got, ok := c.Get(k2, now) + require.True(t, ok) + require.Equal(t, "two", string(got.Body)) + + got, ok = c.Get(k1, now) + require.True(t, ok) + require.Equal(t, "one-v2", string(got.Body)) +}