diff --git a/router/core/graph_server.go b/router/core/graph_server.go index af206997be..2866d18a51 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -681,6 +681,8 @@ type graphMux struct { mux *chi.Mux reused atomic.Bool + planCacheOnEvictEnabled atomic.Bool + planCache *ristretto.Cache[uint64, *planWithMetaData] planFallbackCache *slowplancache.Cache[*planWithMetaData] persistedOperationCache *ristretto.Cache[uint64, NormalizationCacheEntry] @@ -721,10 +723,14 @@ func (s *graphMux) buildOperationCaches(srv *graphServer) (computeSha256 bool, e BufferItems: 64, } if srv.cacheWarmup != nil && srv.cacheWarmup.Enabled && srv.cacheWarmup.InMemoryFallback { + s.planCacheOnEvictEnabled.Store(true) planCacheConfig.OnEvict = func(item *ristretto.Item[*planWithMetaData]) { // This could be called before planFallbackCache is set, but it's not a problem // because there is a nil guard inside, as well as items should not really be evicted // on startup + if !s.planCacheOnEvictEnabled.Load() { + return + } s.planFallbackCache.Set(item.Key, item.Value, item.Value.planningDuration) } } @@ -961,6 +967,11 @@ func (s *graphMux) Shutdown(ctx context.Context) error { // cancel the graph muxes context to close its resources like websocket connections, resolvers, etc. s.cancel() + // ristretto Close() invokes OnEvict for every entry; skip slow-cache migration during shutdown. + s.planCacheOnEvictEnabled.Store(false) + if s.planFallbackCache != nil { + s.planFallbackCache.Wait() + } s.planCache.Close() s.planFallbackCache.Close() s.persistedOperationCache.Close() diff --git a/router/core/operation_planner.go b/router/core/operation_planner.go index f9da57396f..e14c07bf3b 100644 --- a/router/core/operation_planner.go +++ b/router/core/operation_planner.go @@ -18,8 +18,8 @@ import ( ) type planWithMetaData struct { - preparedPlan plan.Plan - operationDocument, schemaDocument *ast.Document + preparedPlan plan.Plan + operationDocument *ast.Document typeFieldUsageInfo []*graphqlschemausage.TypeFieldUsageInfo argumentUsageInfo []*graphqlmetricsv1.ArgumentUsageInfo content string @@ -96,7 +96,6 @@ func (p *OperationPlanner) planOperation(content string, name string, includeQue return &planWithMetaData{ preparedPlan: preparedPlan, operationDocument: &doc, - schemaDocument: p.executor.RouterSchema, }, nil } diff --git a/router/core/router.go b/router/core/router.go index 5dd087acd1..7e6dff4076 100644 --- a/router/core/router.go +++ b/router/core/router.go @@ -627,6 +627,10 @@ func (r *Router) serverTLSConfig() (*tls.Config, error) { // newGraphServer creates a new server. func (r *Router) newServer(ctx context.Context, response *routerconfig.Response) error { + // Extract slow-plan cache entries before building the new graph server, which + // overwrites ReloadPersistentState cache references and before the old graphMux shuts down. + r.reloadPersistentState.OnRouterConfigReload() + server, err := newGraphServer(ctx, r, response, r.proxy) if err != nil { r.logger.Error("Failed to create graph server. Keeping the old server", zap.Error(err)) diff --git a/router/pkg/slowplancache/slow_plan_cache.go b/router/pkg/slowplancache/slow_plan_cache.go index 17fba9aa5f..cd6a0897c7 100644 --- a/router/pkg/slowplancache/slow_plan_cache.go +++ b/router/pkg/slowplancache/slow_plan_cache.go @@ -222,5 +222,14 @@ func (c *Cache[V]) Close() { // This downside is also there in ristretto (if set is called concurrently) // it is even documented in the ristretto code as a comment close(c.writeCh) + + // Drop cached entries so Close releases references to stored values + // (e.g. query plans holding schema AST pointers) before the Cache + // struct is GC'd. + c.entries.Range(func(key, _ any) bool { + c.entries.Delete(key) + return true + }) + c.size = 0 }) } diff --git a/router/pkg/slowplancache/slow_plan_cache_test.go b/router/pkg/slowplancache/slow_plan_cache_test.go index 69734bf771..dacced887d 100644 --- a/router/pkg/slowplancache/slow_plan_cache_test.go +++ b/router/pkg/slowplancache/slow_plan_cache_test.go @@ -411,6 +411,27 @@ func TestCache_DoubleClose(t *testing.T) { }) } +func TestCache_CloseReleasesEntries(t *testing.T) { + t.Parallel() + c, err := New[*testPlan](10, 0) + require.NoError(t, err) + + c.Set(1, &testPlan{content: "q1"}, 10*time.Millisecond) + c.Set(2, &testPlan{content: "q2"}, 20*time.Millisecond) + c.Set(3, &testPlan{content: "q3"}, 30*time.Millisecond) + c.Wait() + + c.Close() + + count := 0 + c.entries.Range(func(_, _ any) bool { + count++ + return true + }) + require.Equal(t, 0, count, "entries sync.Map must be empty after Close") + require.Equal(t, int64(0), c.size) +} + func BenchmarkCache_Set(b *testing.B) { c, err := New[*testPlan](1000, 0) require.NoError(b, err)