Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions internal/protocol/cache_lifetime_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
package protocol

import (
"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 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")
}
}
79 changes: 72 additions & 7 deletions internal/protocol/gateway.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
Expand Down Expand Up @@ -878,6 +884,7 @@ func (g *Gateway) mcpTools(ctx context.Context, capsetID string) ([]map[string]a
}
cacheKey := mcpToolsCacheKey(capsetID, items)
g.mu.Lock()
g.evictExpiredCachesLocked(time.Now())

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

每次缓存访问都在全局锁下全量扫描缓存做过期淘汰,热路径由 O(1) 退化为 O(n)

mcpTools(887 行)与 connectHandler(1268 行)在每次调用、持有全局 g.mu 时都执行 evictExpiredCachesLocked,该函数会完整遍历 mcpCacheAt 与 connectCacheAt 两张 map(每张最多 gatewayCacheMaxEntries=1024 项)。即使缓存全部命中、没有任何过期项,每次请求也要做最多约 2048 次时间比较并串行化在全局互斥锁内。改动前热路径只是 O(1) 的 map 读取;改动后所有网关请求在全局锁下承担 O(cache) 的清扫成本,高 QPS 下会放大锁竞争与请求延迟,且该成本与缓存命中与否无关。

Problem code:

Changed code at internal/protocol/gateway.go:887

Recommendation:
将过期淘汰从每次访问的全量 O(n) 扫描改为摊销/懒淘汰:读取时只检查目标 cacheKey 自身的过期时间,命中即返回;仅在写入新条目或周期性(每 N 次访问/后台定时器)时才做全量清扫。若保留全量扫描,至少将其移出读路径。

if cached := g.mcpToolsCache[cacheKey]; cached != nil {
g.mu.Unlock()
return cloneToolList(cached), nil
Expand Down Expand Up @@ -924,6 +931,11 @@ func (g *Gateway) mcpTools(ctx context.Context, capsetID string) ([]map[string]a
g.mcpToolsCache = map[string][]map[string]any{}
}
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
}
Expand Down Expand Up @@ -1253,6 +1265,7 @@ type connectExposedMethodKey struct{}
func (g *Gateway) connectHandler(item store.ExposedMethod) (http.Handler, error) {
key := connectHandlerCacheKey(item)
g.mu.Lock()
g.evictExpiredCachesLocked(time.Now())
if g.connectCache != nil {
if handler := g.connectCache[key]; handler != nil {
g.mu.Unlock()
Expand Down Expand Up @@ -1301,6 +1314,11 @@ func (g *Gateway) connectHandler(item store.ExposedMethod) (http.Handler, error)
g.connectCache = map[string]http.Handler{}
}
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
}
Expand Down Expand Up @@ -1802,6 +1820,49 @@ 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 (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 {
deleteOldestCache(g.mcpToolsCache, g.mcpCacheAt)
}
for len(g.connectCache) > gatewayCacheMaxEntries {
deleteOldestCache(g.connectCache, g.connectCacheAt)
}
}

func deleteOldestCache[T any](cache map[string]T, created map[string]time.Time) {
var oldestKey string
var oldest time.Time
for key, at := range created {
if oldestKey == "" || at.Before(oldest) {
oldestKey = key
oldest = at
}
}
if oldestKey != "" {
delete(cache, oldestKey)
delete(created, oldestKey)
}
Comment thread
monkeyscan[bot] marked this conversation as resolved.
}

func (g *Gateway) Close() error {
Expand All @@ -1814,6 +1875,10 @@ func (g *Gateway) Close() error {
}
delete(g.conns, key)
}
clear(g.mcpToolsCache)
clear(g.connectCache)
clear(g.mcpCacheAt)
clear(g.connectCacheAt)
return err
}

Expand Down
Loading