fix(router): fix pubsub goroutine leak

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

@@ -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)
 			}
 		})

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 {