Bound Gateway cache lifetime and capacity

internal/protocol/cache_lifetime_test.go

@@ -0,0 +1,63 @@
+package protocol
+
+import (
+	"fmt"
+	"net/http"
+	"testing"
+	"time"
+)
+
+func TestGatewayCacheEntriesExpireAndInvalidateTogether(t *testing.T) {
+	gateway := &Gateway{
+		mcpToolsCache:  map[string][]map[string]any{"expired": {{"name": "tool"}}},
+		connectCache:   map[string]http.Handler{"expired": http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})},
+		mcpCacheAt:     map[string]time.Time{"expired": time.Now().Add(-gatewayCacheTTL)},
+		connectCacheAt: map[string]time.Time{"expired": time.Now().Add(-gatewayCacheTTL)},
+	}
+	gateway.mu.Lock()
+	gateway.evictExpiredCachesLocked(time.Now())
+	gateway.mu.Unlock()
+	if len(gateway.mcpToolsCache) != 0 || len(gateway.connectCache) != 0 {
+		t.Fatalf("expired gateway cache entries remain: mcp=%d connect=%d", len(gateway.mcpToolsCache), len(gateway.connectCache))
+	}
+}
+
+func TestGatewayPruneRemovesOldestEntriesWithinCapacity(t *testing.T) {
+	gateway := &Gateway{mcpToolsCache: make(map[string][]map[string]any), mcpCacheAt: make(map[string]time.Time)}
+	for i := 0; i < gatewayCacheMaxEntries+1; i++ {
+		key := fmt.Sprintf("key-%d", i)
+		gateway.mcpToolsCache[key] = []map[string]any{{"name": key}}
+		gateway.mcpCacheAt[key] = time.Unix(int64(i), 0)
+	}
+	gateway.pruneGatewayCachesLocked()
+	if len(gateway.mcpToolsCache) != gatewayCacheMaxEntries {
+		t.Fatalf("cache entries = %d, want %d", len(gateway.mcpToolsCache), gatewayCacheMaxEntries)
+	}
+	if _, ok := gateway.mcpToolsCache["key-0"]; ok {
+		t.Fatal("oldest cache entry was not pruned")
+	}
+}
+
+func TestGatewayPruneAdvancesWhenTimestampMetadataIsMissing(t *testing.T) {
+	gateway := &Gateway{mcpToolsCache: make(map[string][]map[string]any)}
+	for i := 0; i < gatewayCacheMaxEntries+1; i++ {
+		gateway.mcpToolsCache[fmt.Sprintf("key-%d", i)] = nil
+	}
+	gateway.pruneGatewayCachesLocked()
+	if len(gateway.mcpToolsCache) != gatewayCacheMaxEntries {
+		t.Fatalf("cache entries = %d, want %d", len(gateway.mcpToolsCache), gatewayCacheMaxEntries)
+	}
+}
+
+func TestGatewayInstanceInvalidationClearsSchemaCaches(t *testing.T) {
+	gateway := &Gateway{
+		mcpToolsCache:  map[string][]map[string]any{"cached": {{"name": "tool"}}},
+		connectCache:   map[string]http.Handler{"cached": http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})},
+		mcpCacheAt:     map[string]time.Time{"cached": time.Now()},
+		connectCacheAt: map[string]time.Time{"cached": time.Now()},
+	}
+	gateway.InvalidateInstance("missing-instance")
+	if len(gateway.mcpToolsCache) != 0 || len(gateway.connectCache) != 0 {
+		t.Fatal("instance invalidation left stale schema caches")
+	}
+}

internal/protocol/gateway.go

@@ -53,13 +53,19 @@ type Gateway struct {
 	AccessLogger accessLogger
 	Logger       *slog.Logger
 
-	mu            sync.Mutex
-	conns         map[string]*grpc.ClientConn
-	mcpToolsCache map[string][]map[string]any
-	connectCache  map[string]http.Handler
-}
-
-const DefaultMaxRequestBytes int64 = 1 << 20
+	mu             sync.Mutex
+	conns          map[string]*grpc.ClientConn
+	mcpToolsCache  map[string][]map[string]any
+	connectCache   map[string]http.Handler
+	mcpCacheAt     map[string]time.Time
+	connectCacheAt map[string]time.Time
+}
+
+const (
+	DefaultMaxRequestBytes int64 = 1 << 20
+	gatewayCacheMaxEntries       = 1024
+	gatewayCacheTTL              = 10 * time.Minute
+)
 
 type Catalog struct {
 	CapsetID    string                  `json:"capset_id"`
@@ -878,6 +884,7 @@ func (g *Gateway) mcpTools(ctx context.Context, capsetID string) ([]map[string]a
 	}
 	cacheKey := mcpToolsCacheKey(capsetID, items)
 	g.mu.Lock()
+	evictExpiredCacheEntryLocked(g.mcpToolsCache, g.mcpCacheAt, cacheKey, time.Now())
 	if cached := g.mcpToolsCache[cacheKey]; cached != nil {
 		g.mu.Unlock()
 		return cloneToolList(cached), nil
@@ -923,7 +930,13 @@ func (g *Gateway) mcpTools(ctx context.Context, capsetID string) ([]map[string]a
 	if g.mcpToolsCache == nil {
 		g.mcpToolsCache = map[string][]map[string]any{}
 	}
+	g.evictExpiredCachesLocked(time.Now())
 	g.mcpToolsCache[cacheKey] = cloneToolList(tools)
+	if g.mcpCacheAt == nil {
+		g.mcpCacheAt = map[string]time.Time{}
+	}
+	g.mcpCacheAt[cacheKey] = time.Now()
+	g.pruneGatewayCachesLocked()
 	g.mu.Unlock()
 	return cloneToolList(tools), nil
 }
@@ -1253,6 +1266,7 @@ type connectExposedMethodKey struct{}
 func (g *Gateway) connectHandler(item store.ExposedMethod) (http.Handler, error) {
 	key := connectHandlerCacheKey(item)
 	g.mu.Lock()
+	evictExpiredCacheEntryLocked(g.connectCache, g.connectCacheAt, key, time.Now())
 	if g.connectCache != nil {
 		if handler := g.connectCache[key]; handler != nil {
 			g.mu.Unlock()
@@ -1300,7 +1314,13 @@ func (g *Gateway) connectHandler(item store.ExposedMethod) (http.Handler, error)
 	if g.connectCache == nil {
 		g.connectCache = map[string]http.Handler{}
 	}
+	g.evictExpiredCachesLocked(time.Now())
 	g.connectCache[key] = handler
+	if g.connectCacheAt == nil {
+		g.connectCacheAt = map[string]time.Time{}
+	}
+	g.connectCacheAt[key] = time.Now()
+	g.pruneGatewayCachesLocked()
 	g.mu.Unlock()
 	return handler, nil
 }
@@ -1802,6 +1822,67 @@ func (g *Gateway) InvalidateInstance(instanceID string) {
 		_ = conn.Close()
 		delete(g.conns, key)
 	}
+	clear(g.mcpToolsCache)
+	clear(g.connectCache)
+	clear(g.mcpCacheAt)
+	clear(g.connectCacheAt)
+}
+
+func evictExpiredCacheEntryLocked[T any](cache map[string]T, created map[string]time.Time, key string, now time.Time) {
+	if at, ok := created[key]; ok && now.Sub(at) >= gatewayCacheTTL {
+		delete(created, key)
+		delete(cache, key)
+	}
+}
+
+func (g *Gateway) evictExpiredCachesLocked(now time.Time) {
+	for key, created := range g.mcpCacheAt {
+		if now.Sub(created) >= gatewayCacheTTL {
+			delete(g.mcpCacheAt, key)
+			delete(g.mcpToolsCache, key)
+		}
+	}
+	for key, created := range g.connectCacheAt {
+		if now.Sub(created) >= gatewayCacheTTL {
+			delete(g.connectCacheAt, key)
+			delete(g.connectCache, key)
+		}
+	}
+}
+
+func (g *Gateway) pruneGatewayCachesLocked() {
+	for len(g.mcpToolsCache) > gatewayCacheMaxEntries {
+		if !deleteOldestCache(g.mcpToolsCache, g.mcpCacheAt) {
+			break
+		}
+	}
+	for len(g.connectCache) > gatewayCacheMaxEntries {
+		if !deleteOldestCache(g.connectCache, g.connectCacheAt) {
+			break
+		}
+	}
+}
+
+func deleteOldestCache[T any](cache map[string]T, created map[string]time.Time) bool {
+	var oldestKey string
+	var oldest time.Time
+	for key, at := range created {
+		if oldestKey == "" || at.Before(oldest) {
+			oldestKey = key
+			oldest = at
+		}
+	}
+	if oldestKey == "" {
+		for key := range cache {
+			delete(cache, key)
+			delete(created, key)
+			return true
+		}
+		return false
+	}
+	delete(cache, oldestKey)
+	delete(created, oldestKey)
+	return true
 }
 
 func (g *Gateway) Close() error {
@@ -1814,6 +1895,10 @@ func (g *Gateway) Close() error {
 		}
 		delete(g.conns, key)
 	}
+	clear(g.mcpToolsCache)
+	clear(g.connectCache)
+	clear(g.mcpCacheAt)
+	clear(g.connectCacheAt)
 	return err
 }