diff --git a/router-tests/lifecycle/shutdown_test.go b/router-tests/lifecycle/shutdown_test.go index 36691a2146..2fb37e7ba7 100644 --- a/router-tests/lifecycle/shutdown_test.go +++ b/router-tests/lifecycle/shutdown_test.go @@ -2,6 +2,8 @@ package integration import ( "context" + "net/http" + "sync" "syscall" "testing" "time" @@ -11,7 +13,10 @@ import ( "github.com/wundergraph/cosmo/router-tests/testenv" "github.com/wundergraph/cosmo/router-tests/testutils" "github.com/wundergraph/cosmo/router/core" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/controlplane/configpoller" + "github.com/wundergraph/cosmo/router/pkg/routerconfig" "go.uber.org/goleak" ) @@ -68,3 +73,188 @@ func TestShutdownGoroutineLeaks(t *testing.T) { require.Nil(t, res) } } + +const blockRequestHeader = "x-block-request-id" + +// blockingRequestModule holds any request carrying the blockRequestHeader open +// inside the graph mux middleware chain until the test releases it. Router +// middlewares are mounted after the in-flight counting middleware in +// buildGraphMux, so a parked request counts as in-flight for as long as it is +// blocked. +type blockingRequestModule struct { + // entered receives the request's block id once the request is parked, + // i.e. after the in-flight counter has been incremented. + entered chan string + // release maps a block id to a channel that, once closed, lets the + // request continue. Requests with an unknown id pass through unblocked. + release map[string]chan struct{} +} + +func (m *blockingRequestModule) Middleware(ctx core.RequestContext, next http.Handler) { + if id := ctx.Request().Header.Get(blockRequestHeader); id != "" { + if ch, ok := m.release[id]; ok { + m.entered <- id + <-ch + } + } + next.ServeHTTP(ctx.ResponseWriter(), ctx.Request()) +} + +func (m *blockingRequestModule) Module() core.ModuleInfo { + return core.ModuleInfo{ + ID: "blockingRequestModule", + Priority: 1, + // Return the shared instance so the test keeps access to the + // channels; a fresh instance would lose them to mapstructure's + // zero-value decoding. + New: func() core.Module { return m }, + } +} + +var _ core.RouterMiddlewareHandler = (*blockingRequestModule)(nil) + +// When a hot reload's Changes leave a graph untouched, the new graph server +// reuses that graph's mux and traffic keeps flowing through it. The replaced +// server's shutdown drain waits for requests running in the muxes it tears +// down; requests served by the new server through a reused mux are not the +// replaced server's to wait for. +func TestGraphServerShutdown(t *testing.T) { + t.Parallel() + + t.Run("waits only for requests in muxes it tears down when a hot reload reuses others", func(t *testing.T) { + t.Parallel() + + const ff1 = "experiment-a" + + initial := buildHelloRouterConfig("v1", "Base v1", map[string]string{ff1: "FF1 v1"}) + poller := newFakeConfigPoller(initial) + + mod := &blockingRequestModule{ + entered: make(chan string, 2), + release: map[string]chan struct{}{ + "pre-swap": make(chan struct{}), + "post-swap": make(chan struct{}), + }, + } + + // Release the parked requests exactly once no matter which path the + // test takes: a request still parked at teardown would leak its + // goroutine (goleak) and stall the router's shutdown for the full + // shutdown delay. + releasePreSwap := sync.OnceFunc(func() { close(mod.release["pre-swap"]) }) + releasePostSwap := sync.OnceFunc(func() { close(mod.release["post-swap"]) }) + t.Cleanup(releasePreSwap) + t.Cleanup(releasePostSwap) + + testenv.Run(t, &testenv.Config{ + RouterConfig: &testenv.RouterConfig{ + ConfigPollerFactory: func(_ *nodev1.RouterConfig) configpoller.ConfigPoller { + return poller + }, + }, + RouterOptions: []core.Option{ + core.WithConfigVersionHeader(true), + core.WithCustomModules(mod), + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + type result struct { + res *testenv.TestResponse + err error + } + + startBlockedRequest := func(id string, header http.Header) <-chan result { + header = header.Clone() + if header == nil { + header = http.Header{} + } + header.Set(blockRequestHeader, id) + return testenv.Go(func() result { + res, err := xEnv.MakeGraphQLRequest(testenv.GraphQLRequest{ + Query: `query { hello }`, + Header: header, + }) + return result{res, err} + }) + } + + // Step 1: the initial server serves the base graph at v1. + res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{Query: `query { hello }`}) + require.Equal(t, "v1", res.Response.Header.Get("X-Router-Config-Version")) + + // Step 2: park a request on the ff1 mux. The reload below changes + // ff1, so this mux is torn down by the old server's shutdown — + // the drain must wait for this request. + preSwapDone := startBlockedRequest("pre-swap", http.Header{"X-Feature-Flag": []string{ff1}}) + testenv.AwaitChannelWithT(t, 10*time.Second, mod.entered, func(t *testing.T, got string) { + require.Equal(t, "pre-swap", got) + }, "pre-swap request never reached the blocking middleware") + + // Step 3: emit a reload that changes only the feature flag. The + // base graph is unchanged, so the new server reuses the base mux. + // Emit runs the swap synchronously and blocks until the old + // server's shutdown completes, so it runs on its own goroutine. + next := buildHelloRouterConfig("v2", "Base v1", map[string]string{ff1: "FF1 v2"}) + emitDone := testenv.Go(func() error { + return poller.Emit(t, next, &routerconfig.Changes{ + AddedConfigs: map[string]struct{}{}, + RemovedConfigs: map[string]struct{}{}, + ChangedConfigs: map[string]struct{}{ff1: {}}, + }) + }) + + // Step 4: wait until the new server has been swapped in, observed + // via the rebuilt feature-flag mux serving the v2 version header. + require.EventuallyWithT(t, func(c *assert.CollectT) { + ffRes, err := xEnv.MakeGraphQLRequest(testenv.GraphQLRequest{ + Query: `query { hello }`, + Header: http.Header{"X-Feature-Flag": []string{ff1}}, + }) + if assert.NoError(c, err) { + assert.Equal(c, "v2-"+ff1, ffRes.Response.Header.Get("X-Router-Config-Version")) + } + }, 10*time.Second, 50*time.Millisecond, + "new graph server must be swapped in and serving the rebuilt feature-flag mux") + + // The old server still owns the parked pre-swap request, so its + // shutdown (and with it the Emit call) must not have finished. + select { + case err := <-emitDone: + t.Fatalf("config update returned while a request was still in flight in a mux the old server tears down (err: %v)", err) + default: + } + + // Step 5: park a second request on the base graph. It is served + // by the new server through the reused base mux — the old server + // does not own it and must not wait for it. + postSwapDone := startBlockedRequest("post-swap", nil) + testenv.AwaitChannelWithT(t, 10*time.Second, mod.entered, func(t *testing.T, got string) { + require.Equal(t, "post-swap", got) + }, "post-swap request never reached the blocking middleware") + + // Step 6: release the old server's own request. From this moment + // every request in the muxes the old server tears down is done. + releasePreSwap() + testenv.AwaitChannelWithT(t, 10*time.Second, preSwapDone, func(t *testing.T, preSwap result) { + require.NoError(t, preSwap.err) + require.Equal(t, "v1-"+ff1, preSwap.res.Response.Header.Get("X-Router-Config-Version"), + "the pre-swap request was served by the old server's ff1 mux") + }, "pre-swap request did not complete after being released") + + // Step 7: the old server's shutdown must now complete even though + // the post-swap request is still running through the reused mux. + testenv.AwaitChannelWithT(t, 10*time.Second, emitDone, func(t *testing.T, err error) { + require.NoError(t, err) + }, "old graph server shutdown did not complete after all requests in the muxes it tears down "+ + "had finished; it must not wait on requests served by the new server through a reused mux") + + // Step 8: the post-swap request is still healthy on the new + // server; release it and verify it completes. + releasePostSwap() + testenv.AwaitChannelWithT(t, 10*time.Second, postSwapDone, func(t *testing.T, postSwap result) { + require.NoError(t, postSwap.err) + require.Equal(t, "v1", postSwap.res.Response.Header.Get("X-Router-Config-Version"), + "the post-swap request was served through the reused base mux, whose version header is baked in from v1") + }, "post-swap request did not complete after being released") + }) + }) +} diff --git a/router/core/graph_server.go b/router/core/graph_server.go index d7c324b8ba..f9298913ad 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -95,9 +95,6 @@ type ( baseOtelAttributes []attribute.KeyValue baseRouterConfigVersion string mux *chi.Mux - // inFlightRequests is used to track the number of requests currently being processed - // does not include websocket (hijacked) connections. - inFlightRequests *atomic.Int64 // graphMuxList contains all graph muxes of this graph server. // It's keyed by mux name (feature flag name or empty string for base graph). graphMuxList map[string]*graphMux @@ -225,7 +222,6 @@ func newGraphServer(routerCtx context.Context, r *Router, response *routerconfig playgroundHandler: r.playgroundHandler, traceDialer: traceDialer, baseRouterConfigVersion: response.Config.GetVersion(), - inFlightRequests: &atomic.Int64{}, graphMuxList: make(map[string]*graphMux, 1), instanceData: InstanceData{ HostName: r.hostName, @@ -700,6 +696,13 @@ type graphMux struct { reused atomic.Bool finalized atomic.Bool + // inFlightRequests tracks the number of requests currently being processed + // by this mux. Does not include subscriptions or websocket (hijacked) + // connections. A reused mux keeps serving under the next graph server, so + // the count belongs to the mux and is only drained by the server that + // tears the mux down. + inFlightRequests atomic.Int64 + planCache *ristretto.Cache[uint64, *planWithMetaData] planFallbackCache *slowplancache.Cache[*planWithMetaData] persistedOperationCache *ristretto.Cache[uint64, NormalizationCacheEntry] @@ -1948,11 +1951,11 @@ func (s *graphServer) buildGraphMux( // We don't want to count any type of subscriptions e.g. SSE as in-flight requests because they are long-lived if requestContext != nil && requestContext.operation != nil && requestContext.operation.opType != OperationTypeSubscription { - s.inFlightRequests.Add(1) + gm.inFlightRequests.Add(1) // Counting like this is safe because according to the go http.ServeHTTP documentation // the requests is guaranteed to be finished when ServeHTTP returns - defer s.inFlightRequests.Add(-1) + defer gm.inFlightRequests.Add(-1) } handler.ServeHTTP(w, r) @@ -2157,6 +2160,24 @@ func newGRPCStartupParams(traceConfig *rtrace.Config, ipAnonymization *IPAnonymi return startupConfig } +// inFlightOwnedRequests sums the in-flight requests of the muxes this server +// will actually shut down. Muxes flagged as reused are inherited by the next +// server and keep serving new traffic, so their in-flight requests are not +// this server's to wait for. +func (s *graphServer) inFlightOwnedRequests() int64 { + s.graphMuxListLock.Lock() + defer s.graphMuxListLock.Unlock() + + var n int64 + for _, gm := range s.graphMuxList { + if gm.reused.Load() { + continue + } + n += gm.inFlightRequests.Load() + } + return n +} + // wait waits for all in-flight requests to finish. Similar to http.Server.Shutdown we wait in intervals + jitter // to make the shutdown process more efficient. func (s *graphServer) wait(ctx context.Context) error { @@ -2166,7 +2187,7 @@ func (s *graphServer) wait(ctx context.Context) error { defer timer.Stop() for { - if s.inFlightRequests.Load() == 0 { + if s.inFlightOwnedRequests() == 0 { return nil } select { diff --git a/router/core/graph_server_test.go b/router/core/graph_server_test.go index 066a17e187..fe26ee1274 100644 --- a/router/core/graph_server_test.go +++ b/router/core/graph_server_test.go @@ -825,7 +825,6 @@ func TestGraphServerShutdown(t *testing.T) { prev := &graphServer{ Config: &Config{logger: zap.NewNop()}, graphServerCancel: cancel, - inFlightRequests: &atomic.Int64{}, baseTransport: &http.Transport{}, graphMuxList: map[string]*graphMux{"": baseMux}, } @@ -854,7 +853,6 @@ func TestGraphServerShutdown(t *testing.T) { srv := &graphServer{ Config: &Config{logger: zap.NewNop()}, graphServerCancel: cancel, - inFlightRequests: &atomic.Int64{}, baseTransport: &http.Transport{}, graphMuxList: map[string]*graphMux{"": mux}, }