diff --git a/router-tests/events/nats_events_test.go b/router-tests/events/nats_events_test.go index f79bd6784e..648ae0d86e 100644 --- a/router-tests/events/nats_events_test.go +++ b/router-tests/events/nats_events_test.go @@ -1607,12 +1607,18 @@ func TestNatsEvents(t *testing.T) { Key: "provider_id", String: "my-nats", }) - return myNatsLogs.FilterMessage("NATS connection established").Len() == 4 && - myNatsLogs.FilterMessage("NATS disconnected").Len() == 1 && - myNatsLogs.FilterMessage("NATS connection closed").Len() == 1 && - defaultLogs.FilterMessage("NATS connection established").Len() == 4 && - defaultLogs.FilterMessage("NATS disconnected").Len() == 1 && - defaultLogs.FilterMessage("NATS connection closed").Len() == 1 + + // Verify the connection lifecycle via router logs that for the "default" and "my-nats" pubsub provider. + // 4 nats connections opened (base + ff mux get their own * 2 graph mux generations). + // 2 get closed (old mux generation shuts down). + myNatsProviderLogsComplete := myNatsLogs.FilterMessage("NATS connection established").Len() == 4 && + myNatsLogs.FilterMessage("NATS disconnected").Len() == 2 && + myNatsLogs.FilterMessage("NATS connection closed").Len() == 2 + defaultProviderLogsComplete := defaultLogs.FilterMessage("NATS connection established").Len() == 4 && + defaultLogs.FilterMessage("NATS disconnected").Len() == 2 && + defaultLogs.FilterMessage("NATS connection closed").Len() == 2 + + return myNatsProviderLogsComplete && defaultProviderLogsComplete }, EventWaitTimeout, time.Second) // Then wait for subscriptions to be started again diff --git a/router/core/graph_server.go b/router/core/graph_server.go index fa6cd1408f..2c105ffabf 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -106,7 +106,6 @@ type ( prometheusEngineMetrics *rmetric.EngineMetrics connectionMetrics *rmetric.ConnectionMetrics instanceData InstanceData - pubSubProviders []datasource.Provider traceDialer *TraceDialer connector *grpcconnector.Connector circuitBreakerManager *circuit.Manager @@ -716,6 +715,12 @@ type graphMux struct { otelCacheMetrics *rmetric.CacheMetrics streamMetricStore rmetric.StreamMetricStore prometheusMetricsExporter *graphqlmetrics.PrometheusMetricsExporter + + // pubSubProviders are the EDFS providers built for this mux. They are owned by + // the mux (not the server) so that a mux reused by the next server keeps its + // providers alive: Shutdown skips reused muxes, so their providers are not torn + // down until the mux itself is finally discarded. + pubSubProviders []datasource.Provider } // buildOperationCaches creates the caches for the graph mux. @@ -1027,6 +1032,15 @@ func (s *graphMux) Shutdown(ctx context.Context) error { } } + // Shut down the pubsub providers owned by this mux. A reused mux never reaches + // here (Shutdown skips it), so its providers stay alive for the next server. + const defaultShutdownTimeout = 5 * time.Second + if pErr := providersActionWithTimeout(ctx, s.pubSubProviders, func(ctx context.Context, provider datasource.Provider) error { + return provider.Shutdown(ctx) + }, defaultShutdownTimeout, "pubsub provider shutdown timed out"); pErr != nil { + err = errors.Join(err, pErr) + } + if err != nil { return fmt.Errorf("shutdown graph mux: %w", err) } @@ -1540,10 +1554,11 @@ func (s *graphServer) buildGraphMux( }) } - s.pubSubProviders = providers - if pubSubStartupErr := s.startupPubSubProviders(s.graphServerCtx); pubSubStartupErr != nil { + if pubSubStartupErr := s.startupPubSubProviders(s.graphServerCtx, providers); pubSubStartupErr != nil { return nil, pubSubStartupErr } + // Providers are owned by the mux so they follow its reuse/shutdown lifecycle. + gm.pubSubProviders = providers operationProcessor := NewOperationProcessor(OperationProcessorOptions{ Executor: executor, @@ -2266,10 +2281,7 @@ func (s *graphServer) Shutdown(ctx context.Context) error { subgraphTransport.CloseIdleConnections() } - // Shutdown pubsub providers - if err := s.shutdownPubSubProviders(ctx); err != nil { - finalErr = errors.Join(finalErr, err) - } + // Pubsub providers are shut down per-mux in the loop above (skipping reused muxes). if s.connector != nil { s.logger.Debug("Stopping old plugins") @@ -2284,28 +2296,16 @@ func (s *graphServer) Shutdown(ctx context.Context) error { // startupPubSubProviders starts the given pubsub providers // It returns an error if any of the providers fail to start // or if some providers takes to long to start -func (s *graphServer) startupPubSubProviders(ctx context.Context) error { +func (s *graphServer) startupPubSubProviders(ctx context.Context, providers []datasource.Provider) error { // Default timeout for pubsub provider startup const defaultStartupTimeout = 5 * time.Second - return s.providersActionWithTimeout(ctx, func(ctx context.Context, provider datasource.Provider) error { + return providersActionWithTimeout(ctx, providers, func(ctx context.Context, provider datasource.Provider) error { return provider.Startup(ctx) }, defaultStartupTimeout, "pubsub provider startup timed out") } -// shutdownPubSubProviders shuts down all pubsub providers -// It returns an error if any of the providers fail to shutdown -// or if some providers takes to long to shutdown -func (s *graphServer) shutdownPubSubProviders(ctx context.Context) error { - // Default timeout for pubsub provider shutdown - const defaultShutdownTimeout = 5 * time.Second - - return s.providersActionWithTimeout(ctx, func(ctx context.Context, provider datasource.Provider) error { - return provider.Shutdown(ctx) - }, defaultShutdownTimeout, "pubsub provider shutdown timed out") -} - -func (s *graphServer) providersActionWithTimeout(ctx context.Context, action func(ctx context.Context, provider datasource.Provider) error, timeout time.Duration, timeoutMessage string) error { +func providersActionWithTimeout(ctx context.Context, providers []datasource.Provider, action func(ctx context.Context, provider datasource.Provider) error, timeout time.Duration, timeoutMessage string) error { cancellableCtx, cancel := context.WithCancel(ctx) defer cancel() @@ -2313,7 +2313,7 @@ func (s *graphServer) providersActionWithTimeout(ctx context.Context, action fun defer timer.Stop() providersGroup := new(errgroup.Group) - for _, provider := range s.pubSubProviders { + for _, provider := range providers { providersGroup.Go(func() error { actionDone := make(chan error, 1) go func() { diff --git a/router/core/graph_server_test.go b/router/core/graph_server_test.go index aabc6ae4ff..cd9d5a110c 100644 --- a/router/core/graph_server_test.go +++ b/router/core/graph_server_test.go @@ -2,18 +2,22 @@ package core import ( "cmp" + "context" "net/http" "runtime" "slices" + "sync/atomic" "testing" "weak" + "go.uber.org/zap" + "github.com/go-chi/chi/v5" "github.com/stretchr/testify/require" 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/pubsub/datasource" "github.com/wundergraph/cosmo/router/pkg/routerconfig" - "go.uber.org/zap" ) func TestGetRoutingUrlGroupingForCircuitBreakers(t *testing.T) { @@ -790,6 +794,57 @@ func TestBuildMultiGraphHandler(t *testing.T) { }) } +// reuseTrackingProvider is a pubsub provider that only records whether Shutdown +// was called. The embedded interface satisfies datasource.Provider; every other +// method is unused in this test (and would panic on the nil interface if hit). +// shutdown is atomic because providersActionWithTimeout calls Shutdown from a +// separate goroutine. +type reuseTrackingProvider struct { + datasource.Provider + shutdown atomic.Bool +} + +func (p *reuseTrackingProvider) Shutdown(context.Context) error { + p.shutdown.Store(true) + return nil +} + +func TestGraphServerShutdown(t *testing.T) { + // When the base graph is unchanged, a hot reload reuses the previous server's + // base mux, and that mux keeps serving on the new server. Its pubsub providers + // are owned by the mux, so shutting down the previous server must leave a + // reused mux and the providers it depends on intact. + t.Run("keeps a reused mux's pubsub providers alive after the previous server shuts down", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // The base mux and its pubsub provider, built by the previous server. + provider := &reuseTrackingProvider{} + baseMux := &graphMux{mux: chi.NewMux(), pubSubProviders: []datasource.Provider{provider}} + + prev := &graphServer{ + Config: &Config{logger: zap.NewNop()}, + graphServerCancel: cancel, + inFlightRequests: &atomic.Int64{}, + baseTransport: &http.Transport{}, + graphMuxList: map[string]*graphMux{"": baseMux}, + } + + // The next server reuses the unchanged base mux, inheriting its provider + // rather than rebuilding it. + next := &graphServer{graphMuxList: map[string]*graphMux{}} + next.commitReusedMuxes([]reusedGraphMux{{key: "", mux: baseMux}}) + require.True(t, baseMux.reused.Load(), "base mux must be flagged as reused") + + require.NoError(t, prev.Shutdown(ctx)) + + // The reused mux lives on in the next server, and its provider stays up. + require.Same(t, baseMux, next.graphMuxList[""], "reused mux must live on in the next server") + require.False(t, provider.shutdown.Load(), + "provider for a reused mux must stay up after the previous server shuts down") + }) +} + func toSet[T comparable](slice ...T) map[T]bool { set := make(map[T]bool, len(slice)) for _, v := range slice {