diff --git a/.changesets/fix_tninesling_demand_control_perf.md b/.changesets/fix_tninesling_demand_control_perf.md new file mode 100644 index 0000000000..e892390f87 --- /dev/null +++ b/.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 diff --git a/apollo-router/src/plugins/demand_control/cost_calculator/schema.rs b/apollo-router/src/plugins/demand_control/cost_calculator/schema.rs index d59243f4d5..b5e620609a 100644 --- a/apollo-router/src/plugins/demand_control/cost_calculator/schema.rs +++ b/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, +} + +impl InputDefinition { + fn new( + schema: &ValidFederationSchema, + field_definition: &InputValueDefinition, + ) -> Result { + 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, + list_size_directive: Option, + requires_directive: Option, + arguments: HashMap, +} + +impl FieldDefinition { + fn new( + schema: &ValidFederationSchema, + parent_type_name: &Name, + field_definition: &apollo_compiler::ast::FieldDefinition, + ) -> Result { + 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>, - type_field_list_size_directives: HashMap>, - type_field_requires_directives: HashMap>, + input_field_definitions: HashMap>, + output_field_definitions: HashMap>, } impl DemandControlledSchema { pub(crate) fn new(schema: Arc>) -> Result { let fed_schema = ValidFederationSchema::new((*schema).clone())?; - let mut type_field_cost_directives: HashMap> = + let mut input_field_definitions: HashMap> = HashMap::new(); - let mut type_field_list_size_directives: HashMap> = - HashMap::new(); - let mut type_field_requires_directives: HashMap> = + let mut output_field_definitions: HashMap> = 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 { - // 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> for DemandControlledSchema { 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 d8887e86a9..7881b5abf8 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 @@ -2,7 +2,6 @@ use std::sync::Arc; use ahash::HashMap; use apollo_compiler::ast; -use apollo_compiler::ast::InputValueDefinition; use apollo_compiler::ast::NamedType; use apollo_compiler::executable::ExecutableDocument; use apollo_compiler::executable::Field; @@ -12,12 +11,12 @@ use apollo_compiler::executable::Operation; use apollo_compiler::executable::Selection; use apollo_compiler::executable::SelectionSet; use apollo_compiler::schema::ExtendedType; -use apollo_compiler::Node; use serde_json_bytes::Value; use super::directives::IncludeDirective; use super::directives::SkipDirective; use super::schema::DemandControlledSchema; +use super::schema::InputDefinition; use super::DemandControlError; use crate::graphql::Response; use crate::graphql::ResponseVisitor; @@ -28,6 +27,7 @@ use crate::query_planner::DeferredNode; use crate::query_planner::PlanNode; use crate::query_planner::Primary; use crate::query_planner::QueryPlan; +use crate::spec::TYPENAME; pub(crate) struct StaticCostCalculator { list_size: u32, @@ -44,49 +44,37 @@ struct ScoringContext<'a> { fn score_argument( argument: &apollo_compiler::ast::Value, - argument_definition: &Node, + argument_definition: &InputDefinition, schema: &DemandControlledSchema, variables: &Object, ) -> Result { - let ty = schema - .types - .get(argument_definition.ty.inner_named_type()) - .ok_or_else(|| { - DemandControlError::QueryParseFailure(format!( - "Argument {} was found in query, but its type ({}) was not found in the schema", - argument_definition.name, - argument_definition.ty.inner_named_type() - )) - })?; - let cost_directive = schema.argument_cost_directive(argument_definition, ty); - - match (argument, ty) { + match (argument, argument_definition.ty()) { (_, ExtendedType::Interface(_)) | (_, ExtendedType::Object(_)) | (_, ExtendedType::Union(_)) => Err(DemandControlError::QueryParseFailure( format!( "Argument {} has type {}, but objects, interfaces, and unions are disallowed in this position", - argument_definition.name, - argument_definition.ty.inner_named_type() + argument_definition.name(), + argument_definition.ty().name() ) )), - (ast::Value::Object(inner_args), ExtendedType::InputObject(inner_arg_defs)) => { - let mut cost = cost_directive.map_or(1.0, |cost| cost.weight()); + (ast::Value::Object(inner_args), ExtendedType::InputObject(_)) => { + let mut cost = argument_definition.cost_directive().map_or(1.0, |cost| cost.weight()); for (arg_name, arg_val) in inner_args { - let arg_def = inner_arg_defs.fields.get(arg_name).ok_or_else(|| { + let arg_def = schema.input_field_definition(argument_definition.ty().name(), arg_name).ok_or_else(|| { DemandControlError::QueryParseFailure(format!( "Argument {} was found in query, but its type ({}) was not found in the schema", - argument_definition.name, - argument_definition.ty.inner_named_type() + arg_name, + argument_definition.ty().name() )) })?; - cost += score_argument(arg_val, arg_def, schema, variables,)?; + cost += score_argument(arg_val, arg_def, schema, variables)?; } Ok(cost) } (ast::Value::List(inner_args), _) => { - let mut cost = cost_directive.map_or(0.0, |cost| cost.weight()); + let mut cost = argument_definition.cost_directive().map_or(0.0, |cost| cost.weight()); for arg_val in inner_args { cost += score_argument(arg_val, argument_definition, schema, variables)?; } @@ -102,46 +90,34 @@ fn score_argument( } } (ast::Value::Null, _) => Ok(0.0), - _ => Ok(cost_directive.map_or(0.0, |cost| cost.weight())) + _ => Ok(argument_definition.cost_directive().map_or(0.0, |cost| cost.weight())) } } fn score_variable( variable: &Value, - argument_definition: &Node, + argument_definition: &InputDefinition, schema: &DemandControlledSchema, ) -> Result { - let ty = schema - .types - .get(argument_definition.ty.inner_named_type()) - .ok_or_else(|| { - DemandControlError::QueryParseFailure(format!( - "Argument {} was found in query, but its type ({}) was not found in the schema", - argument_definition.name, - argument_definition.ty.inner_named_type() - )) - })?; - let cost_directive = schema.argument_cost_directive(argument_definition, ty); - - match (variable, ty) { + match (variable, argument_definition.ty()) { (_, ExtendedType::Interface(_)) | (_, ExtendedType::Object(_)) | (_, ExtendedType::Union(_)) => Err(DemandControlError::QueryParseFailure( format!( "Argument {} has type {}, but objects, interfaces, and unions are disallowed in this position", - argument_definition.name, - argument_definition.ty.inner_named_type() + argument_definition.name(), + argument_definition.ty().name() ) )), - (Value::Object(inner_args), ExtendedType::InputObject(inner_arg_defs)) => { - let mut cost = cost_directive.map_or(1.0, |cost| cost.weight()); + (Value::Object(inner_args), ExtendedType::InputObject(_)) => { + let mut cost = argument_definition.cost_directive().map_or(1.0, |cost| cost.weight()); for (arg_name, arg_val) in inner_args { - let arg_def = inner_arg_defs.fields.get(arg_name.as_str()).ok_or_else(|| { + let arg_def = schema.input_field_definition(argument_definition.ty().name(), arg_name.as_str()).ok_or_else(|| { DemandControlError::QueryParseFailure(format!( "Argument {} was found in query, but its type ({}) was not found in the schema", - argument_definition.name, - argument_definition.ty.inner_named_type() + argument_definition.name(), + argument_definition.ty().name() )) })?; cost += score_variable(arg_val, arg_def, schema)?; @@ -149,14 +125,14 @@ fn score_variable( Ok(cost) } (Value::Array(inner_args), _) => { - let mut cost = cost_directive.map_or(0.0, |cost| cost.weight()); + let mut cost = argument_definition.cost_directive().map_or(0.0, |cost| cost.weight()); for arg_val in inner_args { cost += score_variable(arg_val, argument_definition, schema)?; } Ok(cost) } (Value::Null, _) => Ok(0.0), - _ => Ok(cost_directive.map_or(0.0, |cost| cost.weight())) + _ => Ok(argument_definition.cost_directive().map_or(0.0, |cost| cost.weight())) } } @@ -198,24 +174,23 @@ impl StaticCostCalculator { parent_type: &NamedType, list_size_from_upstream: Option, ) -> Result { + if field.name == TYPENAME { + return Ok(0.0); + } if StaticCostCalculator::skipped_by_directives(field) { return Ok(0.0); } - // We need to look up the `FieldDefinition` from the supergraph schema instead of using `field.definition` - // because `field.definition` was generated from the API schema, which strips off the directives we need. - let definition = ctx.schema.type_field(parent_type, &field.name)?; - let ty = field.inner_type_def(ctx.schema).ok_or_else(|| { - DemandControlError::QueryParseFailure(format!( - "Field {} was found in query, but its type is missing from the schema.", - field.name - )) - })?; - - let list_size_directive = match ctx + let definition = ctx .schema - .type_field_list_size_directive(parent_type, &field.name) - { + .output_field_definition(parent_type, &field.name) + .ok_or_else(|| { + DemandControlError::QueryParseFailure(format!( + "Field {} was found in query, but its type is missing from the schema.", + field.name + )) + })?; + let list_size_directive = match definition.list_size_directive() { Some(dir) => ListSizeDirective::new(dir, field, ctx.variables).map(Some), None => Ok(None), }?; @@ -235,12 +210,12 @@ impl StaticCostCalculator { // Determine the cost for this particular field. Scalars are free, non-scalars are not. // For fields with selections, add in the cost of the selections as well. - let mut type_cost = if let Some(cost_directive) = ctx - .schema - .type_field_cost_directive(parent_type, &field.name) - { + let mut type_cost = if let Some(cost_directive) = definition.cost_directive() { cost_directive.weight() - } else if ty.is_interface() || ty.is_object() || ty.is_union() { + } else if definition.ty().is_interface() + || definition.ty().is_object() + || definition.ty().is_union() + { 1.0 } else { 0.0 @@ -274,10 +249,7 @@ impl StaticCostCalculator { // If the field is marked with `@requires`, the required selection may not be included // in the query's selection. Adding that requirement's cost to the field ensures it's // accounted for. - let requirements = ctx - .schema - .type_field_requires_directive(parent_type, &field.name) - .map(|d| &d.fields); + let requirements = definition.requires_directive().map(|d| &d.fields); if let Some(selection_set) = requirements { requirements_cost = self.score_selection_set( ctx, @@ -559,39 +531,78 @@ impl<'schema> ResponseCostCalculator<'schema> { pub(crate) fn new(schema: &'schema DemandControlledSchema) -> Self { Self { cost: 0.0, schema } } -} -impl ResponseVisitor for ResponseCostCalculator<'_> { - fn visit_field( + fn score_response_field( &mut self, request: &ExecutableDocument, variables: &Object, parent_ty: &NamedType, field: &Field, value: &Value, + include_argument_score: bool, ) { - self.visit_list_item(request, variables, parent_ty, field, value); + if let Some(definition) = self.schema.output_field_definition(parent_ty, &field.name) { + match value { + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => { + self.cost += definition + .cost_directive() + .map_or(0.0, |cost| cost.weight()); + } + Value::Array(items) => { + for item in items { + self.visit_list_item(request, variables, parent_ty, field, item); + } + } + Value::Object(children) => { + self.cost += definition + .cost_directive() + .map_or(1.0, |cost| cost.weight()); + self.visit_selections(request, variables, &field.selection_set, children); + } + } - let definition = self.schema.type_field(parent_ty, &field.name); - for argument in &field.arguments { - if let Ok(Some(argument_definition)) = definition - .as_ref() - .map(|def| def.argument_by_name(&argument.name)) - { - if let Ok(score) = - score_argument(&argument.value, argument_definition, self.schema, variables) - { - self.cost += score; + if include_argument_score { + for argument in &field.arguments { + if let Some(argument_definition) = definition.argument_by_name(&argument.name) { + if let Ok(score) = score_argument( + &argument.value, + argument_definition, + self.schema, + variables, + ) { + self.cost += score; + } + } else { + tracing::warn!( + "Failed to get schema definition for argument {}.{}({}:). The resulting response cost will be a partial result.", + parent_ty, + field.name, + argument.name, + ) + } } - } else { - tracing::warn!( - "Failed to get schema definition for argument {} of field {}. The resulting actual cost will be a partial result.", - argument.name, - field.name - ) } + } else { + tracing::warn!( + "Failed to get schema definition for field {}.{}. The resulting response cost will be a partial result.", + parent_ty, + field.name, + ) } } +} + +impl ResponseVisitor for ResponseCostCalculator<'_> { + fn visit_field( + &mut self, + request: &ExecutableDocument, + variables: &Object, + parent_ty: &NamedType, + field: &Field, + value: &Value, + ) { + self.score_response_field(request, variables, parent_ty, field, value, true); + } fn visit_list_item( &mut self, @@ -601,24 +612,7 @@ impl ResponseVisitor for ResponseCostCalculator<'_> { field: &apollo_compiler::executable::Field, value: &Value, ) { - let cost_directive = self - .schema - .type_field_cost_directive(parent_ty, &field.name); - - match value { - Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => { - self.cost += cost_directive.map_or(0.0, |cost| cost.weight()); - } - Value::Array(items) => { - for item in items { - self.visit_list_item(request, variables, parent_ty, field, item); - } - } - Value::Object(children) => { - self.cost += cost_directive.map_or(1.0, |cost| cost.weight()); - self.visit_selections(request, variables, &field.selection_set, children); - } - } + self.score_response_field(request, variables, parent_ty, field, value, false); } }