initial stab at making the router prewarm on pqs

.changesets/feat_pq_warmup.md

@@ -0,0 +1,13 @@
+### Warm query plan cache using persisted queries on startup ([Issue #5334](https://github.com/apollographql/router/issues/5334))
+
+This adds support for the router to use [persisted queries](https://www.apollographql.com/docs/graphos/operations/persisted-queries/) to warm the query plan cache upon startup using a new `experimental_prewarm_query_plan_cache` configuration option under `persisted_queries`.
+
+To enable:
+
+```yml
+persisted_queries:
+  enabled: true
+  experimental_prewarm_query_plan_cache: true
+```
+
+By [@lleadbet](https://github.com/lleadbet) in https://github.com/apollographql/router/pull/5340

apollo-router/src/configuration/persisted_queries.rs

@@ -14,6 +14,9 @@ pub struct PersistedQueries {
 
     /// Restricts execution of operations that are not found in the Persisted Query List
     pub safelist: PersistedQueriesSafelist,
+
+    /// Experimental feature to prewarm the query plan cache with persisted queries
+    pub experimental_prewarm_query_plan_cache: bool,
 }
 
 #[cfg(test)]
@@ -24,11 +27,14 @@ impl PersistedQueries {
         enabled: Option<bool>,
         log_unknown: Option<bool>,
         safelist: Option<PersistedQueriesSafelist>,
+        experimental_prewarm_query_plan_cache: Option<bool>,
     ) -> Self {
         Self {
             enabled: enabled.unwrap_or_else(default_pq),
             safelist: safelist.unwrap_or_default(),
             log_unknown: log_unknown.unwrap_or_else(default_log_unknown),
+            experimental_prewarm_query_plan_cache: experimental_prewarm_query_plan_cache
+                .unwrap_or_else(default_prewarm_query_plan_cache),
         }
     }
 }
@@ -62,6 +68,7 @@ impl Default for PersistedQueries {
             enabled: default_pq(),
             safelist: PersistedQueriesSafelist::default(),
             log_unknown: default_log_unknown(),
+            experimental_prewarm_query_plan_cache: default_prewarm_query_plan_cache(),
         }
     }
 }
@@ -90,3 +97,7 @@ const fn default_require_id() -> bool {
 const fn default_log_unknown() -> bool {
     false
 }
+
+const fn default_prewarm_query_plan_cache() -> bool {
+    false
+}

apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap

@@ -1,6 +1,5 @@
 ---
 source: apollo-router/src/configuration/tests.rs
-assertion_line: 26
 expression: "&schema"
 ---
 {
@@ -4149,6 +4148,11 @@ expression: "&schema"
           "description": "Activates Persisted Queries (disabled by default)",
           "type": "boolean"
         },
+        "experimental_prewarm_query_plan_cache": {
+          "default": false,
+          "description": "Experimental feature to prewarm the query plan cache with persisted queries",
+          "type": "boolean"
+        },
         "log_unknown": {
           "default": false,
           "description": "Enabling this field configures the router to log any freeform GraphQL request that is not in the persisted query list",

apollo-router/src/query_planner/caching_query_planner.rs

@@ -153,9 +153,10 @@ where
         &mut self,
         query_analysis: &QueryAnalysisLayer,
         persisted_query_layer: &PersistedQueryLayer,
-        previous_cache: InMemoryCachePlanner,
+        previous_cache: Option<InMemoryCachePlanner>,
         count: Option<usize>,
         experimental_reuse_query_plans: bool,
+        experimental_pql_prewarm: bool,
     ) {
         let _timer = Timer::new(|duration| {
             ::tracing::info!(
@@ -172,68 +173,79 @@ where
                 }),
         );
 
-        let mut cache_keys = {
-            let cache = previous_cache.lock().await;
-
-            let count = count.unwrap_or(cache.len() / 3);
-
-            cache
-                .iter()
-                .map(
-                    |(
-                        CachingQueryKey {
-                            query,
-                            operation,
-                            hash,
-                            metadata,
-                            plan_options,
-                            config_mode: _,
-                            schema_id: _,
-                            introspection: _,
+        let mut cache_keys = match previous_cache {
+            Some(ref previous_cache) => {
+                let cache = previous_cache.lock().await;
+
+                let count = count.unwrap_or(cache.len() / 3);
+
+                cache
+                    .iter()
+                    .map(
+                        |(
+                            CachingQueryKey {
+                                query,
+                                operation,
+                                hash,
+                                metadata,
+                                plan_options,
+                                config_mode: _,
+                                schema_id: _,
+                                introspection: _,
+                            },
+                            _,
+                        )| WarmUpCachingQueryKey {
+                            query: query.clone(),
+                            operation: operation.clone(),
+                            hash: Some(hash.clone()),
+                            metadata: metadata.clone(),
+                            plan_options: plan_options.clone(),
+                            config_mode: self.config_mode.clone(),
+                            introspection: self.introspection,
                         },
-                        _,
-                    )| WarmUpCachingQueryKey {
-                        query: query.clone(),
-                        operation: operation.clone(),
-                        hash: Some(hash.clone()),
-                        metadata: metadata.clone(),
-                        plan_options: plan_options.clone(),
-                        config_mode: self.config_mode.clone(),
-                        introspection: self.introspection,
-                    },
-                )
-                .take(count)
-                .collect::<Vec<_>>()
+                    )
+                    .take(count)
+                    .collect::<Vec<_>>()
+            }
+            None => Vec::new(),
         };
 
         cache_keys.shuffle(&mut thread_rng());
 
+        let should_warm_with_pqs =
+            (experimental_pql_prewarm && previous_cache.is_none()) || previous_cache.is_some();
         let persisted_queries_operations = persisted_query_layer.all_operations();
 
-        let capacity = cache_keys.len()
-            + persisted_queries_operations
-                .as_ref()
-                .map(|ops| ops.len())
-                .unwrap_or(0);
+        let capacity = if should_warm_with_pqs {
+            cache_keys.len()
+                + persisted_queries_operations
+                    .as_ref()
+                    .map(|ops| ops.len())
+                    .unwrap_or(0)
+        } else {
+            cache_keys.len()
+        };
         tracing::info!(
             "warming up the query plan cache with {} queries, this might take a while",
             capacity
         );
 
         // persisted queries are added first because they should get a lower priority in the LRU cache,
         // since a lot of them may be there to support old clients
-        let mut all_cache_keys = Vec::with_capacity(capacity);
-        if let Some(queries) = persisted_queries_operations {
-            for query in queries {
-                all_cache_keys.push(WarmUpCachingQueryKey {
-                    query,
-                    operation: None,
-                    hash: None,
-                    metadata: CacheKeyMetadata::default(),
-                    plan_options: PlanOptions::default(),
-                    config_mode: self.config_mode.clone(),
-                    introspection: self.introspection,
-                });
+        let mut all_cache_keys: Vec<WarmUpCachingQueryKey> = Vec::with_capacity(capacity);
+        if should_warm_with_pqs {
+            if let Some(queries) = persisted_queries_operations {
+                for query in queries {
+                    all_cache_keys.push(WarmUpCachingQueryKey {
+                        query,
+                        operation: None,
+                        hash: None,
+                        metadata: CacheKeyMetadata::default(),
+                        plan_options: PlanOptions::default(),
+                        config_mode: self.config_mode.clone(),
+                        introspection: self.introspection,
+                    });
+                }
             }
         }
 
@@ -269,19 +281,22 @@ where
             };
 
             if experimental_reuse_query_plans {
-                // if the query hash did not change with the schema update, we can reuse the previously cached entry
-                if let Some(hash) = hash {
-                    if hash == doc.hash {
-                        if let Some(entry) =
-                            { previous_cache.lock().await.get(&caching_key).cloned() }
-                        {
-                            self.cache.insert_in_memory(caching_key, entry).await;
-                            reused += 1;
-                            continue;
+                // check if prewarming via seeing if the previous cache exists (aka a reloaded router); if reloading, try to reuse the
+                if let Some(ref previous_cache) = previous_cache {
+                    // if the query hash did not change with the schema update, we can reuse the previously cached entry
+                    if let Some(hash) = hash {
+                        if hash == doc.hash {
+                            if let Some(entry) =
+                                { previous_cache.lock().await.get(&caching_key).cloned() }
+                            {
+                                self.cache.insert_in_memory(caching_key, entry).await;
+                                reused += 1;
+                                continue;
+                            }
                         }
                     }
                 }
-            }
+            };
 
             let entry = self
                 .cache

apollo-router/src/router_factory.rs

@@ -262,12 +262,31 @@ impl YamlRouterFactory {
                 .warm_up_query_planner(
                     &query_analysis_layer,
                     &persisted_query_layer,
-                    previous_cache,
+                    Some(previous_cache),
                     configuration.supergraph.query_planning.warmed_up_queries,
                     configuration
                         .supergraph
                         .query_planning
                         .experimental_reuse_query_plans,
+                    configuration
+                        .persisted_queries
+                        .experimental_prewarm_query_plan_cache,
+                )
+                .await;
+        } else {
+            supergraph_creator
+                .warm_up_query_planner(
+                    &query_analysis_layer,
+                    &persisted_query_layer,
+                    None,
+                    configuration.supergraph.query_planning.warmed_up_queries,
+                    configuration
+                        .supergraph
+                        .query_planning
+                        .experimental_reuse_query_plans,
+                    configuration
+                        .persisted_queries
+                        .experimental_prewarm_query_plan_cache,
                 )
                 .await;
         };

apollo-router/src/services/supergraph/service.rs

@@ -905,9 +905,10 @@ impl SupergraphCreator {
         &mut self,
         query_parser: &QueryAnalysisLayer,
         persisted_query_layer: &PersistedQueryLayer,
-        previous_cache: InMemoryCachePlanner,
+        previous_cache: Option<InMemoryCachePlanner>,
         count: Option<usize>,
         experimental_reuse_query_plans: bool,
+        experimental_pql_prewarm: bool,
     ) {
         self.query_planner_service
             .warm_up(
@@ -916,6 +917,7 @@ impl SupergraphCreator {
                 previous_cache,
                 count,
                 experimental_reuse_query_plans,
+                experimental_pql_prewarm,
             )
             .await
     }

docs/source/configuration/persisted-queries.mdx

@@ -66,6 +66,18 @@ persisted_queries:
 
 If used with the [`safelist`](#safelist) option, the router logs unregistered and rejected operations. With [`safelist.required_id`](#require_id) off, the only rejected operations are unregistered ones. If [`safelist.required_id`](#require_id) is turned on, operations can be rejected even when registered because they use operation IDs rather than operation strings.
 
+#### `experimental_prewarm_query_plan_cache`
+
+<ExperimentalFeature />
+
+Adding `experimental_prewarm_query_plan_cache: true` to `persisted_queries` configures the router to prewarm the query plan cache on startup using the persisted query list. All subsequent requests benefit from the warmed cache, reducing latency and enhancing performance.
+
+```yaml title="router.yaml"
+persisted_queries:
+  enabled: true
+  experimental_prewarm_query_plan_cache: true
+```
+
 #### `safelist`
 
 Adding `safelist: true` to `persisted_queries` causes the router to reject any operations that haven't been registered to your PQL.