diff --git a/apollo-router/src/plugins/demand_control/cost_calculator/static_cost.rs b/apollo-router/src/plugins/demand_control/cost_calculator/static_cost.rs index 1156c6df7d..977ecd4757 100644 --- a/apollo-router/src/plugins/demand_control/cost_calculator/static_cost.rs +++ b/apollo-router/src/plugins/demand_control/cost_calculator/static_cost.rs @@ -464,7 +464,7 @@ mod tests { let config: Arc = 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(); diff --git a/apollo-router/src/plugins/test.rs b/apollo-router/src/plugins/test.rs index 523e07f253..700e1ce921 100644 --- a/apollo-router/src/plugins/test.rs +++ b/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 PluginTestHarness { 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 { ( diff --git a/apollo-router/src/query_planner/bridge_query_planner.rs b/apollo-router/src/query_planner/bridge_query_planner.rs index f2348ef11b..d4fcb1331c 100644 --- a/apollo-router/src/query_planner/bridge_query_planner.rs +++ b/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>), Both { js: Arc>, @@ -115,6 +115,7 @@ impl PlannerMode { schema: &Schema, configuration: &Configuration, old_planner: Option>>, + rust_planner: Option>, ) -> Result { 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>, 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, configuration: Arc, - old_planner: Option>>, + old_js_planner: Option>>, + rust_planner: Option>, ) -> Result { - 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(); diff --git a/apollo-router/src/query_planner/bridge_query_planner_pool.rs b/apollo-router/src/query_planner/bridge_query_planner_pool.rs index 19e80b539e..a306f19b6b 100644 --- a/apollo-router/src/query_planner/bridge_query_planner_pool.rs +++ b/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>>, + js_planners: Vec>>, sender: Sender<( QueryPlannerRequest, oneshot::Sender>, @@ -48,11 +49,13 @@ impl BridgeQueryPlannerPool { } pub(crate) async fn new_from_planners( - old_planners: Vec>>, + old_js_planners: Vec>>, schema: Arc, configuration: Arc, size: NonZeroUsize, ) -> Result { + let rust_planner = PlannerMode::maybe_rust(&schema, &configuration)?; + let mut join_set = JoinSet::new(); let (sender, receiver) = bounded::<( @@ -60,15 +63,16 @@ impl BridgeQueryPlannerPool { oneshot::Sender>, )>(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>> { - self.planners.clone() + self.js_planners.clone() } pub(crate) fn schema(&self) -> Arc {