Share a single Rust QP instance when using a pool of JS workers

apollo-router/src/plugins/demand_control/cost_calculator/static_cost.rs

@@ -464,7 +464,7 @@ mod tests {
         let config: Arc<Configuration> = Arc::new(Default::default());
         let (schema, query) = parse_schema_and_operation(schema_str, query_str, &config);
 
-        let mut planner = BridgeQueryPlanner::new(schema.into(), config.clone(), None)
+        let mut planner = BridgeQueryPlanner::new(schema.into(), config.clone(), None, None)
             .await
             .unwrap();
 

apollo-router/src/plugins/test.rs

@@ -13,6 +13,7 @@ use crate::plugin::DynPlugin;
 use crate::plugin::Plugin;
 use crate::plugin::PluginInit;
 use crate::query_planner::BridgeQueryPlanner;
+use crate::query_planner::PlannerMode;
 use crate::services::execution;
 use crate::services::http;
 use crate::services::router;
@@ -92,9 +93,11 @@ impl<T: Plugin> PluginTestHarness<T> {
             let schema = Schema::parse(schema, &config).unwrap();
             let sdl = schema.raw_sdl.clone();
             let supergraph = schema.supergraph_schema().clone();
-            let planner = BridgeQueryPlanner::new(schema.into(), Arc::new(config), None)
-                .await
-                .unwrap();
+            let rust_planner = PlannerMode::maybe_rust(&schema, &config).unwrap();
+            let planner =
+                BridgeQueryPlanner::new(schema.into(), Arc::new(config), None, rust_planner)
+                    .await
+                    .unwrap();
             (sdl, supergraph, planner.subgraph_schemas())
         } else {
             (

apollo-router/src/query_planner/bridge_query_planner.rs

@@ -81,7 +81,7 @@ pub(crate) struct BridgeQueryPlanner {
 }
 
 #[derive(Clone)]
-enum PlannerMode {
+pub(crate) enum PlannerMode {
     Js(Arc<Planner<QueryPlanResult>>),
     Both {
         js: Arc<Planner<QueryPlanResult>>,
@@ -115,6 +115,7 @@ impl PlannerMode {
         schema: &Schema,
         configuration: &Configuration,
         old_planner: Option<Arc<Planner<QueryPlanResult>>>,
+        rust_planner: Option<Arc<QueryPlanner>>,
     ) -> Result<Self, ServiceBuildError> {
         Ok(match configuration.experimental_query_planner_mode {
             QueryPlannerMode::New => Self::Rust {
@@ -124,18 +125,33 @@ impl PlannerMode {
                     old_planner,
                 )
                 .await?,
-                rust: Self::rust(schema, configuration)?,
+                rust: rust_planner
+                    .expect("expected Rust QP instance for `experimental_query_planner_mode: new`"),
             },
             QueryPlannerMode::Legacy => {
                 Self::Js(Self::js(&schema.raw_sdl, configuration, old_planner).await?)
             }
             QueryPlannerMode::Both => Self::Both {
                 js: Self::js(&schema.raw_sdl, configuration, old_planner).await?,
-                rust: Self::rust(schema, configuration)?,
+                rust: rust_planner.expect(
+                    "expected Rust QP instance for `experimental_query_planner_mode: both`",
+                ),
             },
         })
     }
 
+    pub(crate) fn maybe_rust(
+        schema: &Schema,
+        configuration: &Configuration,
+    ) -> Result<Option<Arc<QueryPlanner>>, ServiceBuildError> {
+        match configuration.experimental_query_planner_mode {
+            QueryPlannerMode::Legacy => Ok(None),
+            QueryPlannerMode::New | QueryPlannerMode::Both => {
+                Ok(Some(Self::rust(schema, configuration)?))
+            }
+        }
+    }
+
     fn rust(
         schema: &Schema,
         configuration: &Configuration,
@@ -325,9 +341,11 @@ impl BridgeQueryPlanner {
     pub(crate) async fn new(
         schema: Arc<Schema>,
         configuration: Arc<Configuration>,
-        old_planner: Option<Arc<Planner<QueryPlanResult>>>,
+        old_js_planner: Option<Arc<Planner<QueryPlanResult>>>,
+        rust_planner: Option<Arc<QueryPlanner>>,
     ) -> Result<Self, ServiceBuildError> {
-        let planner = PlannerMode::new(&schema, &configuration, old_planner).await?;
+        let planner =
+            PlannerMode::new(&schema, &configuration, old_js_planner, rust_planner).await?;
 
         let subgraph_schemas = Arc::new(planner.subgraphs().await?);
 
@@ -991,7 +1009,7 @@ mod tests {
             let sdl = include_str!("../testdata/minimal_supergraph.graphql");
             let config = Arc::default();
             let schema = Schema::parse(sdl, &config).unwrap();
-            let _planner = BridgeQueryPlanner::new(schema.into(), config, None)
+            let _planner = BridgeQueryPlanner::new(schema.into(), config, None, None)
                 .await
                 .unwrap();
 
@@ -1008,7 +1026,7 @@ mod tests {
             let sdl = include_str!("../testdata/minimal_fed2_supergraph.graphql");
             let config = Arc::default();
             let schema = Schema::parse(sdl, &config).unwrap();
-            let _planner = BridgeQueryPlanner::new(schema.into(), config, None)
+            let _planner = BridgeQueryPlanner::new(schema.into(), config, None, None)
                 .await
                 .unwrap();
 
@@ -1027,7 +1045,7 @@ mod tests {
         let schema = Arc::new(Schema::parse(EXAMPLE_SCHEMA, &Default::default()).unwrap());
         let query = include_str!("testdata/unknown_introspection_query.graphql");
 
-        let planner = BridgeQueryPlanner::new(schema.clone(), Default::default(), None)
+        let planner = BridgeQueryPlanner::new(schema.clone(), Default::default(), None, None)
             .await
             .unwrap();
 
@@ -1127,7 +1145,7 @@ mod tests {
         let configuration = Arc::new(configuration);
 
         let schema = Schema::parse(EXAMPLE_SCHEMA, &configuration).unwrap();
-        let planner = BridgeQueryPlanner::new(schema.into(), configuration.clone(), None)
+        let planner = BridgeQueryPlanner::new(schema.into(), configuration.clone(), None, None)
             .await
             .unwrap();
 
@@ -1435,7 +1453,7 @@ mod tests {
         let configuration = Arc::new(configuration);
 
         let schema = Schema::parse(schema, &configuration).unwrap();
-        let planner = BridgeQueryPlanner::new(schema.into(), configuration.clone(), None)
+        let planner = BridgeQueryPlanner::new(schema.into(), configuration.clone(), None, None)
             .await
             .unwrap();
 

apollo-router/src/query_planner/bridge_query_planner_pool.rs

@@ -19,6 +19,7 @@ use super::QueryPlanResult;
 use crate::error::QueryPlannerError;
 use crate::error::ServiceBuildError;
 use crate::metrics::meter_provider;
+use crate::query_planner::PlannerMode;
 use crate::services::QueryPlannerRequest;
 use crate::services::QueryPlannerResponse;
 use crate::spec::Schema;
@@ -28,7 +29,7 @@ static CHANNEL_SIZE: usize = 1_000;
 
 #[derive(Clone)]
 pub(crate) struct BridgeQueryPlannerPool {
-    planners: Vec<Arc<Planner<QueryPlanResult>>>,
+    js_planners: Vec<Arc<Planner<QueryPlanResult>>>,
     sender: Sender<(
         QueryPlannerRequest,
         oneshot::Sender<Result<QueryPlannerResponse, QueryPlannerError>>,
@@ -48,27 +49,30 @@ impl BridgeQueryPlannerPool {
     }
 
     pub(crate) async fn new_from_planners(
-        old_planners: Vec<Arc<Planner<QueryPlanResult>>>,
+        old_js_planners: Vec<Arc<Planner<QueryPlanResult>>>,
         schema: Arc<Schema>,
         configuration: Arc<Configuration>,
         size: NonZeroUsize,
     ) -> Result<Self, ServiceBuildError> {
+        let rust_planner = PlannerMode::maybe_rust(&schema, &configuration)?;
+
         let mut join_set = JoinSet::new();
 
         let (sender, receiver) = bounded::<(
             QueryPlannerRequest,
             oneshot::Sender<Result<QueryPlannerResponse, QueryPlannerError>>,
         )>(CHANNEL_SIZE);
 
-        let mut old_planners_iterator = old_planners.into_iter();
+        let mut old_js_planners_iterator = old_js_planners.into_iter();
 
         (0..size.into()).for_each(|_| {
             let schema = schema.clone();
             let configuration = configuration.clone();
+            let rust_planner = rust_planner.clone();
 
-            let old_planner = old_planners_iterator.next();
+            let old_planner = old_js_planners_iterator.next();
             join_set.spawn(async move {
-                BridgeQueryPlanner::new(schema, configuration, old_planner).await
+                BridgeQueryPlanner::new(schema, configuration, old_planner, rust_planner).await
             });
         });
 
@@ -122,7 +126,7 @@ impl BridgeQueryPlannerPool {
             .init();
 
         Ok(Self {
-            planners,
+            js_planners: planners,
             sender,
             schema,
             subgraph_schemas,
@@ -131,7 +135,7 @@ impl BridgeQueryPlannerPool {
     }
 
     pub(crate) fn planners(&self) -> Vec<Arc<Planner<QueryPlanResult>>> {
-        self.planners.clone()
+        self.js_planners.clone()
     }
 
     pub(crate) fn schema(&self) -> Arc<Schema> {