diff --git a/router-tests/protocol/config_hot_reload_test.go b/router-tests/protocol/config_hot_reload_test.go index 6221ebd5df..df03ecd498 100644 --- a/router-tests/protocol/config_hot_reload_test.go +++ b/router-tests/protocol/config_hot_reload_test.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" + "net/http" "os" "sync/atomic" "testing" @@ -582,6 +583,118 @@ func TestFlakyConfigHotReloadPoller(t *testing.T) { }) } +func TestConfigHotReloadGraphServerSwap(t *testing.T) { + t.Parallel() + + t.Run("verify only ws connections on swapped muxes are closed", func(t *testing.T) { + // verifies that after a config update that changes only one feature flag's mux, websocket connections on + // unchanged muxes (base graph and unmodified feature flags) remain active, while + // connections on the changed mux are closed by the server. + + pm := ConfigPollerMock{ + ready: make(chan struct{}), + } + + testenv.Run(t, &testenv.Config{ + RouterConfig: &testenv.RouterConfig{ + ConfigPollerFactory: func(cfg *nodev1.RouterConfig) configpoller.ConfigPoller { + // Add "myff2" as a second feature flag (clone of "myff") so the router + // starts with three distinct muxes: base graph, myff, and myff2. + if cfg.FeatureFlagConfigs != nil { + if myff, ok := cfg.FeatureFlagConfigs.ConfigByFeatureFlagName["myff"]; ok { + cfg.FeatureFlagConfigs.ConfigByFeatureFlagName["myff2"] = &nodev1.FeatureFlagRouterExecutionConfig{ + EngineConfig: myff.EngineConfig, + Version: "myff2-initial", + Subgraphs: myff.Subgraphs, + } + } + } + pm.initConfig = cfg + return &pm + }, + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + // Wait for the config poller to be ready before establishing connections. + <-pm.ready + + // subscribe dials a WebSocket, starts a currentTime subscription, and reads + // one initial message to confirm the subscription is live before returning. + subscribe := func(header http.Header) *websocket.Conn { + t.Helper() + conn := xEnv.InitGraphQLWebSocketConnection(header, nil, nil) + err := conn.WriteJSON(&testenv.WebSocketMessage{ + ID: "1", + Type: "subscribe", + Payload: []byte(`{"query":"subscription { currentTime { unixTime timeStamp }}"}`), + }) + require.NoError(t, err) + var msg testenv.WebSocketMessage + err = testenv.WSReadJSON(t, conn, &msg) + require.NoError(t, err) + require.Equal(t, "next", msg.Type) + return conn + } + + // Establish 2 connections on each of the three muxes (6 total). + baseConn1 := subscribe(nil) + baseConn2 := subscribe(nil) + myffConn1 := subscribe(http.Header{"X-Feature-Flag": []string{"myff"}}) + myffConn2 := subscribe(http.Header{"X-Feature-Flag": []string{"myff"}}) + myff2Conn1 := subscribe(http.Header{"X-Feature-Flag": []string{"myff2"}}) + myff2Conn2 := subscribe(http.Header{"X-Feature-Flag": []string{"myff2"}}) + + // Trigger a config update where only "myff" has changed. + // The base graph mux and "myff2" mux are unchanged and will be reused. + pm.initConfig.FeatureFlagConfigs.ConfigByFeatureFlagName["myff"].Version = "myff-v2" + require.NoError(t, pm.updateConfig(&routerconfig.Response{ + Config: pm.initConfig, + Changes: &routerconfig.Changes{ + ChangedConfigs: map[string]struct{}{"myff": {}}, + }, + })) + + // assertConnectionClosed drains any pending subscription messages then expects + // a WebSocket close error, confirming the server closed the connection. + // A single absolute deadline is set on the connection before the loop so that + // continuous data messages cannot prevent the timeout from firing. + assertConnectionClosed := func(conn *websocket.Conn) { + t.Helper() + conn.SetReadDeadline(time.Now().Add(10 * time.Second)) + defer conn.SetReadDeadline(time.Time{}) + for { + var msg testenv.WebSocketMessage + err := conn.ReadJSON(&msg) + if err != nil { + var wsErr *websocket.CloseError + require.ErrorAs(t, err, &wsErr, "expected websocket close error, got: %v", err) + return + } + // A data message arrived before the close — keep draining. + require.Equal(t, "next", msg.Type) + } + } + + // The "myff" connections must be closed because their mux was rebuilt. + assertConnectionClosed(myffConn1) + assertConnectionClosed(myffConn2) + + // The base graph and "myff2" connections must still receive data because their + // muxes were reused and their per-mux contexts were not cancelled. + for _, conn := range []*websocket.Conn{baseConn1, baseConn2, myff2Conn1, myff2Conn2} { + var msg testenv.WebSocketMessage + err := testenv.WSReadJSON(t, conn, &msg) + require.NoError(t, err) + require.Equal(t, "next", msg.Type) + } + + // Close the remaining connections. + for _, conn := range []*websocket.Conn{baseConn1, baseConn2, myff2Conn1, myff2Conn2} { + require.NoError(t, conn.Close()) + } + }) + }) +} + func writeTestConfig(t *testing.T, version string, path string) { t.Helper() diff --git a/router/core/graph_server.go b/router/core/graph_server.go index 6fc67b68ca..5994b4444d 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -80,8 +80,9 @@ type ( // All fields are shared between all feature muxes. On shutdown, all graph instances are shutdown. graphServer struct { *Config - context context.Context - cancelFunc context.CancelFunc + graphServerCtx context.Context + graphServerCancel context.CancelFunc + routerCtx context.Context storageProviders *config.StorageProviders engineStats statistics.EngineStatistics playgroundHandler func(http.Handler) http.Handler @@ -135,7 +136,7 @@ type buildMultiGraphHandlerOptions struct { } // newGraphServer creates a new server instance. -func newGraphServer(ctx context.Context, r *Router, response *routerconfig.Response, proxy ProxyFunc) (*graphServer, error) { +func newGraphServer(routerCtx context.Context, r *Router, response *routerconfig.Response, proxy ProxyFunc) (*graphServer, error) { /* Older versions of composition will not populate a compatibility version. * Currently, all "old" router execution configurations are compatible as there have been no breaking * changes. @@ -181,10 +182,11 @@ func newGraphServer(ctx context.Context, r *Router, response *routerconfig.Respo } } - ctx, cancel := context.WithCancel(ctx) + graphServerCtx, graphServerCancel := context.WithCancel(routerCtx) s := &graphServer{ - context: ctx, - cancelFunc: cancel, + graphServerCtx: graphServerCtx, + graphServerCancel: graphServerCancel, + routerCtx: routerCtx, Config: &r.Config, engineStats: r.EngineStats, baseTransport: baseTransport, @@ -312,7 +314,8 @@ func newGraphServer(ctx context.Context, r *Router, response *routerconfig.Respo if needNewBaseGraphMux { // build new base grap mux - gm, err = s.buildGraphMux(ctx, BuildGraphMuxOptions{ + s.logger.Debug("Will build a new base graph mux for new graph server") + gm, err = s.buildGraphMux(BuildGraphMuxOptions{ RouterConfigVersion: s.baseRouterConfigVersion, EngineConfig: response.Config.GetEngineConfig(), ConfigSubgraphs: response.Config.GetSubgraphs(), @@ -323,6 +326,7 @@ func newGraphServer(ctx context.Context, r *Router, response *routerconfig.Respo return nil, fmt.Errorf("failed to build base mux: %w", err) } } else { + s.logger.Debug("Will reuse old base graph mux for new graph server") gm = mux gm.reused.Store(true) s.graphMuxListLock.Lock() @@ -335,7 +339,7 @@ func newGraphServer(ctx context.Context, r *Router, response *routerconfig.Respo s.logger.Info("Feature flags enabled", zap.Strings("flags", maps.Keys(featureFlagConfigMap))) } - multiGraphHandler, err := s.buildMultiGraphHandler(ctx, buildMultiGraphHandlerOptions{ + multiGraphHandler, err := s.buildMultiGraphHandler(buildMultiGraphHandlerOptions{ baseMux: gm.mux, featureFlagConfigs: featureFlagConfigMap, reloadPersistentState: r.reloadPersistentState, @@ -488,7 +492,6 @@ func getRoutingUrlGroupingForCircuitBreakers( } func (s *graphServer) buildMultiGraphHandler( - ctx context.Context, opts buildMultiGraphHandlerOptions, ) (http.HandlerFunc, error) { if len(opts.featureFlagConfigs) == 0 { @@ -507,6 +510,8 @@ func (s *graphServer) buildMultiGraphHandler( if !hasChanged && !wasAdded { oldGraphMux, exists := opts.currentGraphMuxes[featureFlagName] if exists { + s.logger.Debug("will reuse feature flag mux for new graph server", + zap.String("flag", featureFlagName)) featureFlagToMux[featureFlagName] = oldGraphMux.mux s.graphMuxListLock.Lock() s.graphMuxList[featureFlagName] = oldGraphMux @@ -517,7 +522,10 @@ func (s *graphServer) buildMultiGraphHandler( } } - gm, err := s.buildGraphMux(ctx, BuildGraphMuxOptions{ + s.logger.Debug("will create a new feature flag mux for new graph server", + zap.String("flag", featureFlagName)) + + gm, err := s.buildGraphMux(BuildGraphMuxOptions{ FeatureFlagName: featureFlagName, RouterConfigVersion: executionConfig.GetVersion(), EngineConfig: executionConfig.GetEngineConfig(), @@ -584,6 +592,9 @@ func (s *graphServer) setupEngineStatistics(baseAttributes []attribute.KeyValue) } type graphMux struct { + ctx context.Context + cancel context.CancelFunc + mux *chi.Mux reused atomic.Bool @@ -864,6 +875,9 @@ func (s *graphMux) configureCacheMetrics(srv *graphServer, baseOtelAttributes [] } func (s *graphMux) Shutdown(ctx context.Context) error { + // cancel the graph muxes context to close its resources like websocket connections, resolvers, etc. + s.cancel() + s.planCache.Close() s.planFallbackCache.Close() s.persistedOperationCache.Close() @@ -923,10 +937,13 @@ func (s *graphMux) Shutdown(ctx context.Context) error { // It also creates a new execution plan cache for the mux. The mux is not mounted on the server. // The mux is appended internally to the graph server's list of muxes to clean up later when the server is swapped. func (s *graphServer) buildGraphMux( - ctx context.Context, opts BuildGraphMuxOptions, ) (*graphMux, error) { + graphMuxCtx, graphMuxCancel := context.WithCancel(s.routerCtx) + gm := &graphMux{ + ctx: graphMuxCtx, + cancel: graphMuxCancel, metricStore: rmetric.NewNoopMetrics(), streamMetricStore: rmetric.NewNoopStreamMetricStore(), } @@ -1320,7 +1337,7 @@ func (s *graphServer) buildGraphMux( subgraphTippers[subgraph] = subgraphTransport } - if err := s.setupConnector(ctx, opts.EngineConfig, opts.ConfigSubgraphs, telemetryAttExpressions, tracingAttExpressions); err != nil { + if err := s.setupConnector(s.graphServerCtx, opts.EngineConfig, opts.ConfigSubgraphs, telemetryAttExpressions, tracingAttExpressions); err != nil { return nil, fmt.Errorf("failed to setup plugin host: %w", err) } @@ -1369,7 +1386,7 @@ func (s *graphServer) buildGraphMux( } executor, providers, err := ecb.Build( - ctx, + graphMuxCtx, &ExecutorBuildOptions{ EngineConfig: opts.EngineConfig, Subgraphs: opts.ConfigSubgraphs, @@ -1387,7 +1404,7 @@ func (s *graphServer) buildGraphMux( } s.pubSubProviders = providers - if pubSubStartupErr := s.startupPubSubProviders(ctx); pubSubStartupErr != nil { + if pubSubStartupErr := s.startupPubSubProviders(s.graphServerCtx); pubSubStartupErr != nil { return nil, pubSubStartupErr } @@ -1458,7 +1475,7 @@ func (s *graphServer) buildGraphMux( } warmupConfig.AfterOperation = func(item *CacheWarmupOperationPlanResult) { - gm.metricStore.MeasureOperationPlanningTime(ctx, + gm.metricStore.MeasureOperationPlanningTime(graphMuxCtx, item.PlanningTime, nil, otelmetric.WithAttributes( @@ -1510,7 +1527,7 @@ func (s *graphServer) buildGraphMux( return nil, fmt.Errorf("unexpected cache warmer source provided") } - err = WarmupCaches(ctx, warmupConfig) + err = WarmupCaches(graphMuxCtx, warmupConfig) if err != nil { // We don't want to fail the server if the cache warmup fails s.logger.Error("Failed to warmup caches. It will retry after server restart or graph execution config update", zap.Error(err)) @@ -1537,7 +1554,7 @@ func (s *graphServer) buildGraphMux( }) manifestAfterOperation := func(item *CacheWarmupOperationPlanResult) { - gm.metricStore.MeasureOperationPlanningTime(ctx, + gm.metricStore.MeasureOperationPlanningTime(graphMuxCtx, item.PlanningTime, nil, otelmetric.WithAttributes( @@ -1564,7 +1581,7 @@ func (s *graphServer) buildGraphMux( AfterOperation: manifestAfterOperation, } - err = WarmupCaches(ctx, manifestWarmupConfig) + err = WarmupCaches(graphMuxCtx, manifestWarmupConfig) if err != nil { s.logger.Error("Failed to warmup PQL manifest operations", zap.Error(err)) } @@ -1583,7 +1600,7 @@ func (s *graphServer) buildGraphMux( AfterOperation: manifestAfterOperation, } - if rewarmErr := WarmupCaches(ctx, rewarmConfig); rewarmErr != nil { + if rewarmErr := WarmupCaches(graphMuxCtx, rewarmConfig); rewarmErr != nil { s.logger.Error("Failed to re-warm PQL manifest operations after update", zap.Error(rewarmErr)) } }) @@ -1711,7 +1728,7 @@ func (s *graphServer) buildGraphMux( }) if s.webSocketConfiguration != nil && s.webSocketConfiguration.Enabled { - wsMiddleware := NewWebsocketMiddleware(ctx, WebsocketMiddlewareOptions{ + wsMiddleware := NewWebsocketMiddleware(graphMuxCtx, WebsocketMiddlewareOptions{ OperationProcessor: operationProcessor, OperationBlocker: operationBlocker, Planner: operationPlanner, @@ -1971,8 +1988,8 @@ func (s *graphServer) wait(ctx context.Context) error { // Shutdown does cancel the context after all non-hijacked requests such as WebSockets has been handled. func (s *graphServer) Shutdown(ctx context.Context) error { // Cancel the context after the graceful shutdown is done - // to clean up resources like websocket connections, pools, etc. - defer s.cancelFunc() + // to clean up resources. + defer s.graphServerCancel() s.logger.Debug("Shutdown of graph server initiated. Waiting for in-flight requests to finish.", zap.String("config_version", s.baseRouterConfigVersion),