feat: rewarm cache from planner cache on schema change

router-tests/cache_warmup_test.go

@@ -2,6 +2,8 @@ package integration
 
 import (
 	"context"
+	"github.com/stretchr/testify/assert"
+	"github.com/wundergraph/cosmo/router/pkg/controlplane/configpoller"
 	"net/http"
 	"testing"
 	"time"
@@ -914,6 +916,199 @@ func TestCacheWarmup(t *testing.T) {
 	})
 }
 
+func TestInMemorySwitchoverCaching(t *testing.T) {
+	t.Parallel()
+
+	t.Run("Verify the plan is cached on config restart when in memory switchover is enabled", func(t *testing.T) {
+		t.Parallel()
+
+		pm := ConfigPollerMock{
+			ready: make(chan struct{}),
+		}
+
+		testenv.Run(t, &testenv.Config{
+			RouterOptions: []core.Option{
+				core.WithCacheWarmupConfig(&config.CacheWarmupConfiguration{
+					Enabled: true,
+					Source: config.CacheWarmupSource{
+						InMemorySwitchover: config.CacheWarmupInMemorySwitchover{
+							Enabled: true,
+						},
+					},
+				}),
+				core.WithConfigVersionHeader(true),
+			},
+			RouterConfig: &testenv.RouterConfig{
+				ConfigPollerFactory: func(config *nodev1.RouterConfig) configpoller.ConfigPoller {
+					pm.initConfig = config
+					return &pm
+				},
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+				Query: `{ employees { id } }`,
+			})
+			require.Equal(t, res.Response.StatusCode, 200)
+			require.Equal(t, xEnv.RouterConfigVersionMain(), res.Response.Header.Get("X-Router-Config-Version"))
+			require.JSONEq(t, employeesIDData, res.Body)
+			require.Equal(t, "MISS", res.Response.Header.Get("x-wg-execution-plan-cache"))
+
+			// Wait for the config poller to be ready
+			<-pm.ready
+
+			pm.initConfig.Version = "updated"
+			require.NoError(t, pm.updateConfig(pm.initConfig, "old-1"))
+
+			res = xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+				Query: `{ employees { id } }`,
+			})
+			require.Equal(t, res.Response.StatusCode, 200)
+			require.Equal(t, res.Response.Header.Get("X-Router-Config-Version"), "updated")
+			require.JSONEq(t, employeesIDData, res.Body)
+			require.Equal(t, "HIT", res.Response.Header.Get("x-wg-execution-plan-cache"))
+
+		})
+	})
+
+	t.Run("Verify the plan is not cached on config restart when in cache warmer is disabled", func(t *testing.T) {
+		t.Parallel()
+
+		pm := ConfigPollerMock{
+			ready: make(chan struct{}),
+		}
+
+		testenv.Run(t, &testenv.Config{
+			RouterOptions: []core.Option{
+				core.WithCacheWarmupConfig(&config.CacheWarmupConfiguration{
+					Enabled: false,
+				}),
+				core.WithConfigVersionHeader(true),
+			},
+			RouterConfig: &testenv.RouterConfig{
+				ConfigPollerFactory: func(config *nodev1.RouterConfig) configpoller.ConfigPoller {
+					pm.initConfig = config
+					return &pm
+				},
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+				Query: `{ employees { id } }`,
+			})
+			require.Equal(t, res.Response.StatusCode, 200)
+			require.Equal(t, xEnv.RouterConfigVersionMain(), res.Response.Header.Get("X-Router-Config-Version"))
+			require.Equal(t, "MISS", res.Response.Header.Get("x-wg-execution-plan-cache"))
+
+			// Wait for the config poller to be ready
+			<-pm.ready
+
+			pm.initConfig.Version = "updated"
+			require.NoError(t, pm.updateConfig(pm.initConfig, "old-1"))
+
+			res = xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+				Query: `{ employees { id } }`,
+			})
+			require.Equal(t, res.Response.StatusCode, 200)
+			require.Equal(t, res.Response.Header.Get("X-Router-Config-Version"), "updated")
+			require.Equal(t, "MISS", res.Response.Header.Get("x-wg-execution-plan-cache"))
+		})
+	})
+
+	t.Run("Verify the plan is not cached on config restart when using default cache warmer", func(t *testing.T) {
+		t.Parallel()
+
+		pm := ConfigPollerMock{
+			ready: make(chan struct{}),
+		}
+
+		testenv.Run(t, &testenv.Config{
+			RouterOptions: []core.Option{
+				core.WithCacheWarmupConfig(&config.CacheWarmupConfiguration{
+					Enabled: true,
+					Source: config.CacheWarmupSource{
+						InMemorySwitchover: config.CacheWarmupInMemorySwitchover{
+							Enabled: false,
+						},
+					},
+				}),
+				core.WithConfigVersionHeader(true),
+			},
+			RouterConfig: &testenv.RouterConfig{
+				ConfigPollerFactory: func(config *nodev1.RouterConfig) configpoller.ConfigPoller {
+					pm.initConfig = config
+					return &pm
+				},
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+				Query: `{ employees { id customDetails: details { forename } } }`,
+			})
+			require.Equal(t, res.Response.StatusCode, 200)
+			require.Equal(t, xEnv.RouterConfigVersionMain(), res.Response.Header.Get("X-Router-Config-Version"))
+			require.Equal(t, "MISS", res.Response.Header.Get("x-wg-execution-plan-cache"))
+
+			// Wait for the config poller to be ready
+			<-pm.ready
+
+			pm.initConfig.Version = "updated"
+			require.NoError(t, pm.updateConfig(pm.initConfig, "old-1"))
+
+			res = xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+				Query: `{ employees { id customDetails: details { forename } } }`,
+			})
+			require.Equal(t, res.Response.StatusCode, 200)
+			require.Equal(t, res.Response.Header.Get("X-Router-Config-Version"), "updated")
+			require.Equal(t, "MISS", res.Response.Header.Get("x-wg-execution-plan-cache"))
+		})
+	})
+
+	t.Run("Verify plan is cached when static execution config is reloaded", func(t *testing.T) {
+		t.Parallel()
+
+		// Create a temporary file for the router config
+		configFile := t.TempDir() + "/config.json"
+
+		// Initial config with just the employees subgraph
+		writeTestConfig(t, "initial", configFile)
+
+		testenv.Run(t, &testenv.Config{
+			RouterOptions: []core.Option{
+				core.WithConfigVersionHeader(true),
+				core.WithExecutionConfig(&core.ExecutionConfig{
+					Path:          configFile,
+					Watch:         true,
+					WatchInterval: 100 * time.Millisecond,
+				}),
+				core.WithCacheWarmupConfig(&config.CacheWarmupConfiguration{
+					Enabled: true,
+					Source: config.CacheWarmupSource{
+						InMemorySwitchover: config.CacheWarmupInMemorySwitchover{
+							Enabled: true,
+						},
+					},
+				}),
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+				Query: `query { hello }`,
+			})
+			require.Equal(t, res.Response.StatusCode, 200)
+			require.Equal(t, "initial", res.Response.Header.Get("X-Router-Config-Version"))
+			require.Equal(t, "MISS", res.Response.Header.Get("x-wg-execution-plan-cache"))
+
+			writeTestConfig(t, "updated", configFile)
+
+			require.EventuallyWithT(t, func(t *assert.CollectT) {
+				res = xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+					Query: `query { hello }`,
+				})
+				assert.Equal(t, "updated", res.Response.Header.Get("X-Router-Config-Version"))
+				assert.Equal(t, "HIT", res.Response.Header.Get("x-wg-execution-plan-cache"))
+			}, 2*time.Second, 100*time.Millisecond)
+		})
+	})
+
+}
+
 // findDataPoint finds a data point in a slice of histogram data points by matching
 // the value of WgEnginePlanCacheHit attribute
 func findDataPoint(t *testing.T, dataPoints []metricdata.HistogramDataPoint[float64], cacheHit bool) metricdata.HistogramDataPoint[float64] {

router-tests/go.mod

@@ -204,6 +204,7 @@ replace (
 // Use what is best for your personal workflow. See CONTRIBUTING.md for more information
 
 replace (
+	github.com/dgraph-io/ristretto/v2 => github.com/wundergraph/ristretto/v2 v2.0.0-20260113151003-f7231509ad4d
 	github.com/wundergraph/cosmo/demo => ../demo
 	github.com/wundergraph/cosmo/demo/pkg/subgraphs/projects => ../demo/pkg/subgraphs/projects
 	github.com/wundergraph/cosmo/router => ../router

router-tests/go.sum

@@ -65,10 +65,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/dgraph-io/ristretto/v2 v2.1.0 h1:59LjpOJLNDULHh8MC4UaegN52lC4JnO2dITsie/Pa8I=
-github.com/dgraph-io/ristretto/v2 v2.1.0/go.mod h1:uejeqfYXpUomfse0+lO+13ATz4TypQYLJZzBSAemuB4=
-github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13 h1:fAjc9m62+UWV/WAFKLNi6ZS0675eEUC9y3AlwSbQu1Y=
-github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
+github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da h1:aIftn67I1fkbMa512G+w+Pxci9hJPB8oMnkcP3iZF38=
+github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/dgryski/trifles v0.0.0-20230903005119-f50d829f2e54 h1:SG7nF6SRlWhcT7cNTs5R6Hk4V2lcmLz2NsG2VnInyNo=
@@ -354,6 +352,8 @@ github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083 h1:8/D7f8gKxTB
 github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083/go.mod h1:eOTL6acwctsN4F3b7YE+eE2t8zcJ/doLm9sZzsxxxrE=
 github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.242 h1:0ieQmRxYz0nbJEbaaA4Cx2RPcxlomhQ8KI31uuevWx0=
 github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.242/go.mod h1:mX25ASEQiKamxaFSK6NZihh0oDCigIuzro30up4mFH4=
+github.com/wundergraph/ristretto/v2 v2.0.0-20260113151003-f7231509ad4d h1:FlfUZzJbnc8+bgSKJ2jCqozn9CS08ToZe/JJxh4DPGo=
+github.com/wundergraph/ristretto/v2 v2.0.0-20260113151003-f7231509ad4d/go.mod h1:0KsrXtXvnv0EqnzyowllbVJB8yBonswa2lTCK2gGo9E=
 github.com/xrash/smetrics v0.0.0-20250705151800-55b8f293f342 h1:FnBeRrxr7OU4VvAzt5X7s6266i6cSVkkFPS0TuXWbIg=
 github.com/xrash/smetrics v0.0.0-20250705151800-55b8f293f342/go.mod h1:Ohn+xnUBiLI6FVj/9LpzZWtj1/D6lUovWYBkxHVV3aM=
 github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=

router/cmd/main.go

@@ -150,7 +150,8 @@ func Main() {
 	}
 
 	rs, err := core.NewRouterSupervisor(&core.RouterSupervisorOpts{
-		BaseLogger: baseLogger,
+		BaseLogger:       baseLogger,
+		SwitchoverConfig: core.NewSwitchoverConfig(),
 		ConfigFactory: func() (*config.Config, error) {
 			result, err := config.LoadConfig(*configPathFlag)
 			if err != nil {

router/core/cache_warmup_plans.go

@@ -0,0 +1,34 @@
+package core
+
+import (
+	"context"
+	"github.com/dgraph-io/ristretto/v2"
+	nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1"
+	"go.uber.org/zap"
+)
+
+var _ CacheWarmupSource = (*PlanSource)(nil)
+
+type PlanSource struct {
+	queries []*nodev1.Operation
+}
+
+func NewPlanSource(switchoverCacheWarmerQueries *ristretto.Cache[uint64, *planWithMetaData]) *PlanSource {
+	// Extract the items on creation so that the previous planCache can be garbage collected as we won't hold a reference
+	items := make([]*nodev1.Operation, 0)
+
+	if switchoverCacheWarmerQueries != nil {
+		switchoverCacheWarmerQueries.Iter(func(k any, v *planWithMetaData) (stop bool) {
+			items = append(items, &nodev1.Operation{
+				Request: &nodev1.OperationRequest{Query: v.content},
+			})
+			return false
+		})
+	}
+
+	return &PlanSource{queries: items}
+}
+
+func (c *PlanSource) LoadItems(ctx context.Context, log *zap.Logger) ([]*nodev1.Operation, error) {
+	return c.queries, nil
+}

router/core/graph_server.go

@@ -112,14 +112,15 @@ type BuildGraphMuxOptions struct {
 	EngineConfig        *nodev1.EngineConfiguration
 	ConfigSubgraphs     []*nodev1.Subgraph
 	RoutingUrlGroupings map[string]map[string]bool
+	SwitchoverConfig    *SwitchoverConfig
 }
 
 func (b BuildGraphMuxOptions) IsBaseGraph() bool {
 	return b.FeatureFlagName == ""
 }
 
 // newGraphServer creates a new server instance.
-func newGraphServer(ctx context.Context, r *Router, routerConfig *nodev1.RouterConfig, proxy ProxyFunc) (*graphServer, error) {
+func newGraphServer(ctx context.Context, r *Router, routerConfig *nodev1.RouterConfig, proxy ProxyFunc, switchoverConfig *SwitchoverConfig) (*graphServer, error) {
 	/* Older versions of composition will not populate a compatibility version.
 	 * Currently, all "old" router execution configurations are compatible as there have been no breaking
 	 * changes.
@@ -273,6 +274,7 @@ func newGraphServer(ctx context.Context, r *Router, routerConfig *nodev1.RouterC
 		EngineConfig:        routerConfig.GetEngineConfig(),
 		ConfigSubgraphs:     routerConfig.GetSubgraphs(),
 		RoutingUrlGroupings: routingUrlGroupings,
+		SwitchoverConfig:    switchoverConfig,
 	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to build base mux: %w", err)
@@ -283,7 +285,7 @@ func newGraphServer(ctx context.Context, r *Router, routerConfig *nodev1.RouterC
 		s.logger.Info("Feature flags enabled", zap.Strings("flags", maps.Keys(featureFlagConfigMap)))
 	}
 
-	multiGraphHandler, err := s.buildMultiGraphHandler(ctx, gm.mux, featureFlagConfigMap)
+	multiGraphHandler, err := s.buildMultiGraphHandler(ctx, gm.mux, featureFlagConfigMap, switchoverConfig)
 	if err != nil {
 		return nil, fmt.Errorf("failed to build feature flag handler: %w", err)
 	}
@@ -433,6 +435,7 @@ func (s *graphServer) buildMultiGraphHandler(
 	ctx context.Context,
 	baseMux *chi.Mux,
 	featureFlagConfigs map[string]*nodev1.FeatureFlagRouterExecutionConfig,
+	switchoverConfig *SwitchoverConfig,
 ) (http.HandlerFunc, error) {
 	if len(featureFlagConfigs) == 0 {
 		return baseMux.ServeHTTP, nil
@@ -447,6 +450,7 @@ func (s *graphServer) buildMultiGraphHandler(
 			RouterConfigVersion: executionConfig.GetVersion(),
 			EngineConfig:        executionConfig.GetEngineConfig(),
 			ConfigSubgraphs:     executionConfig.Subgraphs,
+			SwitchoverConfig:    switchoverConfig,
 		})
 		if err != nil {
 			return nil, fmt.Errorf("failed to build mux for feature flag '%s': %w", featureFlagName, err)
@@ -519,12 +523,12 @@ type graphMux struct {
 	validationCache             *ristretto.Cache[uint64, bool]
 	operationHashCache          *ristretto.Cache[uint64, string]
 
-	accessLogsFileLogger   *logging.BufferedLogger
-	metricStore            rmetric.Store
-	prometheusCacheMetrics *rmetric.CacheMetrics
-	otelCacheMetrics       *rmetric.CacheMetrics
-	streamMetricStore      rmetric.StreamMetricStore
-	prometheusMetricsExporter  *graphqlmetrics.PrometheusMetricsExporter
+	accessLogsFileLogger      *logging.BufferedLogger
+	metricStore               rmetric.Store
+	prometheusCacheMetrics    *rmetric.CacheMetrics
+	otelCacheMetrics          *rmetric.CacheMetrics
+	streamMetricStore         rmetric.StreamMetricStore
+	prometheusMetricsExporter *graphqlmetrics.PrometheusMetricsExporter
 }
 
 // buildOperationCaches creates the caches for the graph mux.
@@ -1296,7 +1300,8 @@ func (s *graphServer) buildGraphMux(
 		DisableExposingVariablesContentOnValidationError: s.engineExecutionConfiguration.DisableExposingVariablesContentOnValidationError,
 		ComplexityLimits:                                 s.securityConfiguration.ComplexityLimits,
 	})
-	operationPlanner := NewOperationPlanner(executor, gm.planCache)
+
+	operationPlanner := NewOperationPlanner(executor, gm.planCache, opts.SwitchoverConfig.inMemorySwitchOverCache.enabled)
 
 	// We support the MCP only on the base graph. Feature flags are not supported yet.
 	if opts.IsBaseGraph() && s.mcpServer != nil {
@@ -1350,6 +1355,11 @@ func (s *graphServer) buildGraphMux(
 			warmupConfig.Source = NewFileSystemSource(&FileSystemSourceConfig{
 				RootPath: s.Config.cacheWarmup.Source.Filesystem.Path,
 			})
+		} else if s.Config.cacheWarmup.Source.InMemorySwitchover.Enabled {
+			// We first utilize the plan cache (if it was already set, so not on first starts) to create a list of queries
+			// and reset the plan cache to the new plan cache for this start afterwords
+			warmupConfig.Source = NewPlanSource(opts.SwitchoverConfig.inMemorySwitchOverCache.getPlanCacheForFF(opts.FeatureFlagName))
+			opts.SwitchoverConfig.inMemorySwitchOverCache.setPlanCacheForFF(opts.FeatureFlagName, gm.planCache)
 		} else {
 			cdnSource, err := NewCDNSource(s.Config.cdnConfig.URL, s.graphApiToken, s.logger)
 			if err != nil {