Skip to content

fix(router): keep websockets + resolver alive for reused graph muxes #2838

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dkorittki merged 5 commits into from
May 7, 2026
Merged

fix(router): keep websockets + resolver alive for reused graph muxes #2838

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
605a246
fix: keep websocket connections alive for reused graph muxes
dkorittki May 7, 2026
9a72b90
chore: add tests
dkorittki May 7, 2026
1aca0f1
fix: reset router config
dkorittki May 7, 2026
8d55dc9
fix: cancel context before closing caches, etc.
dkorittki May 7, 2026
6f21c06
chore: move comment into subtest
dkorittki May 7, 2026
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
113 changes: 113 additions & 0 deletions router-tests/protocol/config_hot_reload_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (

"context"
"encoding/json"
"net/http"
"os"
"sync/atomic"
"testing"
Expand Down Expand Up @@ -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)
Comment on lines +625 to +630

@coderabbitai coderabbitai Bot May 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Use testenv.WSWriteJSON for the subscribe handshake.

This is still the success path, so the new hot-reload test should use the websocket helper rather than conn.WriteJSON directly. That keeps the test aligned with the repo’s retry/deadline behavior and reduces flakes.

As per coding guidelines: "Use testenv.WSReadJSON and testenv.WSWriteJSON for WebSocket reads and writes in tests instead of conn.ReadJSON and conn.WriteJSON, as these helpers include retry logic with 2-second deadlines and exponential backoff."

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@router-tests/protocol/config_hot_reload_test.go` around lines 625 - 630,
Replace the direct websocket write call conn.WriteJSON(...) used for the
subscription handshake with the test helper testenv.WSWriteJSON so the test
benefits from the built-in retry/timeout behavior; locate the call that
constructs a testenv.WebSocketMessage (ID "1", Type "subscribe", Payload ...)
and change the invocation to testenv.WSWriteJSON(t, conn,
&testenv.WebSocketMessage{...}) while keeping the same message fields and error
assertions intact (remove the manual require.NoError on conn.WriteJSON since
WSWriteJSON already handles errors).

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()

Expand Down
61 changes: 39 additions & 22 deletions router/core/graph_server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Comment thread
dkorittki marked this conversation as resolved.
Config: &r.Config,
engineStats: r.EngineStats,
baseTransport: baseTransport,
Expand Down Expand Up @@ -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(),
Expand All @@ -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()
Expand All @@ -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,
Expand Down Expand Up @@ -488,7 +492,6 @@ func getRoutingUrlGroupingForCircuitBreakers(
}

func (s *graphServer) buildMultiGraphHandler(
ctx context.Context,
opts buildMultiGraphHandlerOptions,
) (http.HandlerFunc, error) {
if len(opts.featureFlagConfigs) == 0 {
Expand All @@ -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
Expand All @@ -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(),
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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(),
}
Expand Down Expand Up @@ -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)
}

Expand Down Expand Up @@ -1369,7 +1386,7 @@ func (s *graphServer) buildGraphMux(
}

executor, providers, err := ecb.Build(
ctx,
graphMuxCtx,
&ExecutorBuildOptions{
EngineConfig: opts.EngineConfig,
Subgraphs: opts.ConfigSubgraphs,
Expand All @@ -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
}

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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))
Expand All @@ -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(
Expand All @@ -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))
}
Expand All @@ -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))
}
})
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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),
Expand Down
Loading