initial stab at making the router prewarm on pqs

.changesets/feat_pq_warmup.md

@@ -0,0 +1,5 @@
+### 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.
+
+By [@lleadbet](https://github.com/lleadbet) in https://github.com/apollographql/router/pull/5340

apollo-router/src/query_planner/caching_query_planner.rs

@@ -153,7 +153,7 @@ 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,
     ) {
@@ -172,38 +172,41 @@ 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: _,
-                            sdl: _,
-                            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: _,
+                                sdl: _,
+                                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());
@@ -269,19 +272,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

@@ -250,7 +250,20 @@ 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,
+                )
+                .await;
+        } else {
+            supergraph_creator
+                .warm_up_query_planner(
+                    &query_analysis_layer,
+                    &persisted_query_layer,
+                    None,
                     configuration.supergraph.query_planning.warmed_up_queries,
                     configuration
                         .supergraph

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

@@ -909,7 +909,7 @@ 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,
     ) {