Avoid re-computing implementers_map in QueryHashVisitor (#6570)

apollo-router/src/axum_factory/tests.rs

@@ -2273,7 +2273,14 @@ async fn test_supergraph_timeout() {
 
     // we do the entire supergraph rebuilding instead of using `from_supergraph_mock_callback_and_configuration`
     // because we need the plugins to apply on the supergraph
-    let mut plugins = create_plugins(&conf, &schema, planner.subgraph_schemas(), None, None)
+    let subgraph_schemas = Arc::new(
+        planner
+            .subgraph_schemas()
+            .iter()
+            .map(|(k, v)| (k.clone(), v.schema.clone()))
+            .collect(),
+    );
+    let mut plugins = create_plugins(&conf, &schema, subgraph_schemas, None, None)
         .await
         .unwrap();
 

apollo-router/src/plugin/mod.rs

@@ -45,7 +45,6 @@ use tower::ServiceBuilder;
 use crate::graphql;
 use crate::layers::ServiceBuilderExt;
 use crate::notification::Notify;
-use crate::query_planner::fetch::SubgraphSchemas;
 use crate::router_factory::Endpoint;
 use crate::services::execution;
 use crate::services::router;
@@ -75,7 +74,7 @@ pub struct PluginInit<T> {
     pub(crate) supergraph_schema: Arc<Valid<Schema>>,
 
     /// The parsed subgraph schemas from the query planner, keyed by subgraph name
-    pub(crate) subgraph_schemas: Arc<SubgraphSchemas>,
+    pub(crate) subgraph_schemas: Arc<HashMap<String, Arc<Valid<Schema>>>>,
 
     /// Launch ID
     pub(crate) launch_id: Option<Arc<String>>,
@@ -176,7 +175,7 @@ where
         supergraph_sdl: Arc<String>,
         supergraph_schema_id: Arc<String>,
         supergraph_schema: Arc<Valid<Schema>>,
-        subgraph_schemas: Option<Arc<SubgraphSchemas>>,
+        subgraph_schemas: Option<Arc<HashMap<String, Arc<Valid<Schema>>>>>,
         launch_id: Option<Option<Arc<String>>>,
         notify: Notify<String, graphql::Response>,
     ) -> Self {
@@ -201,7 +200,7 @@ where
         supergraph_sdl: Arc<String>,
         supergraph_schema_id: Arc<String>,
         supergraph_schema: Arc<Valid<Schema>>,
-        subgraph_schemas: Option<Arc<SubgraphSchemas>>,
+        subgraph_schemas: Option<Arc<HashMap<String, Arc<Valid<Schema>>>>>,
         launch_id: Option<Arc<String>>,
         notify: Notify<String, graphql::Response>,
     ) -> Result<Self, BoxError> {
@@ -224,7 +223,7 @@ where
         supergraph_sdl: Option<Arc<String>>,
         supergraph_schema_id: Option<Arc<String>>,
         supergraph_schema: Option<Arc<Valid<Schema>>>,
-        subgraph_schemas: Option<Arc<SubgraphSchemas>>,
+        subgraph_schemas: Option<Arc<HashMap<String, Arc<Valid<Schema>>>>>,
         launch_id: Option<Arc<String>>,
         notify: Option<Notify<String, graphql::Response>>,
     ) -> Self {

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

@@ -742,7 +742,7 @@ mod tests {
         let mut demand_controlled_subgraph_schemas = HashMap::new();
         for (subgraph_name, subgraph_schema) in planner.subgraph_schemas().iter() {
             let demand_controlled_subgraph_schema =
-                DemandControlledSchema::new(subgraph_schema.clone()).unwrap();
+                DemandControlledSchema::new(subgraph_schema.schema.clone()).unwrap();
             demand_controlled_subgraph_schemas
                 .insert(subgraph_name.to_string(), demand_controlled_subgraph_schema);
         }

apollo-router/src/plugins/include_subgraph_errors.rs

@@ -214,7 +214,13 @@ mod test {
             .await
             .unwrap();
         let schema = planner.schema();
-        let subgraph_schemas = planner.subgraph_schemas();
+        let subgraph_schemas = Arc::new(
+            planner
+                .subgraph_schemas()
+                .iter()
+                .map(|(k, v)| (k.clone(), v.schema.clone()))
+                .collect(),
+        );
 
         let builder = PluggableSupergraphServiceBuilder::new(planner);
 

apollo-router/src/plugins/test.rs

@@ -110,7 +110,12 @@ impl<T: Into<Box<dyn DynPlugin + 'static>> + 'static> PluginTestHarness<T> {
             .supergraph_schema_id(crate::spec::Schema::schema_id(&supergraph_sdl).into())
             .supergraph_sdl(supergraph_sdl)
             .supergraph_schema(Arc::new(parsed_schema))
-            .subgraph_schemas(subgraph_schemas)
+            .subgraph_schemas(Arc::new(
+                subgraph_schemas
+                    .iter()
+                    .map(|(k, v)| (k.clone(), v.schema.clone()))
+                    .collect(),
+            ))
             .notify(Notify::default())
             .build();
 

apollo-router/src/plugins/traffic_shaping/mod.rs

@@ -531,7 +531,13 @@ mod test {
         let planner = QueryPlannerService::new(schema.clone(), config.clone())
             .await
             .unwrap();
-        let subgraph_schemas = planner.subgraph_schemas();
+        let subgraph_schemas = Arc::new(
+            planner
+                .subgraph_schemas()
+                .iter()
+                .map(|(k, v)| (k.clone(), v.schema.clone()))
+                .collect(),
+        );
 
         let mut builder =
             PluggableSupergraphServiceBuilder::new(planner).with_configuration(config.clone());

apollo-router/src/query_planner/caching_query_planner.rs

@@ -1,10 +1,8 @@
-use std::collections::HashMap;
 use std::hash::Hash;
 use std::hash::Hasher;
 use std::sync::Arc;
 use std::task;
 
-use apollo_compiler::validation::Valid;
 use futures::future::BoxFuture;
 use indexmap::IndexMap;
 use query_planner::QueryPlannerPlugin;
@@ -61,7 +59,7 @@ pub(crate) struct CachingQueryPlanner<T: Clone> {
     >,
     delegate: T,
     schema: Arc<Schema>,
-    subgraph_schemas: Arc<HashMap<String, Arc<Valid<apollo_compiler::Schema>>>>,
+    subgraph_schemas: Arc<SubgraphSchemas>,
     plugins: Arc<Plugins>,
     enable_authorization_directives: bool,
     config_mode_hash: Arc<QueryHash>,
@@ -94,7 +92,7 @@ where
     pub(crate) async fn new(
         delegate: T,
         schema: Arc<Schema>,
-        subgraph_schemas: Arc<HashMap<String, Arc<Valid<apollo_compiler::Schema>>>>,
+        subgraph_schemas: Arc<SubgraphSchemas>,
         configuration: &Configuration,
         plugins: Plugins,
     ) -> Result<CachingQueryPlanner<T>, BoxError> {
@@ -339,9 +337,7 @@ where
 }
 
 impl CachingQueryPlanner<QueryPlannerService> {
-    pub(crate) fn subgraph_schemas(
-        &self,
-    ) -> Arc<HashMap<String, Arc<Valid<apollo_compiler::Schema>>>> {
+    pub(crate) fn subgraph_schemas(&self) -> Arc<SubgraphSchemas> {
         self.delegate.subgraph_schemas()
     }
 

apollo-router/src/query_planner/execution.rs

@@ -1,7 +1,6 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use apollo_compiler::validation::Valid;
 use futures::future::join_all;
 use futures::prelude::*;
 use tokio::sync::broadcast;
@@ -23,6 +22,7 @@ use crate::json_ext::Path;
 use crate::json_ext::Value;
 use crate::json_ext::ValueExt;
 use crate::plugins::subscription::SubscriptionConfig;
+use crate::query_planner::fetch::SubgraphSchemas;
 use crate::query_planner::FlattenNode;
 use crate::query_planner::Primary;
 use crate::query_planner::CONDITION_ELSE_SPAN_NAME;
@@ -50,7 +50,7 @@ impl QueryPlan {
         service_factory: &'a Arc<SubgraphServiceFactory>,
         supergraph_request: &'a Arc<http::Request<Request>>,
         schema: &'a Arc<Schema>,
-        subgraph_schemas: &'a Arc<HashMap<String, Arc<Valid<apollo_compiler::Schema>>>>,
+        subgraph_schemas: &'a Arc<SubgraphSchemas>,
         sender: mpsc::Sender<Response>,
         subscription_handle: Option<SubscriptionHandle>,
         subscription_config: &'a Option<SubscriptionConfig>,
@@ -106,7 +106,7 @@ pub(crate) struct ExecutionParameters<'a> {
     pub(crate) context: &'a Context,
     pub(crate) service_factory: &'a Arc<SubgraphServiceFactory>,
     pub(crate) schema: &'a Arc<Schema>,
-    pub(crate) subgraph_schemas: &'a Arc<HashMap<String, Arc<Valid<apollo_compiler::Schema>>>>,
+    pub(crate) subgraph_schemas: &'a Arc<SubgraphSchemas>,
     pub(crate) supergraph_request: &'a Arc<http::Request<Request>>,
     pub(crate) deferred_fetches: &'a HashMap<String, broadcast::Sender<(Value, Vec<Error>)>>,
     pub(crate) query: &'a Arc<Query>,

apollo-router/src/query_planner/fetch.rs

@@ -1,10 +1,11 @@
-use std::collections::HashMap;
 use std::fmt::Display;
 use std::sync::Arc;
 
 use apollo_compiler::ast;
+use apollo_compiler::collections::HashMap;
 use apollo_compiler::validation::Valid;
 use apollo_compiler::ExecutableDocument;
+use apollo_compiler::Name;
 use indexmap::IndexSet;
 use serde::Deserialize;
 use serde::Serialize;
@@ -93,7 +94,12 @@ impl From<ast::OperationType> for OperationKind {
     }
 }
 
-pub(crate) type SubgraphSchemas = HashMap<String, Arc<Valid<apollo_compiler::Schema>>>;
+pub(crate) type SubgraphSchemas = HashMap<String, SubgraphSchema>;
+
+pub(crate) struct SubgraphSchema {
+    pub(crate) schema: Arc<Valid<apollo_compiler::Schema>>,
+    pub(crate) implementers_map: HashMap<Name, apollo_compiler::schema::Implementers>,
+}
 
 /// A fetch node.
 #[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
@@ -406,7 +412,7 @@ impl FetchNode {
             if let Some(subgraph_schema) =
                 parameters.subgraph_schemas.get(&service_name.to_string())
             {
-                match build_operation_with_aliasing(operation, &ctx_arg, subgraph_schema) {
+                match build_operation_with_aliasing(operation, &ctx_arg, &subgraph_schema.schema) {
                     Ok(op) => {
                         alias_query_string = op.serialize().no_indent().to_string();
                         alias_query_string.as_str()
@@ -680,7 +686,7 @@ impl FetchNode {
         subgraph_schemas: &SubgraphSchemas,
     ) -> Result<(), ValidationErrors> {
         let schema = &subgraph_schemas[self.service_name.as_ref()];
-        self.operation.init_parsed(schema)?;
+        self.operation.init_parsed(&schema.schema)?;
         Ok(())
     }
 
@@ -690,11 +696,12 @@ impl FetchNode {
         supergraph_schema_hash: &str,
     ) -> Result<(), ValidationErrors> {
         let schema = &subgraph_schemas[self.service_name.as_ref()];
-        let doc = self.operation.init_parsed(schema)?;
+        let doc = self.operation.init_parsed(&schema.schema)?;
 
         if let Ok(hash) = QueryHashVisitor::hash_query(
-            schema,
+            &schema.schema,
             supergraph_schema_hash,
+            &schema.implementers_map,
             doc,
             self.operation_name.as_deref(),
         ) {

apollo-router/src/query_planner/query_planner_service.rs

@@ -1,6 +1,5 @@
 //! Calls out to the apollo-federation crate
 
-use std::collections::HashMap;
 use std::fmt::Debug;
 use std::ops::ControlFlow;
 use std::sync::Arc;
@@ -9,7 +8,6 @@ use std::task::Poll;
 use std::time::Instant;
 
 use apollo_compiler::ast;
-use apollo_compiler::validation::Valid;
 use apollo_compiler::Name;
 use apollo_federation::error::FederationError;
 use apollo_federation::error::SingleFederationError;
@@ -42,6 +40,8 @@ use crate::plugins::telemetry::config::ApolloSignatureNormalizationAlgorithm;
 use crate::plugins::telemetry::config::Conf as TelemetryConfig;
 use crate::query_planner::convert::convert_root_query_plan_node;
 use crate::query_planner::fetch::QueryHash;
+use crate::query_planner::fetch::SubgraphSchema;
+use crate::query_planner::fetch::SubgraphSchemas;
 use crate::query_planner::labeler::add_defer_labels;
 use crate::services::layers::query_analysis::ParsedDocument;
 use crate::services::layers::query_analysis::ParsedDocumentInner;
@@ -67,7 +67,7 @@ const INTERNAL_INIT_ERROR: &str = "internal";
 pub(crate) struct QueryPlannerService {
     planner: Arc<QueryPlanner>,
     schema: Arc<Schema>,
-    subgraph_schemas: Arc<HashMap<String, Arc<Valid<apollo_compiler::Schema>>>>,
+    subgraph_schemas: Arc<SubgraphSchemas>,
     configuration: Arc<Configuration>,
     enable_authorization_directives: bool,
     _federation_instrument: ObservableGauge<u64>,
@@ -191,7 +191,15 @@ impl QueryPlannerService {
             planner
                 .subgraph_schemas()
                 .iter()
-                .map(|(name, schema)| (name.to_string(), Arc::new(schema.schema().clone())))
+                .map(|(name, schema)| {
+                    (
+                        name.to_string(),
+                        SubgraphSchema {
+                            implementers_map: schema.schema().implementers_map(),
+                            schema: Arc::new(schema.schema().clone()),
+                        },
+                    )
+                })
                 .collect(),
         );
 
@@ -218,9 +226,7 @@ impl QueryPlannerService {
         self.schema.clone()
     }
 
-    pub(crate) fn subgraph_schemas(
-        &self,
-    ) -> Arc<HashMap<String, Arc<Valid<apollo_compiler::Schema>>>> {
+    pub(crate) fn subgraph_schemas(&self) -> Arc<SubgraphSchemas> {
         self.subgraph_schemas.clone()
     }
 
@@ -383,6 +389,7 @@ impl Service<QueryPlannerRequest> for QueryPlannerService {
                     let hash = QueryHashVisitor::hash_query(
                         this.schema.supergraph_schema(),
                         &this.schema.raw_sdl,
+                        &this.schema.implementers_map,
                         &executable_document,
                         operation_name.as_deref(),
                     )
@@ -508,6 +515,7 @@ impl QueryPlannerService {
             let hash = QueryHashVisitor::hash_query(
                 self.schema.supergraph_schema(),
                 &self.schema.raw_sdl,
+                &self.schema.implementers_map,
                 &executable_document,
                 key.operation_name.as_deref(),
             )
@@ -595,6 +603,8 @@ pub(crate) fn metric_rust_qp_init(init_error_kind: Option<&'static str>) {
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashMap;
+
     use test_log::test;
     use tower::ServiceExt;
 

apollo-router/src/query_planner/subscription.rs

@@ -280,7 +280,7 @@ impl SubscriptionNode {
         subgraph_schemas: &SubgraphSchemas,
     ) -> Result<(), ValidationErrors> {
         let schema = &subgraph_schemas[self.service_name.as_ref()];
-        self.operation.init_parsed(schema)?;
+        self.operation.init_parsed(&schema.schema)?;
         Ok(())
     }
 }

apollo-router/src/query_planner/tests.rs

@@ -1838,8 +1838,14 @@ fn broken_plan_does_not_panic() {
         estimated_size: Default::default(),
     };
     let subgraph_schema = apollo_compiler::Schema::parse_and_validate(subgraph_schema, "").unwrap();
-    let mut subgraph_schemas = HashMap::new();
-    subgraph_schemas.insert("X".to_owned(), Arc::new(subgraph_schema));
+    let mut subgraph_schemas = HashMap::default();
+    subgraph_schemas.insert(
+        "X".to_owned(),
+        query_planner::fetch::SubgraphSchema {
+            implementers_map: subgraph_schema.implementers_map(),
+            schema: Arc::new(subgraph_schema),
+        },
+    );
     let result = Arc::make_mut(&mut plan.root)
         .init_parsed_operations_and_hash_subqueries(&subgraph_schemas, "");
     assert_eq!(

apollo-router/src/router_factory.rs

@@ -326,11 +326,19 @@ impl YamlRouterFactory {
         let span = tracing::info_span!("plugins");
 
         // Process the plugins.
+        let subgraph_schemas = Arc::new(
+            planner
+                .subgraph_schemas()
+                .iter()
+                .map(|(k, v)| (k.clone(), v.schema.clone()))
+                .collect(),
+        );
+
         let plugins: Arc<Plugins> = Arc::new(
             create_plugins(
                 &configuration,
                 &schema,
-                planner.subgraph_schemas(),
+                subgraph_schemas,
                 initial_telemetry_plugin,
                 extra_plugins,
             )