feat: rewarm cache from planner cache on schema change

router-tests/cache_warmup_test.go

@@ -10,24 +10,18 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/wundergraph/cosmo/router-tests/testenv"
+	"github.com/wundergraph/cosmo/router/core"
+	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/controlplane/configpoller"
-
+	"github.com/wundergraph/cosmo/router/pkg/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/sdk/metric"
-	"go.opentelemetry.io/otel/sdk/metric/metricdata/metricdatatest"
-
-	nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1"
-
 	"go.opentelemetry.io/otel/sdk/metric/metricdata"
-
-	"github.com/wundergraph/cosmo/router/pkg/otel"
-
-	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/sdk/metric/metricdata/metricdatatest"
 	"go.uber.org/zap"
-
-	"github.com/wundergraph/cosmo/router-tests/testenv"
-	"github.com/wundergraph/cosmo/router/core"
-	"github.com/wundergraph/cosmo/router/pkg/config"
 )
 
 func TestCacheWarmup(t *testing.T) {
@@ -1141,13 +1135,13 @@ engine:
 			listenString1 := "after1"
 			updateConfig(t, xEnv, ctx, listenString1, getConfigString(listenString1))
 			res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{Query: `query { hello }`})
-			assert.Equal(t, "MISS", res.Response.Header.Get("x-wg-execution-plan-cache"))
+			require.Equal(t, "MISS", res.Response.Header.Get("x-wg-execution-plan-cache"))
 
 			// Verify cache persisted on restart
 			listenString2 := "after2"
 			updateConfig(t, xEnv, ctx, listenString2, getConfigString(listenString2))
 			res = xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{Query: `query { hello }`})
-			assert.Equal(t, "HIT", res.Response.Header.Get("x-wg-execution-plan-cache"))
+			require.Equal(t, "HIT", res.Response.Header.Get("x-wg-execution-plan-cache"))
 		})
 
 		require.NoError(t, err)

router/core/cache_warmup_cdn.go

@@ -10,15 +10,14 @@ import (
 	"net/url"
 
 	nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1"
-	"google.golang.org/protobuf/encoding/protojson"
-
 	"github.com/wundergraph/cosmo/router/internal/httpclient"
 	"github.com/wundergraph/cosmo/router/internal/jwt"
 	"go.opentelemetry.io/otel/codes"
 	semconv12 "go.opentelemetry.io/otel/semconv/v1.12.0"
 	semconv "go.opentelemetry.io/otel/semconv/v1.17.0"
 	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
+	"google.golang.org/protobuf/encoding/protojson"
 )
 
 var _ CacheWarmupSource = (*CDNSource)(nil)

router/core/cache_warmup_plans.go

@@ -2,20 +2,28 @@ package core
 
 import (
 	"context"
+
 	nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1"
 	"go.uber.org/zap"
 )
 
 var _ CacheWarmupSource = (*PlanSource)(nil)
 
+// PlanSource is a very basic cache warmup source that relies on the caller of this type to pass in the
+// queries to be used for cache warming directly
 type PlanSource struct {
 	queries []*nodev1.Operation
 }
 
+// NewPlanSource creates a new PlanSource with the given queries from the caller
 func NewPlanSource(switchoverCacheWarmerQueries []*nodev1.Operation) *PlanSource {
+	if switchoverCacheWarmerQueries == nil {
+		switchoverCacheWarmerQueries = make([]*nodev1.Operation, 0)
+	}
 	return &PlanSource{queries: switchoverCacheWarmerQueries}
 }
 
+// LoadItems loads the items from the plan source when called by the cache warmer
 func (c *PlanSource) LoadItems(_ context.Context, _ *zap.Logger) ([]*nodev1.Operation, error) {
 	return c.queries, nil
 }

router/core/graph_server.go

@@ -120,8 +120,16 @@ func (b BuildGraphMuxOptions) IsBaseGraph() bool {
 	return b.FeatureFlagName == ""
 }
 
+// BuildMultiGraphHandlerOptions contains the configuration options for building a multi-graph handler.
+type BuildMultiGraphHandlerOptions struct {
+	BaseMux                 *chi.Mux
+	FeatureFlagConfigs      map[string]*nodev1.FeatureFlagRouterExecutionConfig
+	SwitchoverConfig        *SwitchoverConfig
+	CosmoCacheWarmerEnabled bool
+}
+
 // newGraphServer creates a new server instance.
-func newGraphServer(ctx context.Context, r *Router, routerConfig *nodev1.RouterConfig, proxy ProxyFunc, switchoverConfig *SwitchoverConfig) (*graphServer, error) {
+func newGraphServer(ctx context.Context, r *Router, routerConfig *nodev1.RouterConfig, proxy ProxyFunc) (*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.
@@ -280,7 +288,7 @@ func newGraphServer(ctx context.Context, r *Router, routerConfig *nodev1.RouterC
 		EngineConfig:            routerConfig.GetEngineConfig(),
 		ConfigSubgraphs:         routerConfig.GetSubgraphs(),
 		RoutingUrlGroupings:     routingUrlGroupings,
-		SwitchoverConfig:        switchoverConfig,
+		SwitchoverConfig:        r.switchoverConfig,
 		CosmoCacheWarmerEnabled: cosmoCacheWarmerEnabled,
 	})
 	if err != nil {
@@ -292,7 +300,12 @@ 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, switchoverConfig, cosmoCacheWarmerEnabled)
+	multiGraphHandler, err := s.buildMultiGraphHandler(ctx, BuildMultiGraphHandlerOptions{
+		BaseMux:                 gm.mux,
+		FeatureFlagConfigs:      featureFlagConfigMap,
+		SwitchoverConfig:        r.switchoverConfig,
+		CosmoCacheWarmerEnabled: cosmoCacheWarmerEnabled,
+	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to build feature flag handler: %w", err)
 	}
@@ -440,26 +453,23 @@ func getRoutingUrlGroupingForCircuitBreakers(
 
 func (s *graphServer) buildMultiGraphHandler(
 	ctx context.Context,
-	baseMux *chi.Mux,
-	featureFlagConfigs map[string]*nodev1.FeatureFlagRouterExecutionConfig,
-	switchoverConfig *SwitchoverConfig,
-	cosmoCacheWarmerEnabled bool,
+	opts BuildMultiGraphHandlerOptions,
 ) (http.HandlerFunc, error) {
-	if len(featureFlagConfigs) == 0 {
-		return baseMux.ServeHTTP, nil
+	if len(opts.FeatureFlagConfigs) == 0 {
+		return opts.BaseMux.ServeHTTP, nil
 	}
 
-	featureFlagToMux := make(map[string]*chi.Mux, len(featureFlagConfigs))
+	featureFlagToMux := make(map[string]*chi.Mux, len(opts.FeatureFlagConfigs))
 
 	// Build all the muxes for the feature flags in serial to avoid any race conditions
-	for featureFlagName, executionConfig := range featureFlagConfigs {
+	for featureFlagName, executionConfig := range opts.FeatureFlagConfigs {
 		gm, err := s.buildGraphMux(ctx, BuildGraphMuxOptions{
 			FeatureFlagName:         featureFlagName,
 			RouterConfigVersion:     executionConfig.GetVersion(),
 			EngineConfig:            executionConfig.GetEngineConfig(),
 			ConfigSubgraphs:         executionConfig.Subgraphs,
-			SwitchoverConfig:        switchoverConfig,
-			CosmoCacheWarmerEnabled: cosmoCacheWarmerEnabled,
+			SwitchoverConfig:        opts.SwitchoverConfig,
+			CosmoCacheWarmerEnabled: opts.CosmoCacheWarmerEnabled,
 		})
 		if err != nil {
 			return nil, fmt.Errorf("failed to build mux for feature flag '%s': %w", featureFlagName, err)
@@ -486,7 +496,7 @@ func (s *graphServer) buildMultiGraphHandler(
 			return
 		}
 
-		baseMux.ServeHTTP(w, r)
+		opts.BaseMux.ServeHTTP(w, r)
 	}, nil
 }
 
@@ -1365,7 +1375,11 @@ func (s *graphServer) buildGraphMux(
 			warmupConfig.Source = NewFileSystemSource(&FileSystemSourceConfig{
 				RootPath: s.Config.cacheWarmup.Source.Filesystem.Path,
 			})
-		// We want to enable this only whenever there is no registering with cosmo
+		// Enable in-memory switchover fallback when:
+		// - Router has cache warmer with inMemorySwitchoverFallback enabled, AND
+		// - Either:
+		//   - Using static execution config (not Cosmo): s.selfRegister == nil
+		//   - OR Cosmo cache warmer is disabled: !opts.CosmoCacheWarmerEnabled
 		case s.cacheWarmup.InMemorySwitchoverFallback && (s.selfRegister == nil || !opts.CosmoCacheWarmerEnabled):
 			// 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

router/core/restart_switchover_config.go

@@ -1,12 +1,11 @@
 package core
 
 import (
-	"go.uber.org/zap"
 	"sync"
 
 	"github.com/dgraph-io/ristretto/v2"
-
 	nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1"
+	"go.uber.org/zap"
 )
 
 type planCache = *ristretto.Cache[uint64, *planWithMetaData]
@@ -24,31 +23,36 @@ func NewSwitchoverConfig(logger *zap.Logger) *SwitchoverConfig {
 	}
 }
 
+// UpdateSwitchoverConfig updates the switchover config based on the provided config.
 func (s *SwitchoverConfig) UpdateSwitchoverConfig(config *Config, isCosmoCacheWarmerEnabled bool) {
 	s.inMemorySwitchOverCache.updateStateFromConfig(config, isCosmoCacheWarmerEnabled)
 }
 
+// CleanupFeatureFlags cleans up anything related to unused feature flags due to being now excluded
+// from the execution config
 func (s *SwitchoverConfig) CleanupFeatureFlags(routerCfg *nodev1.RouterConfig) {
 	s.inMemorySwitchOverCache.cleanupUnusedFeatureFlags(routerCfg)
 }
 
-func (s *SwitchoverConfig) ProcessOnConfigChangeRestart() {
+func (s *SwitchoverConfig) ProcessOnConfigChangeRouterInstanceRestart() {
 	// For cases of router config changes (not execution config), we shut down before creating the
 	// graph mux, because we need to initialize everything from the start
 	// This causes problems in using the previous planCache reference as it gets closed, so we need to
 	// copy it over before it gets closed, and we restart with config changes
 
 	// There can be inflight requests when this is called even though it's called in the restart path,
 	// This is because this is called before the router instance is shutdown before being reloaded
-	s.inMemorySwitchOverCache.processOnConfigChangeRestart()
+	s.inMemorySwitchOverCache.extractQueriesAndOverridePlanCache()
 }
 
+// InMemorySwitchOverCache is a store that stores either queries or references to the planner cache for use with the cache warmer
 type InMemorySwitchOverCache struct {
 	mu                    sync.RWMutex
 	queriesForFeatureFlag map[string]any
 	logger                *zap.Logger
 }
 
+// updateStateFromConfig updates the internal state of the in-memory switchover cache based on the provided config
 func (c *InMemorySwitchOverCache) updateStateFromConfig(config *Config, isCosmoCacheWarmerEnabled bool) {
 	enabled := config.cacheWarmup != nil &&
 		!isCosmoCacheWarmerEnabled && // We only enable in-memory switchover cache if the cosmo cache warmer is not enabled
@@ -71,13 +75,15 @@ func (c *InMemorySwitchOverCache) updateStateFromConfig(config *Config, isCosmoC
 	}
 }
 
+// IsEnabled returns whether the in-memory switchover cache is enabled
 func (c *InMemorySwitchOverCache) IsEnabled() bool {
 	c.mu.RLock()
 	defer c.mu.RUnlock()
 
 	return c.queriesForFeatureFlag != nil
 }
 
+// getPlanCacheForFF gets the plan cache in the []*nodev1.Operation format for a specific feature flag key
 func (c *InMemorySwitchOverCache) getPlanCacheForFF(featureFlagKey string) []*nodev1.Operation {
 	c.mu.RLock()
 	defer c.mu.RUnlock()
@@ -91,16 +97,17 @@ func (c *InMemorySwitchOverCache) getPlanCacheForFF(featureFlagKey string) []*no
 		return convertToNodeOperation(cache)
 	case []*nodev1.Operation:
 		return cache
+	// This would occur during the first start (we add this case to specifically log any other cases)
 	case nil:
-		// This would occur during the first start
-		return make([]*nodev1.Operation, 0)
+		return nil
+	// This should not happen as we cannot have any types other than the above
 	default:
-		// This should not happen as we cannot have any types other than the above
 		c.logger.Error("unexpected type")
-		return make([]*nodev1.Operation, 0)
+		return nil
 	}
 }
 
+// setPlanCacheForFF sets the plan cache for a specific feature flag key
 func (c *InMemorySwitchOverCache) setPlanCacheForFF(featureFlagKey string, cache planCache) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -111,7 +118,8 @@ func (c *InMemorySwitchOverCache) setPlanCacheForFF(featureFlagKey string, cache
 	c.queriesForFeatureFlag[featureFlagKey] = cache
 }
 
-func (c *InMemorySwitchOverCache) processOnConfigChangeRestart() {
+// extractQueriesAndOverridePlanCache extracts the queries from the plan cache and overrides the internal map
+func (c *InMemorySwitchOverCache) extractQueriesAndOverridePlanCache() {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
@@ -128,6 +136,8 @@ func (c *InMemorySwitchOverCache) processOnConfigChangeRestart() {
 	c.queriesForFeatureFlag = switchoverMap
 }
 
+// cleanupUnusedFeatureFlags removes any feature flags from the in-memory switchover cache
+// this is useful in case where the updated execution config excludes a feature flag
 func (c *InMemorySwitchOverCache) cleanupUnusedFeatureFlags(routerCfg *nodev1.RouterConfig) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -136,17 +146,12 @@ func (c *InMemorySwitchOverCache) cleanupUnusedFeatureFlags(routerCfg *nodev1.Ro
 		return
 	}
 
-	ffNames := make(map[string]struct{})
-	for ffName := range routerCfg.FeatureFlagConfigs.ConfigByFeatureFlagName {
-		ffNames[ffName] = struct{}{}
-	}
-
 	for ffName := range c.queriesForFeatureFlag {
 		// Skip the base which is ""
 		if ffName == "" {
 			continue
 		}
-		if _, exists := ffNames[ffName]; !exists {
+		if _, exists := routerCfg.FeatureFlagConfigs.ConfigByFeatureFlagName[ffName]; !exists {
 			delete(c.queriesForFeatureFlag, ffName)
 		}
 	}

router/core/restart_switchover_config_test.go

@@ -3,13 +3,11 @@ package core
 import (
 	"testing"
 
-	"go.uber.org/zap"
-
 	"github.com/dgraph-io/ristretto/v2"
 	"github.com/stretchr/testify/require"
-
 	nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1"
 	"github.com/wundergraph/cosmo/router/pkg/config"
+	"go.uber.org/zap"
 )
 
 func TestInMemorySwitchOverCache_UpdateInMemorySwitchOverCacheForConfigChanges(t *testing.T) {
@@ -184,9 +182,7 @@ func TestInMemorySwitchOverCache_GetPlanCacheForFF(t *testing.T) {
 		}
 
 		result := cache.getPlanCacheForFF("non-existent")
-
-		require.NotNil(t, result)
-		require.Empty(t, result)
+		require.Nil(t, result)
 	})
 
 	t.Run("returns nil when cache is disabled", func(t *testing.T) {
@@ -373,7 +369,7 @@ func TestInMemorySwitchOverCache_ProcessOnConfigChangeRestart(t *testing.T) {
 		cache.queriesForFeatureFlag["ff1"] = mockCache1
 		cache.queriesForFeatureFlag["ff2"] = mockCache2
 
-		cache.processOnConfigChangeRestart()
+		cache.extractQueriesAndOverridePlanCache()
 
 		// Verify both caches have been converted to operation slices
 		require.IsType(t, []*nodev1.Operation{}, cache.queriesForFeatureFlag["ff1"])
@@ -394,7 +390,7 @@ func TestInMemorySwitchOverCache_ProcessOnConfigChangeRestart(t *testing.T) {
 			queriesForFeatureFlag: nil,
 		}
 
-		cache.processOnConfigChangeRestart()
+		cache.extractQueriesAndOverridePlanCache()
 
 		// Should remain nil since processing is skipped
 		require.Nil(t, cache.queriesForFeatureFlag)
@@ -407,7 +403,7 @@ func TestInMemorySwitchOverCache_ProcessOnConfigChangeRestart(t *testing.T) {
 		}
 
 		require.NotPanics(t, func() {
-			cache.processOnConfigChangeRestart()
+			cache.extractQueriesAndOverridePlanCache()
 		})
 
 		require.Empty(t, cache.queriesForFeatureFlag)

router/core/router.go

@@ -597,7 +597,7 @@ func NewRouter(opts ...Option) (*Router, error) {
 
 // newGraphServer creates a new server.
 func (r *Router) newServer(ctx context.Context, cfg *nodev1.RouterConfig) error {
-	server, err := newGraphServer(ctx, r, cfg, r.proxy, r.switchoverConfig)
+	server, err := newGraphServer(ctx, r, cfg, r.proxy)
 	if err != nil {
 		r.logger.Error("Failed to create graph server. Keeping the old server", zap.Error(err))
 		return err

router/core/supervisor.go

@@ -159,7 +159,7 @@ func (rs *RouterSupervisor) Start() error {
 		rs.logger.Debug("Got shutdown signal", zap.Bool("shutdown", shutdown))
 
 		if !shutdown {
-			rs.router.switchoverConfig.ProcessOnConfigChangeRestart()
+			rs.router.switchoverConfig.ProcessOnConfigChangeRouterInstanceRestart()
 		}
 
 		if err := rs.stopRouter(); err != nil {

router/pkg/config/config.go

@@ -987,7 +987,7 @@ type CacheWarmupConfiguration struct {
 	Workers                    int               `yaml:"workers" envDefault:"8" env:"CACHE_WARMUP_WORKERS"`
 	ItemsPerSecond             int               `yaml:"items_per_second" envDefault:"50" env:"CACHE_WARMUP_ITEMS_PER_SECOND"`
 	Timeout                    time.Duration     `yaml:"timeout" envDefault:"30s" env:"CACHE_WARMUP_TIMEOUT"`
-	InMemorySwitchoverFallback bool              `yaml:"in_memory_switchover_fallback" envDefault:"true" env:"CACHE_WARMUP_IN_MEMORY_SWITCHOVER_FALLBACK_ENABLED"`
+	InMemorySwitchoverFallback bool              `yaml:"in_memory_switchover_fallback" envDefault:"true" env:"CACHE_WARMUP_IN_MEMORY_SWITCHOVER_FALLBACK"`
 }
 
 type MCPConfiguration struct {