Demand control lookup optimizations

.changesets/fix_tninesling_demand_control_perf.md

@@ -0,0 +1,9 @@
+### Demand control lookup optimizations ([PR #6450](https://github.com/apollographql/router/pull/6450))
+
+Demand Control can reduce router throughput due to the extra processing required for scoring. This change shifts more data to be computed at plugin initialization and consolidates lookup queries.
+
+- Cost directives for arguments are now stored in a map alongside those for field definitions
+- All precomputed directives are bundled into a struct for each field, along with that field's extended schema type. This reduces 5 individual lookups to a single lookup.
+- Response scoring was looking up each field's definition twice. This is now reduced to a single lookup.
+
+By [@tninesling](https://github.com/tninesling) in https://github.com/apollographql/router/pull/6450

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

@@ -13,112 +13,186 @@ use apollo_federation::link::cost_spec_definition::CostSpecDefinition;
 use apollo_federation::link::cost_spec_definition::ListSizeDirective;
 use apollo_federation::schema::ValidFederationSchema;
 
-use super::directives::RequiresDirective;
+use crate::plugins::demand_control::cost_calculator::directives::RequiresDirective;
 use crate::plugins::demand_control::DemandControlError;
 
+pub(in crate::plugins::demand_control) struct InputDefinition {
+    name: Name,
+    ty: ExtendedType,
+    cost_directive: Option<CostDirective>,
+}
+
+impl InputDefinition {
+    fn new(
+        schema: &ValidFederationSchema,
+        field_definition: &InputValueDefinition,
+    ) -> Result<Self, DemandControlError> {
+        let field_type = schema
+            .schema()
+            .types
+            .get(field_definition.ty.inner_named_type())
+            .ok_or_else(|| {
+                DemandControlError::QueryParseFailure(format!(
+                    "Field {} was found in query, but its type is missing from the schema.",
+                    field_definition.name
+                ))
+            })?;
+        let processed_inputs = InputDefinition {
+            name: field_definition.name.clone(),
+            ty: field_type.clone(),
+            cost_directive: CostSpecDefinition::cost_directive_from_argument(
+                schema,
+                field_definition,
+                field_type,
+            )?,
+        };
+
+        Ok(processed_inputs)
+    }
+
+    pub(in crate::plugins::demand_control) fn name(&self) -> &Name {
+        &self.name
+    }
+
+    pub(in crate::plugins::demand_control) fn ty(&self) -> &ExtendedType {
+        &self.ty
+    }
+
+    pub(in crate::plugins::demand_control) fn cost_directive(&self) -> Option<&CostDirective> {
+        self.cost_directive.as_ref()
+    }
+}
+
+pub(in crate::plugins::demand_control) struct FieldDefinition {
+    ty: ExtendedType,
+    cost_directive: Option<CostDirective>,
+    list_size_directive: Option<ListSizeDirective>,
+    requires_directive: Option<RequiresDirective>,
+    arguments: HashMap<Name, InputDefinition>,
+}
+
+impl FieldDefinition {
+    fn new(
+        schema: &ValidFederationSchema,
+        parent_type_name: &Name,
+        field_definition: &apollo_compiler::ast::FieldDefinition,
+    ) -> Result<Self, DemandControlError> {
+        let field_type = schema
+            .schema()
+            .types
+            .get(field_definition.ty.inner_named_type())
+            .ok_or_else(|| {
+                DemandControlError::QueryParseFailure(format!(
+                    "Field {} was found in query, but its type is missing from the schema.",
+                    field_definition.name,
+                ))
+            })?;
+        let mut processed_field_definition = Self {
+            ty: field_type.clone(),
+            cost_directive: None,
+            list_size_directive: None,
+            requires_directive: None,
+            arguments: HashMap::new(),
+        };
+
+        processed_field_definition.cost_directive =
+            CostSpecDefinition::cost_directive_from_field(schema, field_definition, field_type)?;
+        processed_field_definition.list_size_directive =
+            CostSpecDefinition::list_size_directive_from_field_definition(
+                schema,
+                field_definition,
+            )?;
+        processed_field_definition.requires_directive = RequiresDirective::from_field_definition(
+            field_definition,
+            parent_type_name,
+            schema.schema(),
+        )?;
+
+        for argument in &field_definition.arguments {
+            processed_field_definition.arguments.insert(
+                argument.name.clone(),
+                InputDefinition::new(schema, argument)?,
+            );
+        }
+
+        Ok(processed_field_definition)
+    }
+
+    pub(in crate::plugins::demand_control) fn ty(&self) -> &ExtendedType {
+        &self.ty
+    }
+
+    pub(in crate::plugins::demand_control) fn cost_directive(&self) -> Option<&CostDirective> {
+        self.cost_directive.as_ref()
+    }
+
+    pub(in crate::plugins::demand_control) fn list_size_directive(
+        &self,
+    ) -> Option<&ListSizeDirective> {
+        self.list_size_directive.as_ref()
+    }
+
+    pub(in crate::plugins::demand_control) fn requires_directive(
+        &self,
+    ) -> Option<&RequiresDirective> {
+        self.requires_directive.as_ref()
+    }
+
+    pub(in crate::plugins::demand_control) fn argument_by_name(
+        &self,
+        argument_name: &str,
+    ) -> Option<&InputDefinition> {
+        self.arguments.get(argument_name)
+    }
+}
+
 pub(crate) struct DemandControlledSchema {
     inner: ValidFederationSchema,
-    type_field_cost_directives: HashMap<Name, HashMap<Name, CostDirective>>,
-    type_field_list_size_directives: HashMap<Name, HashMap<Name, ListSizeDirective>>,
-    type_field_requires_directives: HashMap<Name, HashMap<Name, RequiresDirective>>,
+    input_field_definitions: HashMap<Name, HashMap<Name, InputDefinition>>,
+    output_field_definitions: HashMap<Name, HashMap<Name, FieldDefinition>>,
 }
 
 impl DemandControlledSchema {
     pub(crate) fn new(schema: Arc<Valid<Schema>>) -> Result<Self, DemandControlError> {
         let fed_schema = ValidFederationSchema::new((*schema).clone())?;
-        let mut type_field_cost_directives: HashMap<Name, HashMap<Name, CostDirective>> =
+        let mut input_field_definitions: HashMap<Name, HashMap<Name, InputDefinition>> =
             HashMap::new();
-        let mut type_field_list_size_directives: HashMap<Name, HashMap<Name, ListSizeDirective>> =
-            HashMap::new();
-        let mut type_field_requires_directives: HashMap<Name, HashMap<Name, RequiresDirective>> =
+        let mut output_field_definitions: HashMap<Name, HashMap<Name, FieldDefinition>> =
             HashMap::new();
 
         for (type_name, type_) in &schema.types {
-            let field_cost_directives = type_field_cost_directives
-                .entry(type_name.clone())
-                .or_default();
-            let field_list_size_directives = type_field_list_size_directives
-                .entry(type_name.clone())
-                .or_default();
-            let field_requires_directives = type_field_requires_directives
-                .entry(type_name.clone())
-                .or_default();
-
             match type_ {
                 ExtendedType::Interface(ty) => {
-                    for field_name in ty.fields.keys() {
-                        let field_definition = schema.type_field(type_name, field_name)?;
-                        let field_type = schema.types.get(field_definition.ty.inner_named_type()).ok_or_else(|| {
-                            DemandControlError::QueryParseFailure(format!(
-                                "Field {} was found in query, but its type is missing from the schema.",
-                                field_name
-                            ))
-                        })?;
-
-                        if let Some(cost_directive) = CostSpecDefinition::cost_directive_from_field(
-                            &fed_schema,
-                            field_definition,
-                            field_type,
-                        )? {
-                            field_cost_directives.insert(field_name.clone(), cost_directive);
-                        }
-
-                        if let Some(list_size_directive) =
-                            CostSpecDefinition::list_size_directive_from_field_definition(
-                                &fed_schema,
-                                field_definition,
-                            )?
-                        {
-                            field_list_size_directives
-                                .insert(field_name.clone(), list_size_directive);
-                        }
-
-                        if let Some(requires_directive) = RequiresDirective::from_field_definition(
-                            field_definition,
-                            type_name,
-                            &schema,
-                        )? {
-                            field_requires_directives
-                                .insert(field_name.clone(), requires_directive);
-                        }
+                    let type_fields = output_field_definitions
+                        .entry(type_name.clone())
+                        .or_default();
+                    for (field_name, field_definition) in &ty.fields {
+                        type_fields.insert(
+                            field_name.clone(),
+                            FieldDefinition::new(&fed_schema, type_name, field_definition)?,
+                        );
                     }
                 }
                 ExtendedType::Object(ty) => {
-                    for field_name in ty.fields.keys() {
-                        let field_definition = schema.type_field(type_name, field_name)?;
-                        let field_type = schema.types.get(field_definition.ty.inner_named_type()).ok_or_else(|| {
-                            DemandControlError::QueryParseFailure(format!(
-                                "Field {} was found in query, but its type is missing from the schema.",
-                                field_name
-                            ))
-                        })?;
-
-                        if let Some(cost_directive) = CostSpecDefinition::cost_directive_from_field(
-                            &fed_schema,
-                            field_definition,
-                            field_type,
-                        )? {
-                            field_cost_directives.insert(field_name.clone(), cost_directive);
-                        }
-
-                        if let Some(list_size_directive) =
-                            CostSpecDefinition::list_size_directive_from_field_definition(
-                                &fed_schema,
-                                field_definition,
-                            )?
-                        {
-                            field_list_size_directives
-                                .insert(field_name.clone(), list_size_directive);
-                        }
-
-                        if let Some(requires_directive) = RequiresDirective::from_field_definition(
-                            field_definition,
-                            type_name,
-                            &schema,
-                        )? {
-                            field_requires_directives
-                                .insert(field_name.clone(), requires_directive);
-                        }
+                    let type_fields = output_field_definitions
+                        .entry(type_name.clone())
+                        .or_default();
+                    for (field_name, field_definition) in &ty.fields {
+                        type_fields.insert(
+                            field_name.clone(),
+                            FieldDefinition::new(&fed_schema, type_name, field_definition)?,
+                        );
+                    }
+                }
+                ExtendedType::InputObject(ty) => {
+                    let type_fields = input_field_definitions
+                        .entry(type_name.clone())
+                        .or_default();
+                    for (field_name, field_definition) in &ty.fields {
+                        type_fields.insert(
+                            field_name.clone(),
+                            InputDefinition::new(&fed_schema, field_definition)?,
+                        );
                     }
                 }
                 _ => {
@@ -129,54 +203,28 @@ impl DemandControlledSchema {
 
         Ok(Self {
             inner: fed_schema,
-            type_field_cost_directives,
-            type_field_list_size_directives,
-            type_field_requires_directives,
+            input_field_definitions,
+            output_field_definitions,
         })
     }
 
-    pub(in crate::plugins::demand_control) fn type_field_cost_directive(
+    pub(in crate::plugins::demand_control) fn input_field_definition(
         &self,
         type_name: &str,
         field_name: &str,
-    ) -> Option<&CostDirective> {
-        self.type_field_cost_directives
-            .get(type_name)?
-            .get(field_name)
+    ) -> Option<&InputDefinition> {
+        self.input_field_definitions.get(type_name)?.get(field_name)
     }
 
-    pub(in crate::plugins::demand_control) fn type_field_list_size_directive(
+    pub(in crate::plugins::demand_control) fn output_field_definition(
         &self,
         type_name: &str,
         field_name: &str,
-    ) -> Option<&ListSizeDirective> {
-        self.type_field_list_size_directives
-            .get(type_name)?
-            .get(field_name)
-    }
-
-    pub(in crate::plugins::demand_control) fn type_field_requires_directive(
-        &self,
-        type_name: &str,
-        field_name: &str,
-    ) -> Option<&RequiresDirective> {
-        self.type_field_requires_directives
+    ) -> Option<&FieldDefinition> {
+        self.output_field_definitions
             .get(type_name)?
             .get(field_name)
     }
-
-    pub(in crate::plugins::demand_control) fn argument_cost_directive(
-        &self,
-        definition: &InputValueDefinition,
-        ty: &ExtendedType,
-    ) -> Option<CostDirective> {
-        // For now, we ignore FederationError and return None because this should not block the whole scoring
-        // process at runtime. Later, this should be pushed into the constructor and propagate any federation
-        // errors encountered when parsing.
-        CostSpecDefinition::cost_directive_from_argument(&self.inner, definition, ty)
-            .ok()
-            .flatten()
-    }
 }
 
 impl AsRef<Valid<Schema>> for DemandControlledSchema {