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..37cdf2dabc 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -67,6 +67,7 @@ import ( const ( featureFlagHeader = "X-Feature-Flag" featureFlagCookie = "feature_flag" + providerTimeout = 5 * time.Second ) type ( @@ -106,7 +107,6 @@ type ( prometheusEngineMetrics *rmetric.EngineMetrics connectionMetrics *rmetric.ConnectionMetrics instanceData InstanceData - pubSubProviders []datasource.Provider traceDialer *TraceDialer connector *grpcconnector.Connector circuitBreakerManager *circuit.Manager @@ -716,6 +716,8 @@ type graphMux struct { otelCacheMetrics *rmetric.CacheMetrics streamMetricStore rmetric.StreamMetricStore prometheusMetricsExporter *graphqlmetrics.PrometheusMetricsExporter + + pubSubProviders []datasource.Provider } // buildOperationCaches creates the caches for the graph mux. @@ -976,6 +978,25 @@ func (s *graphMux) configureCacheMetrics(srv *graphServer, baseOtelAttributes [] return nil } +// addPubsubProviders appends providers to s. Currently not thread-safe. +func (s *graphMux) addPubsubProviders(providers []datasource.Provider) { + s.pubSubProviders = append(s.pubSubProviders, providers...) +} + +// startPubsubProviders starts all pubsub providers of s. +func (s *graphMux) startPubsubProviders(ctx context.Context) error { + return providersActionWithTimeout(ctx, s.pubSubProviders, func(ctx context.Context, provider datasource.Provider) error { + return provider.Startup(ctx) + }, providerTimeout, "pubsub provider startup timed out") +} + +// stopPubsubProviders stops all pubsub providers of s. +func (s *graphMux) stopPubsubProviders(ctx context.Context) error { + return providersActionWithTimeout(ctx, s.pubSubProviders, func(ctx context.Context, provider datasource.Provider) error { + return provider.Shutdown(ctx) + }, providerTimeout, "pubsub provider shutdown timed out") +} + func (s *graphMux) Shutdown(ctx context.Context) error { // Make sure we do not shutdown the mux multiple times if !s.finalized.CompareAndSwap(false, true) { @@ -1027,6 +1048,11 @@ func (s *graphMux) Shutdown(ctx context.Context) error { } } + pErr := s.stopPubsubProviders(ctx) + if pErr != nil { + err = errors.Join(err, pErr) + } + if err != nil { return fmt.Errorf("shutdown graph mux: %w", err) } @@ -1540,9 +1566,10 @@ func (s *graphServer) buildGraphMux( }) } - s.pubSubProviders = providers - if pubSubStartupErr := s.startupPubSubProviders(s.graphServerCtx); pubSubStartupErr != nil { - return nil, pubSubStartupErr + gm.addPubsubProviders(providers) + pErr := gm.startPubsubProviders(s.graphServerCtx) + if pErr != nil { + return nil, pErr } operationProcessor := NewOperationProcessor(OperationProcessorOptions{ @@ -2266,11 +2293,6 @@ 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) - } - if s.connector != nil { s.logger.Debug("Stopping old plugins") if err := s.connector.StopAllProviders(); err != nil { @@ -2281,48 +2303,21 @@ func (s *graphServer) Shutdown(ctx context.Context) error { return finalErr } -// 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 { - // Default timeout for pubsub provider startup - const defaultStartupTimeout = 5 * time.Second - - return s.providersActionWithTimeout(ctx, 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 { - cancellableCtx, cancel := context.WithCancel(ctx) +func providersActionWithTimeout(ctx context.Context, providers []datasource.Provider, action func(ctx context.Context, provider datasource.Provider) error, timeout time.Duration, timeoutMessage string) error { + timeoutCtx, cancel := context.WithTimeout(ctx, timeout) defer cancel() - timer := time.NewTimer(timeout) - 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() { - actionDone <- action(cancellableCtx, provider) + actionDone <- action(timeoutCtx, provider) }() select { case err := <-actionDone: return err - case <-timer.C: + case <-timeoutCtx.Done(): return errors.New(timeoutMessage) } }) diff --git a/router/core/graph_server_test.go b/router/core/graph_server_test.go index aabc6ae4ff..066a17e187 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,78 @@ 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") + }) + + t.Run("shuts down pubsub providers of a non-reused mux", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + provider := &reuseTrackingProvider{} + mux := &graphMux{mux: chi.NewMux(), pubSubProviders: []datasource.Provider{provider}, cancel: func() {}} + + srv := &graphServer{ + Config: &Config{logger: zap.NewNop()}, + graphServerCancel: cancel, + inFlightRequests: &atomic.Int64{}, + baseTransport: &http.Transport{}, + graphMuxList: map[string]*graphMux{"": mux}, + } + + require.NoError(t, srv.Shutdown(ctx)) + + require.True(t, provider.shutdown.Load(), + "provider for a non-reused mux must be shut down when the server shuts down") + }) +} + func toSet[T comparable](slice ...T) map[T]bool { set := make(map[T]bool, len(slice)) for _, v := range slice {