fix(router): config reload memory leak — graphMux nil + proto reset

router/core/executor.go

@@ -8,6 +8,7 @@ import (
 
 	"go.uber.org/zap"
 
+	"github.com/wundergraph/cosmo/router/pkg/mondaytweaks"
 	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/grpcconnector"
@@ -51,6 +52,19 @@ type Executor struct {
 	TrackUsageInfo  bool
 }
 
+// Close releases schema and planner references held by the executor so a replaced
+// graph mux can be garbage-collected after config reload.
+func (e *Executor) Close() {
+	if e == nil {
+		return
+	}
+	e.ClientSchema = nil
+	e.RouterSchema = nil
+	e.PlanConfig = plan.Configuration{}
+	e.RenameTypeNames = nil
+	e.Resolver = nil
+}
+
 type ExecutorBuildOptions struct {
 	EngineConfig                   *nodev1.EngineConfiguration
 	Subgraphs                      []*nodev1.Subgraph
@@ -62,10 +76,11 @@ type ExecutorBuildOptions struct {
 	TraceClientRequired            bool
 	PluginsEnabled                 bool
 	InstanceData                   InstanceData
+	WebSocketConfiguration         *config.WebSocketConfiguration
 }
 
 func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *ExecutorBuildOptions) (*Executor, []pubsub_datasource.Provider, error) {
-	planConfig, providers, err := b.buildPlannerConfiguration(ctx, opts.EngineConfig, opts.Subgraphs, opts.RouterEngineConfig, opts.PluginsEnabled)
+	planConfig, providers, err := b.buildPlannerConfiguration(ctx, opts)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to build planner configuration: %w", err)
 	}
@@ -215,29 +230,43 @@ func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *Executor
 	}, providers, nil
 }
 
-func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Context, engineConfig *nodev1.EngineConfiguration, subgraphs []*nodev1.Subgraph, routerEngineCfg *RouterEngineConfiguration, pluginsEnabled bool) (*plan.Configuration, []pubsub_datasource.Provider, error) {
+func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Context, opts *ExecutorBuildOptions) (*plan.Configuration, []pubsub_datasource.Provider, error) {
 	// this loader is used to take the engine config and create a plan config
 	// the plan config is what the engine uses to turn a GraphQL Request into an execution plan
 	// the plan config is stateful as it carries connection pools and other things
 
+	subscriptionClientOptions := b.subscriptionClientOptions
+	if subscriptionClientOptions == nil {
+		subscriptionClientOptions = &SubscriptionClientOptions{}
+	}
+	resolvedSubscriptionClientOptions := *subscriptionClientOptions
+	if mondaytweaks.UseNoopUpstreamSubscriptionClientWhenUnused {
+		resolvedSubscriptionClientOptions.UseNoopClient = shouldUseNoopUpstreamSubscriptionClient(
+			opts.EngineConfig.GetGraphqlSchema(),
+			opts.EngineConfig,
+			opts.RouterEngineConfig.Events,
+			opts.WebSocketConfiguration,
+		)
+	}
+
 	loader := NewLoader(ctx, b.trackUsageInfo, NewDefaultFactoryResolver(
 		ctx,
 		b.transportOptions,
-		b.subscriptionClientOptions,
+		&resolvedSubscriptionClientOptions,
 		b.baseTripper,
 		b.subgraphTrippers,
 		b.pluginHost,
 		b.logger,
-		routerEngineCfg.Execution.EnableNetPoll,
+		opts.RouterEngineConfig.Execution.EnableNetPoll,
 		b.instanceData,
 	), b.logger, b.subscriptionHooks)
 
 	// this generates the plan config using the data source factories from the config package
-	planConfig, providers, err := loader.Load(engineConfig, subgraphs, routerEngineCfg, pluginsEnabled)
+	planConfig, providers, err := loader.Load(opts.EngineConfig, opts.Subgraphs, opts.RouterEngineConfig, opts.PluginsEnabled)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to load configuration: %w", err)
 	}
-	debug := &routerEngineCfg.Execution.Debug
+	debug := &opts.RouterEngineConfig.Execution.Debug
 	planConfig.Debug = plan.DebugConfiguration{
 		PrintOperationTransformations: debug.PrintOperationTransformations,
 		PrintOperationEnableASTRefs:   debug.PrintOperationEnableASTRefs,
@@ -248,19 +277,19 @@ func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Con
 		PlanningVisitor:               debug.PlanningVisitor,
 		DatasourceVisitor:             debug.DatasourceVisitor,
 	}
-	planConfig.MinifySubgraphOperations = routerEngineCfg.Execution.MinifySubgraphOperations
+	planConfig.MinifySubgraphOperations = opts.RouterEngineConfig.Execution.MinifySubgraphOperations
 
-	planConfig.EnableOperationNamePropagation = routerEngineCfg.Execution.EnableSubgraphFetchOperationName
+	planConfig.EnableOperationNamePropagation = opts.RouterEngineConfig.Execution.EnableSubgraphFetchOperationName
 
-	planConfig.BuildFetchReasons = routerEngineCfg.Execution.EnableRequireFetchReasons || routerEngineCfg.Execution.ValidateRequiredExternalFields
-	planConfig.ValidateRequiredExternalFields = routerEngineCfg.Execution.ValidateRequiredExternalFields
-	planConfig.RelaxSubgraphOperationFieldSelectionMergingNullability = routerEngineCfg.Execution.RelaxSubgraphOperationFieldSelectionMergingNullability
+	planConfig.BuildFetchReasons = opts.RouterEngineConfig.Execution.EnableRequireFetchReasons || opts.RouterEngineConfig.Execution.ValidateRequiredExternalFields
+	planConfig.ValidateRequiredExternalFields = opts.RouterEngineConfig.Execution.ValidateRequiredExternalFields
+	planConfig.RelaxSubgraphOperationFieldSelectionMergingNullability = opts.RouterEngineConfig.Execution.RelaxSubgraphOperationFieldSelectionMergingNullability
 
 	// Enable cost computation when cost control is enabled
-	if routerEngineCfg.CostControl != nil && routerEngineCfg.CostControl.Enabled {
+	if opts.RouterEngineConfig.CostControl != nil && opts.RouterEngineConfig.CostControl.Enabled {
 		planConfig.ComputeCosts = true
-		planConfig.StaticCostDefaultListSize = routerEngineCfg.CostControl.EstimatedListSize
-		planConfig.IgnoreImplementingTypeWeights = routerEngineCfg.CostControl.IgnoreImplementingTypeWeights
+		planConfig.StaticCostDefaultListSize = opts.RouterEngineConfig.CostControl.EstimatedListSize
+		planConfig.IgnoreImplementingTypeWeights = opts.RouterEngineConfig.CostControl.IgnoreImplementingTypeWeights
 	}
 
 	return planConfig, providers, nil

router/core/executor_test.go

@@ -0,0 +1,37 @@
+package core
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/wundergraph/graphql-go-tools/v2/pkg/ast"
+	"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan"
+)
+
+func TestExecutorCloseReleasesSchemaReferences(t *testing.T) {
+	t.Parallel()
+
+	executor := &Executor{
+		ClientSchema:    &ast.Document{},
+		RouterSchema:    &ast.Document{},
+		PlanConfig:      plan.Configuration{DataSources: []plan.DataSource{nil}},
+		RenameTypeNames: nil,
+	}
+
+	executor.Close()
+
+	require.Nil(t, executor.ClientSchema)
+	require.Nil(t, executor.RouterSchema)
+	require.Empty(t, executor.PlanConfig.DataSources)
+	require.Nil(t, executor.RenameTypeNames)
+	require.Nil(t, executor.Resolver)
+}
+
+func TestExecutorCloseNilSafe(t *testing.T) {
+	t.Parallel()
+
+	var executor *Executor
+	require.NotPanics(t, func() {
+		executor.Close()
+	})
+}

router/core/factoryresolver.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/url"
 	"slices"
+	"sync"
 	"time"
 
 	"github.com/buger/jsonparser"
@@ -17,6 +18,7 @@ import (
 	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/grpcconnector"
+	"github.com/wundergraph/cosmo/router/pkg/mondaytweaks"
 	rmetric "github.com/wundergraph/cosmo/router/pkg/metric"
 	"github.com/wundergraph/cosmo/router/pkg/pubsub"
 	pubsub_datasource "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource"
@@ -78,6 +80,10 @@ type DefaultFactoryResolver struct {
 	transportFactory              ApiTransportFactory
 	defaultSubgraphRequestTimeout time.Duration
 	subscriptionClientOptions     []graphql_datasource.SubscriptionClientOption
+	useNoopSubscriptionClient     bool
+
+	subscriptionClient     graphql_datasource.GraphQLSubscriptionClient
+	subscriptionClientOnce sync.Once
 }
 
 func NewDefaultFactoryResolver(
@@ -131,7 +137,9 @@ func NewDefaultFactoryResolver(
 		graphql_datasource.WithLogger(factoryLogger),
 	}
 
+	useNoopSubscriptionClient := false
 	if subscriptionClientOptions != nil {
+		useNoopSubscriptionClient = subscriptionClientOptions.UseNoopClient
 		if subscriptionClientOptions.PingInterval > 0 {
 			options = append(options, graphql_datasource.WithPingInterval(subscriptionClientOptions.PingInterval))
 		}
@@ -164,6 +172,7 @@ func NewDefaultFactoryResolver(
 		transportFactory:              transportFactory,
 		defaultSubgraphRequestTimeout: transportOptions.SubgraphTransportOptions.RequestTimeout,
 		subscriptionClientOptions:     options,
+		useNoopSubscriptionClient:     useNoopSubscriptionClient,
 	}
 }
 
@@ -183,10 +192,40 @@ func (d *DefaultFactoryResolver) ResolveGraphqlFactory(subgraphName string) (pla
 
 	if d.transportFactory == nil || d.baseTransport == nil {
 		// dummy implementation for plan generator that doesn't make requests
-		subscriptionClient := graphql_datasource.NewGraphQLSubscriptionClient(d.engineCtx,
+		return graphql_datasource.NewFactory(d.engineCtx, http.DefaultClient, d.subscriptionClientForFactory())
+	}
+
+	defaultHTTPClient := &http.Client{
+		Timeout:   d.defaultSubgraphRequestTimeout,
+		Transport: d.transportFactory.RoundTripper(d.baseTransport),
+	}
+
+	if subgraphClient, ok := d.subgraphHTTPClients[subgraphName]; ok {
+		// it's intentional that we're not using the subgraphClient for subscriptions
+		// custom subgraph clients are intended to be used for custom timeouts, which is not relevant for subscriptions
+		return graphql_datasource.NewFactory(d.engineCtx, subgraphClient, d.subscriptionClientForFactory())
+	}
+
+	return graphql_datasource.NewFactory(d.engineCtx, defaultHTTPClient, d.subscriptionClientForFactory())
+}
+
+func (d *DefaultFactoryResolver) subscriptionClientForFactory() graphql_datasource.GraphQLSubscriptionClient {
+	if mondaytweaks.ShareUpstreamSubscriptionClient {
+		return d.sharedSubscriptionClient()
+	}
+	return d.newSubscriptionClient()
+}
+
+func (d *DefaultFactoryResolver) newSubscriptionClient() graphql_datasource.GraphQLSubscriptionClient {
+	if d.useNoopSubscriptionClient {
+		return noopGraphQLSubscriptionClientInstance
+	}
+
+	if d.transportFactory == nil || d.baseTransport == nil {
+		return graphql_datasource.NewGraphQLSubscriptionClient(
+			d.engineCtx,
 			d.subscriptionClientOptions...,
 		)
-		return graphql_datasource.NewFactory(d.engineCtx, http.DefaultClient, subscriptionClient)
 	}
 
 	defaultHTTPClient := &http.Client{
@@ -198,18 +237,49 @@ func (d *DefaultFactoryResolver) ResolveGraphqlFactory(subgraphName string) (pla
 		Transport: d.transportFactory.RoundTripper(d.baseTransport),
 	}
 
-	subscriptionClient := graphql_datasource.NewGraphQLSubscriptionClient(
+	return graphql_datasource.NewGraphQLSubscriptionClient(
 		d.engineCtx,
-		append([]graphql_datasource.SubscriptionClientOption{graphql_datasource.WithUpgradeClient(defaultHTTPClient), graphql_datasource.WithStreamingClient(streamingClient)}, d.subscriptionClientOptions...)...,
+		append([]graphql_datasource.SubscriptionClientOption{
+			graphql_datasource.WithUpgradeClient(defaultHTTPClient),
+			graphql_datasource.WithStreamingClient(streamingClient),
+		}, d.subscriptionClientOptions...)...,
 	)
+}
 
-	if subgraphClient, ok := d.subgraphHTTPClients[subgraphName]; ok {
-		// it's intentional that we're not using the subgraphClient for subscriptions
-		// custom subgraph clients are intended to be used for custom timeouts, which is not relevant for subscriptions
-		return graphql_datasource.NewFactory(d.engineCtx, subgraphClient, subscriptionClient)
-	}
+func (d *DefaultFactoryResolver) sharedSubscriptionClient() graphql_datasource.GraphQLSubscriptionClient {
+	d.subscriptionClientOnce.Do(func() {
+		if d.useNoopSubscriptionClient {
+			d.subscriptionClient = noopGraphQLSubscriptionClientInstance
+			return
+		}
+
+		if d.transportFactory == nil || d.baseTransport == nil {
+			d.subscriptionClient = graphql_datasource.NewGraphQLSubscriptionClient(
+				d.engineCtx,
+				d.subscriptionClientOptions...,
+			)
+			return
+		}
+
+		defaultHTTPClient := &http.Client{
+			Timeout:   d.defaultSubgraphRequestTimeout,
+			Transport: d.transportFactory.RoundTripper(d.baseTransport),
+		}
+
+		streamingClient := &http.Client{
+			Transport: d.transportFactory.RoundTripper(d.baseTransport),
+		}
+
+		d.subscriptionClient = graphql_datasource.NewGraphQLSubscriptionClient(
+			d.engineCtx,
+			append([]graphql_datasource.SubscriptionClientOption{
+				graphql_datasource.WithUpgradeClient(defaultHTTPClient),
+				graphql_datasource.WithStreamingClient(streamingClient),
+			}, d.subscriptionClientOptions...)...,
+		)
+	})
 
-	return graphql_datasource.NewFactory(d.engineCtx, defaultHTTPClient, subscriptionClient)
+	return d.subscriptionClient
 }
 
 func (d *DefaultFactoryResolver) ResolveStaticFactory() (factory plan.PlannerFactory[staticdatasource.Configuration], err error) {