fix: correct pubsub provider lifecycle

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

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,36 +2296,24 @@ 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()
 
 	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() {

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 {