feat: rewarm cache from planner cache on schema change

connect/src/wg/cosmo/node/v1/node_pb.ts

@@ -543,6 +543,11 @@ export class RegistrationInfo extends Message<RegistrationInfo> {
    */
   graphPublicKey = "";
 
+  /**
+   * @generated from field: bool is_cache_warmer_enabled = 3;
+   */
+  isCacheWarmerEnabled = false;
+
   constructor(data?: PartialMessage<RegistrationInfo>) {
     super();
     proto3.util.initPartial(data, this);
@@ -553,6 +558,7 @@ export class RegistrationInfo extends Message<RegistrationInfo> {
   static readonly fields: FieldList = proto3.util.newFieldList(() => [
     { no: 1, name: "account_limits", kind: "message", T: AccountLimits },
     { no: 2, name: "graph_public_key", kind: "scalar", T: 9 /* ScalarType.STRING */ },
+    { no: 3, name: "is_cache_warmer_enabled", kind: "scalar", T: 8 /* ScalarType.BOOL */ },
   ]);
 
   static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): RegistrationInfo {

controlplane/src/core/bufservices/NodeService.ts

@@ -55,6 +55,7 @@ export default function (opts: RouterOptions): Partial<ServiceImpl<typeof NodeSe
             traceSamplingRate: (features['trace-sampling-rate'] as number) ?? 0.1,
           },
           graphPublicKey: publicKey,
+          isCacheWarmerEnabled: !!features['cache-warmer'],
         };
 
         registrationInfoCache.set(authContext.federatedGraphId, registrationInfo);

proto/wg/cosmo/node/v1/node.proto

@@ -43,6 +43,7 @@ message ResponseStatus {
 message RegistrationInfo {
   AccountLimits account_limits = 1;
   string graph_public_key = 2;
+  bool is_cache_warmer_enabled = 3;
 }
 
 message AccountLimits {

router-tests/cache_warmup_test.go

@@ -3,25 +3,25 @@ package integration
 import (
 	"context"
 	"net/http"
+	"os"
+	"path/filepath"
+	"syscall"
 	"testing"
 	"time"
 
-	"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/assert"
 	"github.com/stretchr/testify/require"
-	"go.uber.org/zap"
-
 	"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"
+	"go.opentelemetry.io/otel/sdk/metric/metricdata/metricdatatest"
+	"go.uber.org/zap"
 )
 
 func TestCacheWarmup(t *testing.T) {
@@ -914,6 +914,243 @@ 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,
+					InMemorySwitchoverFallback: 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, 200, res.Response.StatusCode)
+			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, 200, res.Response.StatusCode)
+			require.Equal(t, "updated", res.Response.Header.Get("X-Router-Config-Version"))
+			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, 200, res.Response.StatusCode)
+			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, 200, res.Response.StatusCode)
+			require.Equal(t, "updated", res.Response.Header.Get("X-Router-Config-Version"))
+			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,
+					InMemorySwitchoverFallback: 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, 200, res.Response.StatusCode)
+			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, 200, res.Response.StatusCode)
+			require.Equal(t, "updated", res.Response.Header.Get("X-Router-Config-Version"))
+			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,
+					InMemorySwitchoverFallback: true,
+				}),
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+				Query: `query { hello }`,
+			})
+			require.Equal(t, 200, res.Response.StatusCode)
+			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)
+		})
+	})
+
+	t.Run("Successfully persists cache across config change restarts", func(t *testing.T) {
+		t.Parallel()
+
+		updateConfig := func(t *testing.T, xEnv *testenv.Environment, ctx context.Context, listenString string, config string) {
+			f, err := os.Create(filepath.Join(xEnv.GetRouterProcessCwd(), "config.yaml"))
+			require.NoError(t, err)
+
+			_, err = f.WriteString(config)
+			require.NoError(t, err)
+			require.NoError(t, f.Close())
+
+			err = xEnv.SignalRouterProcess(syscall.SIGHUP)
+			require.NoError(t, err)
+			require.NoError(t, xEnv.WaitForServer(ctx, xEnv.RouterURL+"/"+listenString, 600, 60), "healthcheck post-reload failed")
+		}
+
+		getConfigString := func(listenString string) string {
+			return `
+version: "1"
+
+readiness_check_path: "/` + listenString + `"
+
+cache_warmup:
+  enabled: true
+  in_memory_switchover_fallback: true
+
+engine: 
+  debug:
+    enable_cache_response_headers: true
+`
+		}
+
+		err := testenv.RunRouterBinary(t, &testenv.Config{
+			DemoMode: true,
+		}, testenv.RunRouterBinConfigOptions{}, func(t *testing.T, xEnv *testenv.Environment) {
+			// Verify initial start
+			t.Logf("running router binary, cwd: %s", xEnv.GetRouterProcessCwd())
+			ctx := context.Background()
+			require.NoError(t, xEnv.WaitForServer(ctx, xEnv.RouterURL+"/health/ready", 600, 60), "healthcheck pre-reload failed")
+
+			// Enable cache response headers first
+			listenString1 := "after1"
+			updateConfig(t, xEnv, ctx, listenString1, getConfigString(listenString1))
+			res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{Query: `query { hello }`})
+			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 }`})
+			require.Equal(t, "HIT", res.Response.Header.Get("x-wg-execution-plan-cache"))
+		})
+
+		require.NoError(t, err)
+	})
+
+}
+
 // 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

@@ -63,7 +63,7 @@ require (
 	github.com/coreos/go-systemd/v22 v22.5.0 // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.7 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
-	github.com/dgraph-io/ristretto/v2 v2.1.0 // indirect
+	github.com/dgraph-io/ristretto/v2 v2.4.0 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/docker/cli v28.2.2+incompatible // indirect
 	github.com/docker/distribution v2.8.3+incompatible // indirect

router-tests/go.sum

@@ -65,10 +65,10 @@ 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/dgraph-io/ristretto/v2 v2.4.0 h1:I/w09yLjhdcVD2QV192UJcq8dPBaAJb9pOuMyNy0XlU=
+github.com/dgraph-io/ristretto/v2 v2.4.0/go.mod h1:0KsrXtXvnv0EqnzyowllbVJB8yBonswa2lTCK2gGo9E=
+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=

router/cmd/main.go

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

router/core/cache_warmup_cdn.go

@@ -5,19 +5,19 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1"
-	"google.golang.org/protobuf/encoding/protojson"
 	"io"
 	"net/http"
 	"net/url"
 
+	nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1"
 	"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)