From a70a7c0b95e792df9c2de54e63774c762766a00d Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Tue, 6 Jan 2026 10:39:39 -0500 Subject: [PATCH 01/17] feat: support configuration of subgraph limits/sizes --- ...nfiguration__tests__schema_generation.snap | 59 ++++++++++++++++++- apollo-router/src/configuration/subgraph.rs | 15 +++++ .../cost_calculator/static_cost.rs | 54 +++++++++++++---- .../src/plugins/demand_control/mod.rs | 12 ++++ .../plugins/demand_control/strategy/mod.rs | 8 ++- .../strategy/static_estimated.rs | 10 ++++ 6 files changed, 146 insertions(+), 12 deletions(-) diff --git a/apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap b/apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap index 63f7c22779..9f5b88331b 100644 --- a/apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap +++ b/apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap @@ -9126,11 +9126,20 @@ expression: "&schema" "description": "The maximum cost of a query", "format": "double", "type": "number" + }, + "subgraphs": { + "allOf": [ + { + "$ref": "#/definitions/SubgraphSubgraphStrategyLimitConfiguration" + } + ], + "description": "Per-subgraph cost control" } }, "required": [ "list_size", - "max" + "max", + "subgraphs" ], "type": "object" } @@ -10397,6 +10406,28 @@ expression: "&schema" }, "type": "object" }, + "SubgraphStrategyLimit": { + "properties": { + "list_size": { + "description": "The assumed length of lists returned by the operation for this subgraph.", + "format": "uint32", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "max": { + "description": "The maximum query cost routed to this subgraph.", + "format": "double", + "type": [ + "number", + "null" + ] + } + }, + "type": "object" + }, "SubgraphSubgraphApqConfiguration": { "description": "Configuration options pertaining to the subgraph server component.", "properties": { @@ -10480,6 +10511,32 @@ expression: "&schema" }, "type": "object" }, + "SubgraphSubgraphStrategyLimitConfiguration": { + "description": "Configuration options pertaining to the subgraph server component.", + "properties": { + "all": { + "allOf": [ + { + "$ref": "#/definitions/SubgraphStrategyLimit" + } + ], + "default": { + "list_size": null, + "max": null + }, + "description": "options applying to all subgraphs" + }, + "subgraphs": { + "additionalProperties": { + "$ref": "#/definitions/SubgraphStrategyLimit" + }, + "default": {}, + "description": "per subgraph options", + "type": "object" + } + }, + "type": "object" + }, "SubgraphTlsClientConfiguration": { "description": "Configuration options pertaining to the subgraph server component.", "properties": { diff --git a/apollo-router/src/configuration/subgraph.rs b/apollo-router/src/configuration/subgraph.rs index a371f4ebe2..c63561eca4 100644 --- a/apollo-router/src/configuration/subgraph.rs +++ b/apollo-router/src/configuration/subgraph.rs @@ -96,6 +96,21 @@ where pub(crate) fn get(&self, subgraph_name: &str) -> &T { self.subgraphs.get(subgraph_name).unwrap_or(&self.all) } + + // Create a new `SubgraphConfiguration` by extracting a value `V` from `T` + pub(crate) fn extract( + &self, + extract_fn: fn(&T) -> V, + ) -> SubgraphConfiguration { + SubgraphConfiguration { + all: extract_fn(&self.all), + subgraphs: self + .subgraphs + .iter() + .map(|(k, v)| (k.clone(), extract_fn(v))) + .collect(), + } + } } impl Debug for SubgraphConfiguration 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 70595ab172..6d8e16c4d6 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 @@ -19,6 +19,7 @@ use super::directives::IncludeDirective; use super::directives::SkipDirective; use super::schema::DemandControlledSchema; use super::schema::InputDefinition; +use crate::configuration::subgraph::SubgraphConfiguration; use crate::graphql::Response; use crate::graphql::ResponseVisitor; use crate::json_ext::Object; @@ -31,6 +32,7 @@ use crate::spec::TYPENAME; pub(crate) struct StaticCostCalculator { list_size: u32, + subgraph_list_sizes: Arc>>, supergraph_schema: Arc, subgraph_schemas: Arc>, } @@ -149,14 +151,22 @@ impl StaticCostCalculator { supergraph_schema: Arc, subgraph_schemas: Arc>, list_size: u32, + subgraph_list_sizes: Arc>>, ) -> Self { Self { list_size, + subgraph_list_sizes, supergraph_schema, subgraph_schemas, } } + fn subgraph_list_size(&self, subgraph_name: &str) -> u32 { + self.subgraph_list_sizes + .get(subgraph_name) + .unwrap_or(self.list_size) + } + /// Scores a field within a GraphQL operation, handling some expected cases where /// directives change how the query is fetched. In the case of the federation /// directive `@requires`, the cost of the required selection is added to the @@ -685,7 +695,12 @@ mod tests { .unwrap_or_default(); let schema = DemandControlledSchema::new(Arc::new(schema.supergraph_schema().clone())).unwrap(); - let calculator = StaticCostCalculator::new(Arc::new(schema), Default::default(), 100); + let calculator = StaticCostCalculator::new( + Arc::new(schema), + Default::default(), + 100, + Default::default(), + ); calculator .estimated( @@ -713,7 +728,12 @@ mod tests { .cloned() .unwrap_or_default(); let schema = DemandControlledSchema::new(Arc::new(schema)).unwrap(); - let calculator = StaticCostCalculator::new(Arc::new(schema), Default::default(), 100); + let calculator = StaticCostCalculator::new( + Arc::new(schema), + Default::default(), + 100, + Default::default(), + ); calculator .estimated(&query, &calculator.supergraph_schema, &variables, true) @@ -767,6 +787,7 @@ mod tests { Arc::new(schema), Arc::new(demand_controlled_subgraph_schemas), 100, + Default::default(), ); calculator.planned(&query_plan, &variables).unwrap() @@ -802,6 +823,7 @@ mod tests { Arc::new(schema), Arc::new(demand_controlled_subgraph_schemas), 100, + Default::default(), ); calculator.rust_planned(&query_plan, &variables).unwrap() @@ -823,9 +845,14 @@ mod tests { let response = Response::from_bytes(Bytes::from(response_bytes)).unwrap(); let schema = DemandControlledSchema::new(Arc::new(schema.supergraph_schema().clone())).unwrap(); - StaticCostCalculator::new(Arc::new(schema), Default::default(), 100) - .actual(&query.executable, &response, &variables) - .unwrap() + StaticCostCalculator::new( + Arc::new(schema), + Default::default(), + 100, + Default::default(), + ) + .actual(&query.executable, &response, &variables) + .unwrap() } /// Actual cost of an operation on a plain, non-federated schema. @@ -851,9 +878,14 @@ mod tests { let response = Response::from_bytes(Bytes::from(response_bytes)).unwrap(); let schema = DemandControlledSchema::new(Arc::new(schema)).unwrap(); - StaticCostCalculator::new(Arc::new(schema), Default::default(), 100) - .actual(&query, &response, &variables) - .unwrap() + StaticCostCalculator::new( + Arc::new(schema), + Default::default(), + 100, + Default::default(), + ) + .actual(&query, &response, &variables) + .unwrap() } #[test] @@ -1050,7 +1082,8 @@ mod tests { DemandControlledSchema::new(Arc::new(schema.supergraph_schema().clone())).unwrap(), ); - let calculator = StaticCostCalculator::new(schema.clone(), Default::default(), 100); + let calculator = + StaticCostCalculator::new(schema.clone(), Default::default(), 100, Default::default()); let conservative_estimate = calculator .estimated( &query.executable, @@ -1060,7 +1093,8 @@ mod tests { ) .unwrap(); - let calculator = StaticCostCalculator::new(schema.clone(), Default::default(), 5); + let calculator = + StaticCostCalculator::new(schema.clone(), Default::default(), 5, Default::default()); let narrow_estimate = calculator .estimated( &query.executable, diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index 4314e8cd1f..4efee85b67 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -27,6 +27,7 @@ use tower::ServiceBuilder; use tower::ServiceExt; use crate::Context; +use crate::configuration::subgraph::SubgraphConfiguration; use crate::error::Error; use crate::graphql; use crate::graphql::IntoGraphQLErrors; @@ -73,6 +74,8 @@ pub(crate) enum StrategyConfig { list_size: u32, /// The maximum cost of a query max: f64, + /// Per-subgraph cost control + subgraphs: SubgraphConfiguration, }, #[cfg(test)] @@ -82,6 +85,15 @@ pub(crate) enum StrategyConfig { }, } +#[derive(Clone, Default, Debug, Serialize, Deserialize, JsonSchema)] +pub(crate) struct SubgraphStrategyLimit { + /// The assumed length of lists returned by the operation for this subgraph. + list_size: Option, + + /// The maximum query cost routed to this subgraph. + max: Option, +} + #[derive(Copy, Clone, Debug, Serialize, Deserialize, JsonSchema, Eq, PartialEq)] #[serde(deny_unknown_fields, rename_all = "snake_case")] pub(crate) enum Mode { diff --git a/apollo-router/src/plugins/demand_control/strategy/mod.rs b/apollo-router/src/plugins/demand_control/strategy/mod.rs index 738b935b12..afb5f72437 100644 --- a/apollo-router/src/plugins/demand_control/strategy/mod.rs +++ b/apollo-router/src/plugins/demand_control/strategy/mod.rs @@ -93,12 +93,18 @@ impl StrategyFactory { pub(crate) fn create(&self) -> Strategy { let strategy: Arc = match &self.config.strategy { - StrategyConfig::StaticEstimated { list_size, max } => Arc::new(StaticEstimated { + StrategyConfig::StaticEstimated { + list_size, + max, + subgraphs, + } => Arc::new(StaticEstimated { max: *max, + subgraph_maxes: Arc::new(subgraphs.extract(|strategy| strategy.max)), cost_calculator: StaticCostCalculator::new( self.supergraph_schema.clone(), self.subgraph_schemas.clone(), *list_size, + Arc::new(subgraphs.extract(|strategy| strategy.list_size)), ), }), #[cfg(test)] diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index ea442b3a69..eb0867c368 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -1,5 +1,8 @@ +use std::sync::Arc; + use apollo_compiler::ExecutableDocument; +use crate::configuration::subgraph::SubgraphConfiguration; use crate::graphql; use crate::plugins::demand_control::DemandControlError; use crate::plugins::demand_control::cost_calculator::static_cost::StaticCostCalculator; @@ -11,9 +14,16 @@ use crate::services::subgraph; pub(crate) struct StaticEstimated { // The estimated value of the demand pub(crate) max: f64, + pub(crate) subgraph_maxes: Arc>>, pub(crate) cost_calculator: StaticCostCalculator, } +impl StaticEstimated { + fn subgraph_max(&self, subgraph_name: &str) -> Option { + *self.subgraph_maxes.get(subgraph_name) + } +} + impl StrategyImpl for StaticEstimated { fn on_execution_request(&self, request: &execution::Request) -> Result<(), DemandControlError> { self.cost_calculator From 9af693df33e5b0ed26084c99f73730f673b653bb Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Tue, 6 Jan 2026 11:09:57 -0500 Subject: [PATCH 02/17] feat: check that per-subgraph max is non-negative --- .../src/plugins/demand_control/mod.rs | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index 4efee85b67..e11928a595 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -1,6 +1,8 @@ //! Demand control plugin. //! This plugin will use the cost calculation algorithm to determine if a query should be allowed to execute. //! On the request path it will use estimated + +use std::collections::HashSet; use std::future; use std::ops::ControlFlow; use std::sync::Arc; @@ -94,6 +96,41 @@ pub(crate) struct SubgraphStrategyLimit { max: Option, } +impl StrategyConfig { + fn validate(&self, subgraph_names: HashSet<&String>) -> Result<(), BoxError> { + #[derive(thiserror::Error, Debug)] + enum Error { + #[error("Maximum per-subgraph query cost for `{0}` is negative")] + NegativeQueryCost(String), + } + + #[allow(irrefutable_let_patterns)] + // need to destructure StrategyConfig and ignore StrategyConfig::Test + let StrategyConfig::StaticEstimated { subgraphs, .. } = self else { + return Ok(()); + }; + + if subgraphs.all.max.is_some_and(|s| s < 0.0) { + return Err(Error::NegativeQueryCost("all".to_string()).into()); + } + + for (subgraph_name, subgraph_config) in subgraphs.subgraphs.iter() { + if !subgraph_names.contains(subgraph_name) { + tracing::warn!( + "Subgraph `{subgraph_name}` missing from schema but was specified in per-subgraph demand cost; it will be ignored" + ); + continue; + } + + if subgraph_config.max.is_some_and(|s| s < 0.0) { + return Err(Error::NegativeQueryCost(subgraph_name.to_string()).into()); + } + } + + Ok(()) + } +} + #[derive(Copy, Clone, Debug, Serialize, Deserialize, JsonSchema, Eq, PartialEq)] #[serde(deny_unknown_fields, rename_all = "snake_case")] pub(crate) enum Mode { @@ -360,6 +397,10 @@ impl Plugin for DemandControl { .insert(subgraph_name.clone(), demand_controlled_subgraph_schema); } + // validate that per-subgraph maxes are all non-negative + let subgraph_names = init.subgraph_schemas.keys().collect(); + init.config.strategy.validate(subgraph_names)?; + Ok(DemandControl { strategy_factory: StrategyFactory::new( init.config.clone(), From 5cf25b90a07888bf1eca858d5c093ebbc3d3b9aa Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Tue, 6 Jan 2026 15:24:42 -0500 Subject: [PATCH 03/17] test: add tests for configuration --- .../fixtures/invalid_per_subgraph.yaml | 13 ++++ .../fixtures/per_subgraph_inheritance.yaml | 16 +++++ .../fixtures/per_subgraph_no_inheritance.yaml | 15 ++++ .../plugins/demand_control/strategy/mod.rs | 36 +++++++--- .../strategy/static_estimated.rs | 72 +++++++++++++++++++ 5 files changed, 142 insertions(+), 10 deletions(-) create mode 100644 apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml create mode 100644 apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml create mode 100644 apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml diff --git a/apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml b/apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml new file mode 100644 index 0000000000..8ddabcb984 --- /dev/null +++ b/apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml @@ -0,0 +1,13 @@ +demand_control: + enabled: true + mode: enforce + strategy: + static_estimated: + list_size: 1 + max: 10 + subgraphs: + all: + list_size: 3 + subgraphs: + products: + max: -1 diff --git a/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml new file mode 100644 index 0000000000..57d4acc553 --- /dev/null +++ b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml @@ -0,0 +1,16 @@ +demand_control: + enabled: true + mode: enforce + strategy: + static_estimated: + list_size: 1 + max: 10 + subgraphs: + all: + list_size: 3 + max: 5 + subgraphs: + products: + list_size: 5 + reviews: + max: 2 diff --git a/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml new file mode 100644 index 0000000000..bebb4fbf7c --- /dev/null +++ b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml @@ -0,0 +1,15 @@ +demand_control: + enabled: true + mode: enforce + strategy: + static_estimated: + list_size: 1 + max: 10 + subgraphs: + all: + list_size: 3 + subgraphs: + products: + list_size: 5 + reviews: + max: 2 diff --git a/apollo-router/src/plugins/demand_control/strategy/mod.rs b/apollo-router/src/plugins/demand_control/strategy/mod.rs index afb5f72437..37ee805681 100644 --- a/apollo-router/src/plugins/demand_control/strategy/mod.rs +++ b/apollo-router/src/plugins/demand_control/strategy/mod.rs @@ -4,11 +4,13 @@ use ahash::HashMap; use apollo_compiler::ExecutableDocument; use crate::Context; +use crate::configuration::subgraph::SubgraphConfiguration; use crate::graphql; use crate::plugins::demand_control::DemandControlConfig; use crate::plugins::demand_control::DemandControlError; use crate::plugins::demand_control::Mode; use crate::plugins::demand_control::StrategyConfig; +use crate::plugins::demand_control::SubgraphStrategyLimit; use crate::plugins::demand_control::cost_calculator::schema::DemandControlledSchema; use crate::plugins::demand_control::cost_calculator::static_cost::StaticCostCalculator; use crate::plugins::demand_control::strategy::static_estimated::StaticEstimated; @@ -91,22 +93,36 @@ impl StrategyFactory { } } + // Function extracted for use in tests - allows us to build a `StaticEstimated` directly rather + // than a `impl StrategyImpl` + fn create_static_estimated_strategy( + &self, + list_size: u32, + max: f64, + subgraphs: &SubgraphConfiguration, + ) -> StaticEstimated { + let subgraph_list_sizes = Arc::new(subgraphs.extract(|strategy| strategy.list_size)); + let subgraph_maxes = Arc::new(subgraphs.extract(|strategy| strategy.max)); + let cost_calculator = StaticCostCalculator::new( + self.supergraph_schema.clone(), + self.subgraph_schemas.clone(), + list_size, + subgraph_list_sizes, + ); + StaticEstimated { + max, + subgraph_maxes, + cost_calculator, + } + } + pub(crate) fn create(&self) -> Strategy { let strategy: Arc = match &self.config.strategy { StrategyConfig::StaticEstimated { list_size, max, subgraphs, - } => Arc::new(StaticEstimated { - max: *max, - subgraph_maxes: Arc::new(subgraphs.extract(|strategy| strategy.max)), - cost_calculator: StaticCostCalculator::new( - self.supergraph_schema.clone(), - self.subgraph_schemas.clone(), - *list_size, - Arc::new(subgraphs.extract(|strategy| strategy.list_size)), - ), - }), + } => Arc::new(self.create_static_estimated_strategy(*list_size, *max, subgraphs)), #[cfg(test)] StrategyConfig::Test { stage, error } => Arc::new(test::Test { stage: stage.clone(), diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index eb0867c368..0edefe66af 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -85,3 +85,75 @@ impl StrategyImpl for StaticEstimated { Ok(()) } } + +#[cfg(test)] +mod tests { + use tower::BoxError; + + use super::StaticEstimated; + use crate::plugins::demand_control::DemandControl; + use crate::plugins::demand_control::StrategyConfig; + use crate::plugins::test::PluginTestHarness; + + async fn load_config_and_extract_strategy( + config: &'static str, + ) -> Result { + let schema_str = + include_str!("../cost_calculator/fixtures/basic_supergraph_schema.graphql"); + let plugin = PluginTestHarness::::builder() + .config(config) + .schema(schema_str) + .build() + .await?; + + let StrategyConfig::StaticEstimated { + list_size, + max, + ref subgraphs, + } = plugin.config.strategy + else { + panic!("must provide static_estimated config"); + }; + let strategy = plugin + .strategy_factory + .create_static_estimated_strategy(list_size, max, &subgraphs); + Ok(strategy) + } + + #[tokio::test] + async fn test_per_subgraph_configuration_inheritance() { + let config = include_str!("../fixtures/per_subgraph_inheritance.yaml"); + + let strategy = load_config_and_extract_strategy(config).await.unwrap(); + assert_eq!(strategy.subgraph_max("reviews").unwrap(), 2.0); + assert_eq!(strategy.subgraph_max("products").unwrap(), 5.0); + assert_eq!(strategy.subgraph_max("users").unwrap(), 5.0); + } + + #[tokio::test] + async fn test_per_subgraph_configuration_no_inheritance() { + let config = include_str!("../fixtures/per_subgraph_no_inheritance.yaml"); + + let strategy = load_config_and_extract_strategy(config).await.unwrap(); + assert_eq!(strategy.subgraph_max("reviews").unwrap(), 2.0); + assert!(strategy.subgraph_max("products").is_none()); + assert!(strategy.subgraph_max("users").is_none()); + } + + #[tokio::test] + async fn test_invalid_per_subgraph_configuration() { + let config = include_str!("../fixtures/invalid_per_subgraph.yaml"); + let strategy_result = load_config_and_extract_strategy(config).await; + + match strategy_result { + Ok(strategy) => { + eprintln!("{:?}", strategy.subgraph_maxes); + panic!("Expected error") + } + Err(err) => assert_eq!( + &err.to_string(), + "Maximum per-subgraph query cost for `products` is negative" + ), + }; + } +} From 97d6bb3717030656a0a77136de11fcf9966ea691 Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Wed, 7 Jan 2026 10:25:36 -0500 Subject: [PATCH 04/17] feat: calculate cost by subgraph --- .../cost_calculator/static_cost.rs | 251 +++++++++++++++--- .../src/plugins/demand_control/mod.rs | 2 +- .../strategy/static_estimated.rs | 3 +- 3 files changed, 213 insertions(+), 43 deletions(-) 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 6d8e16c4d6..b2b4d4d4e5 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 @@ -1,3 +1,4 @@ +use std::ops::AddAssign; use std::sync::Arc; use ahash::HashMap; @@ -30,6 +31,48 @@ use crate::query_planner::Primary; use crate::query_planner::QueryPlan; use crate::spec::TYPENAME; +#[derive(Default, Debug)] +pub(crate) struct CostBySubgraph(HashMap); +impl CostBySubgraph { + fn new(subgraph: String, value: f64) -> Self { + let mut cost = Self::default(); + cost.insert(subgraph, value); + cost + } + + fn get(&self, subgraph: &str) -> Option { + self.0.get(subgraph).copied() + } + + fn insert(&mut self, subgraph: String, value: f64) { + self.0.insert(subgraph, value); + } + + pub(crate) fn total(&self) -> f64 { + self.0.values().sum() + } + + fn max_by_subgraph(mut self, other: Self) -> Self { + for (subgraph, value) in other.0 { + if let Some(existing_value) = self.get(&subgraph) { + self.insert(subgraph, existing_value.max(value)); + } else { + self.insert(subgraph, value); + } + } + self + } +} + +impl AddAssign for CostBySubgraph { + fn add_assign(&mut self, rhs: Self) { + for (subgraph, value) in rhs.0 { + let entry = self.0.entry(subgraph).or_default(); + *entry += value; + } + } +} + pub(crate) struct StaticCostCalculator { list_size: u32, subgraph_list_sizes: Arc>>, @@ -191,6 +234,7 @@ impl StaticCostCalculator { field: &Field, parent_type: &NamedType, list_size_from_upstream: Option, + subgraph: Option<&str>, ) -> Result { // When we pre-process the schema, __typename isn't included. So, we short-circuit here to avoid failed lookups. if field.name == TYPENAME { @@ -223,6 +267,8 @@ impl StaticCostCalculator { .and_then(|dir| dir.expected_size) { expected_size + } else if let Some(subgraph) = subgraph { + self.subgraph_list_size(subgraph) as i32 } else { self.list_size as i32 }; @@ -244,6 +290,7 @@ impl StaticCostCalculator { &field.selection_set, field.ty().inner_named_type(), list_size_directive.as_ref(), + subgraph, )?; let mut arguments_cost = 0.0; @@ -275,6 +322,7 @@ impl StaticCostCalculator { selection_set, parent_type, list_size_directive.as_ref(), + subgraph, )?; } } @@ -298,6 +346,7 @@ impl StaticCostCalculator { ctx: &ScoringContext, fragment_spread: &FragmentSpread, list_size_directive: Option<&ListSizeDirective>, + subgraph: Option<&str>, ) -> Result { let fragment = fragment_spread.fragment_def(ctx.query).ok_or_else(|| { DemandControlError::QueryParseFailure(format!( @@ -310,6 +359,7 @@ impl StaticCostCalculator { &fragment.selection_set, fragment.type_condition(), list_size_directive, + subgraph, ) } @@ -319,6 +369,7 @@ impl StaticCostCalculator { inline_fragment: &InlineFragment, parent_type: &NamedType, list_size_directive: Option<&ListSizeDirective>, + subgraph: Option<&str>, ) -> Result { self.score_selection_set( ctx, @@ -328,6 +379,7 @@ impl StaticCostCalculator { .as_ref() .unwrap_or(parent_type), list_size_directive, + subgraph, ) } @@ -335,6 +387,7 @@ impl StaticCostCalculator { &self, operation: &Operation, ctx: &ScoringContext, + subgraph: Option<&str>, ) -> Result { let mut cost = if operation.is_mutation() { 10.0 } else { 0.0 }; @@ -345,7 +398,13 @@ impl StaticCostCalculator { ))); }; - cost += self.score_selection_set(ctx, &operation.selection_set, root_type_name, None)?; + cost += self.score_selection_set( + ctx, + &operation.selection_set, + root_type_name, + None, + subgraph, + )?; Ok(cost) } @@ -356,6 +415,7 @@ impl StaticCostCalculator { selection: &Selection, parent_type: &NamedType, list_size_directive: Option<&ListSizeDirective>, + subgraph: Option<&str>, ) -> Result { match selection { Selection::Field(f) => self.score_field( @@ -363,10 +423,13 @@ impl StaticCostCalculator { f, parent_type, list_size_directive.and_then(|dir| dir.size_of(f)), + subgraph, ), - Selection::FragmentSpread(s) => self.score_fragment_spread(ctx, s, list_size_directive), + Selection::FragmentSpread(s) => { + self.score_fragment_spread(ctx, s, list_size_directive, subgraph) + } Selection::InlineFragment(i) => { - self.score_inline_fragment(ctx, i, parent_type, list_size_directive) + self.score_inline_fragment(ctx, i, parent_type, list_size_directive, subgraph) } } } @@ -377,10 +440,17 @@ impl StaticCostCalculator { selection_set: &SelectionSet, parent_type_name: &NamedType, list_size_directive: Option<&ListSizeDirective>, + subgraph: Option<&str>, ) -> Result { let mut cost = 0.0; for selection in selection_set.selections.iter() { - cost += self.score_selection(ctx, selection, parent_type_name, list_size_directive)?; + cost += self.score_selection( + ctx, + selection, + parent_type_name, + list_size_directive, + subgraph, + )?; } Ok(cost) } @@ -403,7 +473,7 @@ impl StaticCostCalculator { &self, plan_node: &PlanNode, variables: &Object, - ) -> Result { + ) -> Result { match plan_node { PlanNode::Sequence { nodes } => self.summed_score_of_nodes(nodes, variables), PlanNode::Parallel { nodes } => self.summed_score_of_nodes(nodes, variables), @@ -434,7 +504,7 @@ impl StaticCostCalculator { subgraph: &str, operation: &SerializableDocument, variables: &Object, - ) -> Result { + ) -> Result { tracing::debug!("On subgraph {}, scoring operation: {}", subgraph, operation); let schema = self.subgraph_schemas.get(subgraph).ok_or_else(|| { @@ -446,7 +516,8 @@ impl StaticCostCalculator { let operation = operation .as_parsed() .map_err(DemandControlError::SubgraphOperationNotInitialized)?; - self.estimated(operation, schema, variables, false) + let estimated_cost = self.estimated(operation, schema, variables, false, Some(subgraph))?; + Ok(CostBySubgraph::new(subgraph.to_string(), estimated_cost)) } fn max_score_of_nodes( @@ -454,15 +525,15 @@ impl StaticCostCalculator { left: &Option>, right: &Option>, variables: &Object, - ) -> Result { + ) -> Result { match (left, right) { - (None, None) => Ok(0.0), + (None, None) => Ok(CostBySubgraph::default()), (None, Some(right)) => self.score_plan_node(right, variables), (Some(left), None) => self.score_plan_node(left, variables), (Some(left), Some(right)) => { let left_score = self.score_plan_node(left, variables)?; let right_score = self.score_plan_node(right, variables)?; - Ok(left_score.max(right_score)) + Ok(left_score.max_by_subgraph(right_score)) } } } @@ -472,8 +543,8 @@ impl StaticCostCalculator { primary: &Primary, deferred: &Vec, variables: &Object, - ) -> Result { - let mut score = 0.0; + ) -> Result { + let mut score = CostBySubgraph::default(); if let Some(node) = &primary.node { score += self.score_plan_node(node, variables)?; } @@ -489,8 +560,8 @@ impl StaticCostCalculator { &self, nodes: &Vec, variables: &Object, - ) -> Result { - let mut sum = 0.0; + ) -> Result { + let mut sum = CostBySubgraph::default(); for node in nodes { sum += self.score_plan_node(node, variables)?; } @@ -503,6 +574,7 @@ impl StaticCostCalculator { schema: &DemandControlledSchema, variables: &Object, should_estimate_requires: bool, + subgraph: Option<&str>, ) -> Result { let mut cost = 0.0; let ctx = ScoringContext { @@ -512,10 +584,10 @@ impl StaticCostCalculator { should_estimate_requires, }; if let Some(op) = &query.operations.anonymous { - cost += self.score_operation(op, &ctx)?; + cost += self.score_operation(op, &ctx, subgraph)?; } for (_name, op) in query.operations.named.iter() { - cost += self.score_operation(op, &ctx)?; + cost += self.score_operation(op, &ctx, subgraph)?; } Ok(cost) } @@ -524,7 +596,7 @@ impl StaticCostCalculator { &self, query_plan: &QueryPlan, variables: &Object, - ) -> Result { + ) -> Result { self.score_plan_node(&query_plan.root, variables) } @@ -668,7 +740,7 @@ mod tests { &self, query_plan: &apollo_federation::query_plan::QueryPlan, variables: &Object, - ) -> Result { + ) -> Result { let js_planner_node: PlanNode = query_plan.node.as_ref().unwrap().into(); self.score_plan_node(&js_planner_node, variables) } @@ -708,6 +780,7 @@ mod tests { &calculator.supergraph_schema, &variables, true, + None, ) .unwrap() } @@ -736,11 +809,21 @@ mod tests { ); calculator - .estimated(&query, &calculator.supergraph_schema, &variables, true) + .estimated( + &query, + &calculator.supergraph_schema, + &variables, + true, + None, + ) .unwrap() } - async fn planned_cost_js(schema_str: &str, query_str: &str, variables_str: &str) -> f64 { + async fn planned_cost_js( + schema_str: &str, + query_str: &str, + variables_str: &str, + ) -> CostBySubgraph { let config: Arc = Arc::new(Default::default()); let (schema, query) = parse_schema_and_operation(schema_str, query_str, &config); let variables = serde_json::from_str::(variables_str) @@ -793,7 +876,7 @@ mod tests { calculator.planned(&query_plan, &variables).unwrap() } - fn planned_cost_rust(schema_str: &str, query_str: &str, variables_str: &str) -> f64 { + fn planned_cost_rust(schema_str: &str, query_str: &str, variables_str: &str) -> CostBySubgraph { let config: Arc = Arc::new(Default::default()); let (schema, query) = parse_schema_and_operation(schema_str, query_str, &config); let variables = serde_json::from_str::(variables_str) @@ -1006,8 +1089,14 @@ mod tests { let variables = "{}"; assert_eq!(basic_estimated_cost(schema, query, variables), 102.0); - assert_eq!(planned_cost_js(schema, query, variables).await, 102.0); - assert_eq!(planned_cost_rust(schema, query, variables), 102.0); + + let cost_js = planned_cost_js(schema, query, variables).await; + assert_eq!(cost_js.total(), 102.0); + assert_eq!(cost_js.get("products").unwrap(), 102.0); + + let cost_rust = planned_cost_rust(schema, query, variables); + assert_eq!(cost_rust.total(), 102.0); + assert_eq!(cost_rust.get("products").unwrap(), 102.0); } #[test(tokio::test)] @@ -1029,8 +1118,17 @@ mod tests { let response = include_bytes!("./fixtures/federated_ships_required_response.json"); assert_eq!(estimated_cost(schema, query, variables), 10200.0); - assert_eq!(planned_cost_js(schema, query, variables).await, 10400.0); - assert_eq!(planned_cost_rust(schema, query, variables), 10400.0); + + let cost_js = planned_cost_js(schema, query, variables).await; + assert_eq!(cost_js.total(), 10400.0); + assert_eq!(cost_js.get("users").unwrap(), 10100.0); + assert_eq!(cost_js.get("vehicles").unwrap(), 300.0); + + let cost_rust = planned_cost_rust(schema, query, variables); + assert_eq!(cost_rust.total(), 10400.0); + assert_eq!(cost_rust.get("users").unwrap(), 10100.0); + assert_eq!(cost_rust.get("vehicles").unwrap(), 300.0); + assert_eq!(actual_cost(schema, query, variables, response), 2.0); } @@ -1042,8 +1140,17 @@ mod tests { let response = include_bytes!("./fixtures/federated_ships_fragment_response.json"); assert_eq!(estimated_cost(schema, query, variables), 300.0); - assert_eq!(planned_cost_js(schema, query, variables).await, 400.0); - assert_eq!(planned_cost_rust(schema, query, variables), 400.0); + + let cost_js = planned_cost_js(schema, query, variables).await; + assert_eq!(cost_js.total(), 400.0); + assert_eq!(cost_js.get("users").unwrap(), 200.0); + assert_eq!(cost_js.get("vehicles").unwrap(), 200.0); + + let cost_rust = planned_cost_rust(schema, query, variables); + assert_eq!(cost_rust.total(), 400.0); + assert_eq!(cost_rust.get("users").unwrap(), 200.0); + assert_eq!(cost_rust.get("vehicles").unwrap(), 200.0); + assert_eq!(actual_cost(schema, query, variables, response), 6.0); } @@ -1055,8 +1162,17 @@ mod tests { let response = include_bytes!("./fixtures/federated_ships_fragment_response.json"); assert_eq!(estimated_cost(schema, query, variables), 300.0); - assert_eq!(planned_cost_js(schema, query, variables).await, 400.0); - assert_eq!(planned_cost_rust(schema, query, variables), 400.0); + + let cost_js = planned_cost_js(schema, query, variables).await; + assert_eq!(cost_js.total(), 400.0); + assert_eq!(cost_js.get("users").unwrap(), 200.0); + assert_eq!(cost_js.get("vehicles").unwrap(), 200.0); + + let cost_rust = planned_cost_rust(schema, query, variables); + assert_eq!(cost_rust.total(), 400.0); + assert_eq!(cost_rust.get("users").unwrap(), 200.0); + assert_eq!(cost_rust.get("vehicles").unwrap(), 200.0); + assert_eq!(actual_cost(schema, query, variables, response), 6.0); } @@ -1068,8 +1184,17 @@ mod tests { let response = include_bytes!("./fixtures/federated_ships_deferred_response.json"); assert_eq!(estimated_cost(schema, query, variables), 10200.0); - assert_eq!(planned_cost_js(schema, query, variables).await, 10400.0); - assert_eq!(planned_cost_rust(schema, query, variables), 10400.0); + + let cost_js = planned_cost_js(schema, query, variables).await; + assert_eq!(cost_js.total(), 10400.0); + assert_eq!(cost_js.get("users").unwrap(), 10100.0); + assert_eq!(cost_js.get("vehicles").unwrap(), 300.0); + + let cost_rust = planned_cost_rust(schema, query, variables); + assert_eq!(cost_rust.total(), 10400.0); + assert_eq!(cost_rust.get("users").unwrap(), 10100.0); + assert_eq!(cost_rust.get("vehicles").unwrap(), 300.0); + assert_eq!(actual_cost(schema, query, variables, response), 2.0); } @@ -1090,6 +1215,7 @@ mod tests { &calculator.supergraph_schema, &Default::default(), true, + None, ) .unwrap(); @@ -1101,6 +1227,7 @@ mod tests { &calculator.supergraph_schema, &Default::default(), true, + None, ) .unwrap(); @@ -1132,8 +1259,17 @@ mod tests { let response = include_bytes!("./fixtures/custom_cost_response.json"); assert_eq!(estimated_cost(schema, query, variables), 127.0); - assert_eq!(planned_cost_js(schema, query, variables).await, 127.0); - assert_eq!(planned_cost_rust(schema, query, variables), 127.0); + + let cost_js = planned_cost_js(schema, query, variables).await; + assert_eq!(cost_js.total(), 127.0); + assert_eq!(cost_js.get("subgraphWithListSize").unwrap(), 6.0); + assert_eq!(cost_js.get("subgraphWithCost").unwrap(), 121.0); + + let cost_rust = planned_cost_rust(schema, query, variables); + assert_eq!(cost_rust.total(), 127.0); + assert_eq!(cost_rust.get("subgraphWithListSize").unwrap(), 6.0); + assert_eq!(cost_rust.get("subgraphWithCost").unwrap(), 121.0); + assert_eq!(actual_cost(schema, query, variables, response), 125.0); } @@ -1145,8 +1281,17 @@ mod tests { let response = include_bytes!("./fixtures/custom_cost_response.json"); assert_eq!(estimated_cost(schema, query, variables), 127.0); - assert_eq!(planned_cost_js(schema, query, variables).await, 127.0); - assert_eq!(planned_cost_rust(schema, query, variables), 127.0); + + let cost_js = planned_cost_js(schema, query, variables).await; + assert_eq!(cost_js.total(), 127.0); + assert_eq!(cost_js.get("subgraphWithListSize").unwrap(), 6.0); + assert_eq!(cost_js.get("subgraphWithCost").unwrap(), 121.0); + + let cost_rust = planned_cost_rust(schema, query, variables); + assert_eq!(cost_rust.total(), 127.0); + assert_eq!(cost_rust.get("subgraphWithListSize").unwrap(), 6.0); + assert_eq!(cost_rust.get("subgraphWithCost").unwrap(), 121.0); + assert_eq!(actual_cost(schema, query, variables, response), 125.0); } @@ -1159,8 +1304,17 @@ mod tests { let response = include_bytes!("./fixtures/custom_cost_response.json"); assert_eq!(estimated_cost(schema, query, variables), 132.0); - assert_eq!(planned_cost_js(schema, query, variables).await, 132.0); - assert_eq!(planned_cost_rust(schema, query, variables), 132.0); + + let cost_js = planned_cost_js(schema, query, variables).await; + assert_eq!(cost_js.total(), 132.0); + assert_eq!(cost_js.get("subgraphWithListSize").unwrap(), 11.0); + assert_eq!(cost_js.get("subgraphWithCost").unwrap(), 121.0); + + let cost_rust = planned_cost_rust(schema, query, variables); + assert_eq!(cost_rust.total(), 132.0); + assert_eq!(cost_rust.get("subgraphWithListSize").unwrap(), 11.0); + assert_eq!(cost_rust.get("subgraphWithCost").unwrap(), 121.0); + assert_eq!(actual_cost(schema, query, variables, response), 125.0); } @@ -1173,8 +1327,17 @@ mod tests { let response = include_bytes!("./fixtures/custom_cost_response.json"); assert_eq!(estimated_cost(schema, query, variables), 127.0); - assert_eq!(planned_cost_js(schema, query, variables).await, 127.0); - assert_eq!(planned_cost_rust(schema, query, variables), 127.0); + + let cost_js = planned_cost_js(schema, query, variables).await; + assert_eq!(cost_js.total(), 127.0); + assert_eq!(cost_js.get("subgraphWithListSize").unwrap(), 6.0); + assert_eq!(cost_js.get("subgraphWithCost").unwrap(), 121.0); + + let cost_rust = planned_cost_rust(schema, query, variables); + assert_eq!(cost_rust.total(), 127.0); + assert_eq!(cost_rust.get("subgraphWithListSize").unwrap(), 6.0); + assert_eq!(cost_rust.get("subgraphWithCost").unwrap(), 121.0); + assert_eq!(actual_cost(schema, query, variables, response), 125.0); } @@ -1206,7 +1369,13 @@ mod tests { let variables = "{}"; assert_eq!(estimated_cost(schema, query, variables), 1.0); - assert_eq!(planned_cost_js(schema, query, variables).await, 1.0); - assert_eq!(planned_cost_rust(schema, query, variables), 1.0); + + let cost_js = planned_cost_js(schema, query, variables).await; + assert_eq!(cost_js.total(), 1.0); + assert_eq!(cost_js.get("subgraph").unwrap(), 1.0); + + let cost_rust = planned_cost_rust(schema, query, variables); + assert_eq!(cost_rust.total(), 1.0); + assert_eq!(cost_rust.get("subgraph").unwrap(), 1.0); } } diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index e11928a595..8151651de0 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -76,7 +76,7 @@ pub(crate) enum StrategyConfig { list_size: u32, /// The maximum cost of a query max: f64, - /// Per-subgraph cost control + /// Cost control by subgraph subgraphs: SubgraphConfiguration, }, diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index 0edefe66af..698db31e1a 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -31,7 +31,8 @@ impl StrategyImpl for StaticEstimated { &request.query_plan, &request.supergraph_request.body().variables, ) - .and_then(|cost| { + .and_then(|cost_by_subgraph| { + let cost = cost_by_subgraph.total(); request .context .insert_cost_strategy("static_estimated".to_string())?; From 406647e3306d40fc1d45b8b1a70589987e7e05b8 Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Wed, 7 Jan 2026 10:41:34 -0500 Subject: [PATCH 05/17] feat: add cost by subgraph to context --- .../demand_control/cost_calculator/static_cost.rs | 4 +++- apollo-router/src/plugins/demand_control/mod.rs | 13 +++++++++++++ .../demand_control/strategy/static_estimated.rs | 3 +++ 3 files changed, 19 insertions(+), 1 deletion(-) 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 b2b4d4d4e5..b8c2f967af 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 @@ -13,6 +13,8 @@ use apollo_compiler::executable::Selection; use apollo_compiler::executable::SelectionSet; use apollo_compiler::schema::ExtendedType; use apollo_federation::query_plan::serializable_document::SerializableDocument; +use serde::Deserialize; +use serde::Serialize; use serde_json_bytes::Value; use super::DemandControlError; @@ -31,7 +33,7 @@ use crate::query_planner::Primary; use crate::query_planner::QueryPlan; use crate::spec::TYPENAME; -#[derive(Default, Debug)] +#[derive(Default, Debug, Serialize, Deserialize)] pub(crate) struct CostBySubgraph(HashMap); impl CostBySubgraph { fn new(subgraph: String, value: f64) -> Self { diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index 8151651de0..bd128321cf 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -38,6 +38,7 @@ use crate::layers::ServiceBuilderExt; use crate::plugin::Plugin; use crate::plugin::PluginInit; use crate::plugins::demand_control::cost_calculator::schema::DemandControlledSchema; +use crate::plugins::demand_control::cost_calculator::static_cost::CostBySubgraph; use crate::plugins::demand_control::strategy::Strategy; use crate::plugins::demand_control::strategy::StrategyFactory; use crate::plugins::telemetry::tracing::apollo_telemetry::emit_error_event; @@ -54,6 +55,9 @@ pub(crate) const COST_ACTUAL_KEY: &str = "apollo::demand_control::actual_cost"; pub(crate) const COST_RESULT_KEY: &str = "apollo::demand_control::result"; pub(crate) const COST_STRATEGY_KEY: &str = "apollo::demand_control::strategy"; +pub(crate) const COST_BY_SUBGRAPH_ESTIMATED_KEY: &str = + "apollo::demand_control::estimated_cost_by_subgraph"; + /// Algorithm for calculating the cost of an incoming query. #[derive(Clone, Debug, Deserialize, JsonSchema)] #[serde(deny_unknown_fields, rename_all = "snake_case")] @@ -300,6 +304,15 @@ impl Context { .map_err(|e| DemandControlError::ContextSerializationError(e.to_string())) } + pub(crate) fn insert_estimated_cost_by_subgraph( + &self, + cost: CostBySubgraph, + ) -> Result<(), DemandControlError> { + self.insert(COST_BY_SUBGRAPH_ESTIMATED_KEY, cost) + .map_err(|e| DemandControlError::ContextSerializationError(e.to_string()))?; + Ok(()) + } + pub(crate) fn insert_actual_cost(&self, cost: f64) -> Result<(), DemandControlError> { self.insert(COST_ACTUAL_KEY, cost) .map_err(|e| DemandControlError::ContextSerializationError(e.to_string()))?; diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index 698db31e1a..539029d43d 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -38,6 +38,9 @@ impl StrategyImpl for StaticEstimated { .insert_cost_strategy("static_estimated".to_string())?; request.context.insert_cost_result("COST_OK".to_string())?; request.context.insert_estimated_cost(cost)?; + request + .context + .insert_estimated_cost_by_subgraph(cost_by_subgraph)?; if cost > self.max { let error = DemandControlError::EstimatedCostTooExpensive { From 301e4fcd56b015a764ce5fce1e2966654f0052ac Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Wed, 7 Jan 2026 11:40:55 -0500 Subject: [PATCH 06/17] feat: reject queries which exceed subgraph limit --- .../cost_calculator/static_cost.rs | 4 ++ .../src/plugins/demand_control/mod.rs | 45 +++++++++++++++++++ .../strategy/static_estimated.rs | 32 ++++++++++--- 3 files changed, 75 insertions(+), 6 deletions(-) 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 b8c2f967af..8a03bf795a 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 @@ -64,6 +64,10 @@ impl CostBySubgraph { } self } + + pub(crate) fn iter(&self) -> impl Iterator + '_ { + self.0.iter() + } } impl AddAssign for CostBySubgraph { diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index bd128321cf..88f31acced 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -165,6 +165,15 @@ pub(crate) enum DemandControlError { /// The maximum cost of the query max_cost: f64, }, + /// subgraph {subgraph} query estimated cost {estimated_cost} exceeded configured maximum {max_cost} + EstimatedSubgraphCostTooExpensive { + /// The subgraph in question + subgraph: String, + /// The estimated cost of the query + estimated_cost: f64, + /// The maximum cost of the query + max_cost: f64, + }, /// auery actual cost {actual_cost} exceeded configured maximum {max_cost} #[allow(dead_code)] ActualCostTooExpensive { @@ -181,6 +190,8 @@ pub(crate) enum DemandControlError { ContextSerializationError(String), /// {0} FederationError(FederationError), + /// multiple query costs exceeded + MultipleCostsTooExpensive(Box>), } impl IntoGraphQLErrors for DemandControlError { @@ -201,6 +212,24 @@ impl IntoGraphQLErrors for DemandControlError { .build(), ]) } + DemandControlError::EstimatedSubgraphCostTooExpensive { + ref subgraph, + estimated_cost, + max_cost, + } => { + let mut extensions = Object::new(); + // TODO: better extensions names? + extensions.insert("cost.subgraph", subgraph.clone().into()); + extensions.insert("cost.subgraph.estimated", estimated_cost.into()); + extensions.insert("cost.subgraph.max", max_cost.into()); + Ok(vec![ + graphql::Error::builder() + .extension_code(self.code()) + .extensions(extensions) + .message(self.to_string()) + .build(), + ]) + } DemandControlError::ActualCostTooExpensive { actual_cost, max_cost, @@ -240,6 +269,13 @@ impl IntoGraphQLErrors for DemandControlError { .message(self.to_string()) .build(), ]), + DemandControlError::MultipleCostsTooExpensive(errors) => { + let mut graphql_errors = Vec::with_capacity(errors.len()); + for error in errors.into_iter() { + graphql_errors.extend(error.into_graphql_errors()?); + } + Ok(graphql_errors) + } } } } @@ -248,6 +284,9 @@ impl DemandControlError { fn code(&self) -> &'static str { match self { DemandControlError::EstimatedCostTooExpensive { .. } => "COST_ESTIMATED_TOO_EXPENSIVE", + DemandControlError::EstimatedSubgraphCostTooExpensive { .. } => { + "SUBGRAPH_COST_ESTIMATED_TOO_EXPENSIVE" + } DemandControlError::ActualCostTooExpensive { .. } => "COST_ACTUAL_TOO_EXPENSIVE", DemandControlError::QueryParseFailure(_) => "COST_QUERY_PARSE_FAILURE", DemandControlError::SubgraphOperationNotInitialized(_) => { @@ -255,6 +294,12 @@ impl DemandControlError { } DemandControlError::ContextSerializationError(_) => "COST_CONTEXT_SERIALIZATION_ERROR", DemandControlError::FederationError(_) => "FEDERATION_ERROR", + DemandControlError::MultipleCostsTooExpensive(errors) => { + errors.first().map(|err| err.code()).unwrap_or( + // should be unreachable, but provide this as a fallback + "COST_ESTIMATED_TOO_EXPENSIVE", + ) + } } } } diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index 539029d43d..1f1a1fc4ab 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -32,26 +32,46 @@ impl StrategyImpl for StaticEstimated { &request.supergraph_request.body().variables, ) .and_then(|cost_by_subgraph| { + let mut errors = vec![]; + let cost = cost_by_subgraph.total(); + + if cost > self.max { + errors.push(DemandControlError::EstimatedCostTooExpensive { + estimated_cost: cost, + max_cost: self.max, + }); + } + + // see if any individual subgraph exceeded its limit + for (subgraph, subgraph_cost) in cost_by_subgraph.iter() { + if let Some(max) = self.subgraph_max(subgraph) + && *subgraph_cost > max + { + errors.push(DemandControlError::EstimatedSubgraphCostTooExpensive { + subgraph: subgraph.clone(), + estimated_cost: *subgraph_cost, + max_cost: max, + }); + } + } + request .context .insert_cost_strategy("static_estimated".to_string())?; - request.context.insert_cost_result("COST_OK".to_string())?; request.context.insert_estimated_cost(cost)?; request .context .insert_estimated_cost_by_subgraph(cost_by_subgraph)?; - if cost > self.max { - let error = DemandControlError::EstimatedCostTooExpensive { - estimated_cost: cost, - max_cost: self.max, - }; + if !errors.is_empty() { + let error = DemandControlError::MultipleCostsTooExpensive(Box::new(errors)); request .context .insert_cost_result(error.code().to_string())?; Err(error) } else { + request.context.insert_cost_result("COST_OK".to_string())?; Ok(()) } }) From 8d2259f9a91224d83afab675bdbf946b33899d9a Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Wed, 7 Jan 2026 12:18:32 -0500 Subject: [PATCH 07/17] feat: publish cost_by_subgraph_estimated to studio --- .../plugins/demand_control/cost_calculator/static_cost.rs | 4 ++-- apollo-router/src/plugins/demand_control/mod.rs | 7 +++++++ .../src/plugins/telemetry/metrics/apollo/studio.rs | 6 ++++++ apollo-router/src/plugins/telemetry/mod.rs | 4 ++++ 4 files changed, 19 insertions(+), 2 deletions(-) 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 8a03bf795a..922cdeb1cb 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 @@ -33,10 +33,10 @@ use crate::query_planner::Primary; use crate::query_planner::QueryPlan; use crate::spec::TYPENAME; -#[derive(Default, Debug, Serialize, Deserialize)] +#[derive(Clone, Default, Debug, Serialize, Deserialize)] pub(crate) struct CostBySubgraph(HashMap); impl CostBySubgraph { - fn new(subgraph: String, value: f64) -> Self { + pub(crate) fn new(subgraph: String, value: f64) -> Self { let mut cost = Self::default(); cost.insert(subgraph, value); cost diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index 88f31acced..5711873078 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -358,6 +358,13 @@ impl Context { Ok(()) } + pub(crate) fn get_estimated_cost_by_subgraph( + &self, + ) -> Result, DemandControlError> { + self.get::<&str, CostBySubgraph>(COST_BY_SUBGRAPH_ESTIMATED_KEY) + .map_err(|e| DemandControlError::ContextSerializationError(e.to_string())) + } + pub(crate) fn insert_actual_cost(&self, cost: f64) -> Result<(), DemandControlError> { self.insert(COST_ACTUAL_KEY, cost) .map_err(|e| DemandControlError::ContextSerializationError(e.to_string()))?; diff --git a/apollo-router/src/plugins/telemetry/metrics/apollo/studio.rs b/apollo-router/src/plugins/telemetry/metrics/apollo/studio.rs index 1566c6b451..eb4769e46c 100644 --- a/apollo-router/src/plugins/telemetry/metrics/apollo/studio.rs +++ b/apollo-router/src/plugins/telemetry/metrics/apollo/studio.rs @@ -11,6 +11,7 @@ use super::histogram::ListLengthHistogram; use crate::apollo_studio_interop::AggregatedExtendedReferenceStats; use crate::apollo_studio_interop::ExtendedReferenceStats; use crate::apollo_studio_interop::ReferencedEnums; +use crate::plugins::demand_control::cost_calculator::static_cost::CostBySubgraph; use crate::plugins::telemetry::apollo::LicensedOperationCountByType; use crate::plugins::telemetry::apollo_exporter::proto::reports::EnumStats; use crate::plugins::telemetry::apollo_exporter::proto::reports::InputFieldStats; @@ -373,6 +374,7 @@ impl From for LimitsStats { pub(crate) struct SingleLimitsStats { pub(crate) strategy: Option, pub(crate) cost_estimated: Option, + pub(crate) cost_by_subgraph_estimated: Option, pub(crate) cost_actual: Option, pub(crate) depth: u64, pub(crate) height: u64, @@ -621,6 +623,10 @@ mod test { limits_stats: SingleLimitsStats { strategy: Some("test".to_string()), cost_estimated: Some(10.0), + cost_by_subgraph_estimated: Some(CostBySubgraph::new( + "products".to_string(), + 10.0, + )), cost_actual: Some(7.0), depth: 2, height: 4, diff --git a/apollo-router/src/plugins/telemetry/mod.rs b/apollo-router/src/plugins/telemetry/mod.rs index af72429b77..1f8c644d95 100644 --- a/apollo-router/src/plugins/telemetry/mod.rs +++ b/apollo-router/src/plugins/telemetry/mod.rs @@ -1474,6 +1474,10 @@ impl Telemetry { strategy: strategy.and_then(|s| serde_json::to_string(&s.mode).ok()), cost_estimated: context.get_estimated_cost().ok().flatten(), cost_actual: context.get_actual_cost().ok().flatten(), + cost_by_subgraph_estimated: context + .get_estimated_cost_by_subgraph() + .ok() + .flatten(), // These limits are related to the Traffic Shaping feature, unrelated to the Demand Control plugin depth: query_limits.map_or(0, |ql| ql.depth as u64), From 63e1694d8aac26b3876a86a6cb7b97b87612d7b5 Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Wed, 7 Jan 2026 15:01:46 -0500 Subject: [PATCH 08/17] lint: remove unnecessary indirection --- apollo-router/src/plugins/demand_control/mod.rs | 2 +- .../src/plugins/demand_control/strategy/static_estimated.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index 5711873078..023e56f4f4 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -191,7 +191,7 @@ pub(crate) enum DemandControlError { /// {0} FederationError(FederationError), /// multiple query costs exceeded - MultipleCostsTooExpensive(Box>), + MultipleCostsTooExpensive(Vec), } impl IntoGraphQLErrors for DemandControlError { diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index 1f1a1fc4ab..b45797d8bf 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -65,7 +65,7 @@ impl StrategyImpl for StaticEstimated { .insert_estimated_cost_by_subgraph(cost_by_subgraph)?; if !errors.is_empty() { - let error = DemandControlError::MultipleCostsTooExpensive(Box::new(errors)); + let error = DemandControlError::MultipleCostsTooExpensive(errors); request .context .insert_cost_result(error.code().to_string())?; @@ -140,7 +140,7 @@ mod tests { }; let strategy = plugin .strategy_factory - .create_static_estimated_strategy(list_size, max, &subgraphs); + .create_static_estimated_strategy(list_size, max, subgraphs); Ok(strategy) } From 9078b0dfe659d1ff1679ab655870d33372ccb101 Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Wed, 7 Jan 2026 15:43:20 -0500 Subject: [PATCH 09/17] fix: don't prevent entire query, just subgraph portion --- .../cost_calculator/static_cost.rs | 6 +- .../src/plugins/demand_control/mod.rs | 41 ++++++++----- .../strategy/static_estimated.rs | 57 ++++++++++--------- 3 files changed, 57 insertions(+), 47 deletions(-) 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 922cdeb1cb..4e32cc8b25 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 @@ -42,7 +42,7 @@ impl CostBySubgraph { cost } - fn get(&self, subgraph: &str) -> Option { + pub(crate) fn get(&self, subgraph: &str) -> Option { self.0.get(subgraph).copied() } @@ -64,10 +64,6 @@ impl CostBySubgraph { } self } - - pub(crate) fn iter(&self) -> impl Iterator + '_ { - self.0.iter() - } } impl AddAssign for CostBySubgraph { diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index 023e56f4f4..d875041b2a 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -57,6 +57,7 @@ pub(crate) const COST_STRATEGY_KEY: &str = "apollo::demand_control::strategy"; pub(crate) const COST_BY_SUBGRAPH_ESTIMATED_KEY: &str = "apollo::demand_control::estimated_cost_by_subgraph"; +pub(crate) const COST_BY_SUBGRAPH_RESULT_KEY: &str = "apollo::demand_control::result_by_subgraph"; /// Algorithm for calculating the cost of an incoming query. #[derive(Clone, Debug, Deserialize, JsonSchema)] @@ -190,8 +191,6 @@ pub(crate) enum DemandControlError { ContextSerializationError(String), /// {0} FederationError(FederationError), - /// multiple query costs exceeded - MultipleCostsTooExpensive(Vec), } impl IntoGraphQLErrors for DemandControlError { @@ -269,13 +268,6 @@ impl IntoGraphQLErrors for DemandControlError { .message(self.to_string()) .build(), ]), - DemandControlError::MultipleCostsTooExpensive(errors) => { - let mut graphql_errors = Vec::with_capacity(errors.len()); - for error in errors.into_iter() { - graphql_errors.extend(error.into_graphql_errors()?); - } - Ok(graphql_errors) - } } } } @@ -294,12 +286,6 @@ impl DemandControlError { } DemandControlError::ContextSerializationError(_) => "COST_CONTEXT_SERIALIZATION_ERROR", DemandControlError::FederationError(_) => "FEDERATION_ERROR", - DemandControlError::MultipleCostsTooExpensive(errors) => { - errors.first().map(|err| err.code()).unwrap_or( - // should be unreachable, but provide this as a fallback - "COST_ESTIMATED_TOO_EXPENSIVE", - ) - } } } } @@ -365,6 +351,16 @@ impl Context { .map_err(|e| DemandControlError::ContextSerializationError(e.to_string())) } + pub(crate) fn get_estimated_cost_for_subgraph( + &self, + subgraph: &str, + ) -> Result, DemandControlError> { + let cost_by_subgraph_opt = self + .get::<&str, CostBySubgraph>(COST_BY_SUBGRAPH_ESTIMATED_KEY) + .map_err(|e| DemandControlError::ContextSerializationError(e.to_string()))?; + Ok(cost_by_subgraph_opt.and_then(|cost_by_subgraph| cost_by_subgraph.get(subgraph))) + } + pub(crate) fn insert_actual_cost(&self, cost: f64) -> Result<(), DemandControlError> { self.insert(COST_ACTUAL_KEY, cost) .map_err(|e| DemandControlError::ContextSerializationError(e.to_string()))?; @@ -393,6 +389,21 @@ impl Context { .map_err(|e| DemandControlError::ContextSerializationError(e.to_string())) } + pub(crate) fn insert_cost_by_subgraph_result( + &self, + subgraph: String, + result: String, + ) -> Result<(), DemandControlError> { + self.upsert( + COST_BY_SUBGRAPH_RESULT_KEY, + |mut map: HashMap| { + map.insert(subgraph, result); + map + }, + ) + .map_err(|e| DemandControlError::ContextSerializationError(e.to_string())) + } + pub(crate) fn insert_cost_strategy(&self, strategy: String) -> Result<(), DemandControlError> { self.insert(COST_STRATEGY_KEY, strategy) .map_err(|e| DemandControlError::ContextSerializationError(e.to_string()))?; diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index b45797d8bf..6c6bb45d1c 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -32,30 +32,7 @@ impl StrategyImpl for StaticEstimated { &request.supergraph_request.body().variables, ) .and_then(|cost_by_subgraph| { - let mut errors = vec![]; - let cost = cost_by_subgraph.total(); - - if cost > self.max { - errors.push(DemandControlError::EstimatedCostTooExpensive { - estimated_cost: cost, - max_cost: self.max, - }); - } - - // see if any individual subgraph exceeded its limit - for (subgraph, subgraph_cost) in cost_by_subgraph.iter() { - if let Some(max) = self.subgraph_max(subgraph) - && *subgraph_cost > max - { - errors.push(DemandControlError::EstimatedSubgraphCostTooExpensive { - subgraph: subgraph.clone(), - estimated_cost: *subgraph_cost, - max_cost: max, - }); - } - } - request .context .insert_cost_strategy("static_estimated".to_string())?; @@ -64,8 +41,11 @@ impl StrategyImpl for StaticEstimated { .context .insert_estimated_cost_by_subgraph(cost_by_subgraph)?; - if !errors.is_empty() { - let error = DemandControlError::MultipleCostsTooExpensive(errors); + if cost > self.max { + let error = DemandControlError::EstimatedCostTooExpensive { + estimated_cost: cost, + max_cost: self.max, + }; request .context .insert_cost_result(error.code().to_string())?; @@ -77,8 +57,31 @@ impl StrategyImpl for StaticEstimated { }) } - fn on_subgraph_request(&self, _request: &subgraph::Request) -> Result<(), DemandControlError> { - Ok(()) + fn on_subgraph_request(&self, request: &subgraph::Request) -> Result<(), DemandControlError> { + let subgraph_name = request.subgraph_name.clone(); + let subgraph_cost = request + .context + .get_estimated_cost_for_subgraph(&subgraph_name)?; + + if let Some(subgraph_cost) = subgraph_cost + && let Some(subgraph_max) = self.subgraph_max(&subgraph_name) + && subgraph_cost > subgraph_max + { + let error = DemandControlError::EstimatedSubgraphCostTooExpensive { + subgraph: subgraph_name.clone(), + estimated_cost: subgraph_cost, + max_cost: subgraph_max, + }; + request + .context + .insert_cost_by_subgraph_result(subgraph_name, error.code().to_string())?; + Err(error) + } else { + request + .context + .insert_cost_by_subgraph_result(subgraph_name, "COST_OK".to_string())?; + Ok(()) + } } fn on_subgraph_response( From 43533c2c3d534d93b416c7834f7fdc5e64beeba9 Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Wed, 7 Jan 2026 18:00:06 -0500 Subject: [PATCH 10/17] feat: measure actual cost per subgraph response --- .../src/plugins/demand_control/mod.rs | 28 ++++++++++++++++- .../plugins/demand_control/strategy/mod.rs | 7 ++++- .../strategy/static_estimated.rs | 30 +++++++++++++++++-- .../plugins/demand_control/strategy/test.rs | 1 + 4 files changed, 62 insertions(+), 4 deletions(-) diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index d875041b2a..49875f53a5 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -57,6 +57,8 @@ pub(crate) const COST_STRATEGY_KEY: &str = "apollo::demand_control::strategy"; pub(crate) const COST_BY_SUBGRAPH_ESTIMATED_KEY: &str = "apollo::demand_control::estimated_cost_by_subgraph"; +pub(crate) const COST_BY_SUBGRAPH_ACTUAL_KEY: &str = + "apollo::demand_control::actual_cost_by_subgraph"; pub(crate) const COST_BY_SUBGRAPH_RESULT_KEY: &str = "apollo::demand_control::result_by_subgraph"; /// Algorithm for calculating the cost of an incoming query. @@ -82,6 +84,7 @@ pub(crate) enum StrategyConfig { /// The maximum cost of a query max: f64, /// Cost control by subgraph + #[serde(default)] subgraphs: SubgraphConfiguration, }, @@ -372,6 +375,29 @@ impl Context { .map_err(|e| DemandControlError::ContextSerializationError(e.to_string())) } + pub(crate) fn update_actual_cost_by_subgraph( + &self, + cost: CostBySubgraph, + ) -> Result<(), DemandControlError> { + // combine this cost with the cost that already exists in the context + self.upsert( + COST_BY_SUBGRAPH_ACTUAL_KEY, + |mut existing_cost: CostBySubgraph| { + existing_cost += cost; + existing_cost + }, + ) + .map_err(|e| DemandControlError::ContextSerializationError(e.to_string()))?; + Ok(()) + } + + pub(crate) fn get_actual_cost_by_subgraph( + &self, + ) -> Result, DemandControlError> { + self.get::<&str, CostBySubgraph>(COST_BY_SUBGRAPH_ACTUAL_KEY) + .map_err(|e| DemandControlError::ContextSerializationError(e.to_string())) + } + pub(crate) fn get_cost_delta(&self) -> Result, DemandControlError> { let estimated = self.get_estimated_cost()?; let actual = self.get_actual_cost()?; @@ -630,7 +656,7 @@ impl Plugin for DemandControl { |(subgraph_name, req): (String, Arc>), fut| async move { let resp: subgraph::Response = fut.await?; let strategy = resp.context.get_demand_control_context().map(|c| c.strategy).expect("must have strategy"); - Ok(match strategy.on_subgraph_response(req.as_ref(), &resp) { + Ok(match strategy.on_subgraph_response(subgraph_name.clone(), req.as_ref(), &resp) { Ok(_) => resp, Err(err) => subgraph::Response::builder() .errors( diff --git a/apollo-router/src/plugins/demand_control/strategy/mod.rs b/apollo-router/src/plugins/demand_control/strategy/mod.rs index 37ee805681..774a44594e 100644 --- a/apollo-router/src/plugins/demand_control/strategy/mod.rs +++ b/apollo-router/src/plugins/demand_control/strategy/mod.rs @@ -52,10 +52,14 @@ impl Strategy { pub(crate) fn on_subgraph_response( &self, + subgraph_name: String, request: &ExecutableDocument, response: &subgraph::Response, ) -> Result<(), DemandControlError> { - match self.inner.on_subgraph_response(request, response) { + match self + .inner + .on_subgraph_response(subgraph_name, request, response) + { Err(e) if self.mode == Mode::Enforce => Err(e), _ => Ok(()), } @@ -142,6 +146,7 @@ pub(crate) trait StrategyImpl: Send + Sync { fn on_subgraph_response( &self, + subgraph_name: String, request: &ExecutableDocument, response: &subgraph::Response, ) -> Result<(), DemandControlError>; diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index 6c6bb45d1c..22fa5b383b 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -5,6 +5,7 @@ use apollo_compiler::ExecutableDocument; use crate::configuration::subgraph::SubgraphConfiguration; use crate::graphql; use crate::plugins::demand_control::DemandControlError; +use crate::plugins::demand_control::cost_calculator::static_cost::CostBySubgraph; use crate::plugins::demand_control::cost_calculator::static_cost::StaticCostCalculator; use crate::plugins::demand_control::strategy::StrategyImpl; use crate::services::execution; @@ -57,6 +58,7 @@ impl StrategyImpl for StaticEstimated { }) } + /// Reject subgraph requests when the total subgraph cost exceeds the subgraph max. fn on_subgraph_request(&self, request: &subgraph::Request) -> Result<(), DemandControlError> { let subgraph_name = request.subgraph_name.clone(); let subgraph_cost = request @@ -84,11 +86,35 @@ impl StrategyImpl for StaticEstimated { } } + /// Parse response to measure actual cost for subgraph. + /// + /// TODO: make this configurable; don't want to incur penalty if we don't want to fn on_subgraph_response( &self, - _request: &ExecutableDocument, - _response: &subgraph::Response, + subgraph_name: String, + request: &ExecutableDocument, + response: &subgraph::Response, ) -> Result<(), DemandControlError> { + if let Some(graphql_response) = response + .response + .body() + .data + .as_ref() + .and_then(|v| graphql::Response::from_value(v.clone()).ok()) + { + let cost = self.cost_calculator.actual( + request, + &graphql_response, + &response + .context + .extensions() + .with_lock(|lock| lock.get().cloned()) + .unwrap_or_default(), + )?; + response + .context + .update_actual_cost_by_subgraph(CostBySubgraph::new(subgraph_name, cost))?; + } Ok(()) } diff --git a/apollo-router/src/plugins/demand_control/strategy/test.rs b/apollo-router/src/plugins/demand_control/strategy/test.rs index 3755b8a5d2..d4c9363317 100644 --- a/apollo-router/src/plugins/demand_control/strategy/test.rs +++ b/apollo-router/src/plugins/demand_control/strategy/test.rs @@ -52,6 +52,7 @@ impl StrategyImpl for Test { fn on_subgraph_response( &self, + _subgraph_name: String, _request: &ExecutableDocument, response: &Response, ) -> Result<(), DemandControlError> { From ec0aed792e7f1180cadb3fbd9a16479222ad834c Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Thu, 8 Jan 2026 14:25:25 -0500 Subject: [PATCH 11/17] test: add integration tests for two supergraphs --- .../strategy/static_estimated.rs | 34 +- .../src/plugins/mock_subgraphs/mod.rs | 12 + .../tests/integration/demand_control.rs | 522 ++++++++++++++++++ apollo-router/tests/integration/mod.rs | 1 + 4 files changed, 550 insertions(+), 19 deletions(-) create mode 100644 apollo-router/tests/integration/demand_control.rs diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index 22fa5b383b..6e5a077a55 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -95,26 +95,22 @@ impl StrategyImpl for StaticEstimated { request: &ExecutableDocument, response: &subgraph::Response, ) -> Result<(), DemandControlError> { - if let Some(graphql_response) = response - .response - .body() - .data - .as_ref() - .and_then(|v| graphql::Response::from_value(v.clone()).ok()) - { - let cost = self.cost_calculator.actual( - request, - &graphql_response, - &response - .context - .extensions() - .with_lock(|lock| lock.get().cloned()) - .unwrap_or_default(), - )?; - response + let subgraph_response_body = response.response.body(); + // TODO: cost calculator cannot currently calculate the cost of an entities fetch + let cost = self.cost_calculator.actual( + request, + subgraph_response_body, + &response .context - .update_actual_cost_by_subgraph(CostBySubgraph::new(subgraph_name, cost))?; - } + .extensions() + .with_lock(|lock| lock.get().cloned()) + .unwrap_or_default(), + )?; + + response + .context + .update_actual_cost_by_subgraph(CostBySubgraph::new(subgraph_name, cost))?; + Ok(()) } diff --git a/apollo-router/src/plugins/mock_subgraphs/mod.rs b/apollo-router/src/plugins/mock_subgraphs/mod.rs index 37a37be5fa..09bb05f0e8 100644 --- a/apollo-router/src/plugins/mock_subgraphs/mod.rs +++ b/apollo-router/src/plugins/mock_subgraphs/mod.rs @@ -140,6 +140,18 @@ impl PluginPrivate for MockSubgraphsPlugin { .build() }; let response = response.body(body).unwrap(); + request + .context + .upsert( + "apollo::experimental_mock_subgraphs::subgraph_call_count", + |mut v: HashMap| { + let subgraph_value = + v.entry(request.subgraph_name.clone()).or_default(); + *subgraph_value += 1; + v + }, + ) + .unwrap(); Ok(subgraph::Response::new_from_response( response, request.context, diff --git a/apollo-router/tests/integration/demand_control.rs b/apollo-router/tests/integration/demand_control.rs new file mode 100644 index 0000000000..0941b4802e --- /dev/null +++ b/apollo-router/tests/integration/demand_control.rs @@ -0,0 +1,522 @@ +use apollo_router::Context; +use apollo_router::graphql; + +const CODE_OK: &str = "COST_OK"; +const CODE_TOO_EXPENSIVE: &str = "COST_ESTIMATED_TOO_EXPENSIVE"; +const CODE_SUBGRAPH_TOO_EXPENSIVE: &str = "SUBGRAPH_COST_ESTIMATED_TOO_EXPENSIVE"; + +fn get_strategy(context: &Context) -> Option { + context + .get::<_, String>("apollo::demand_control::strategy") + .expect("can't deserialize") +} + +fn get_result(context: &Context) -> Option { + context + .get::<_, String>("apollo::demand_control::result") + .expect("can't deserialize") +} + +fn get_result_by_subgraph(context: &Context) -> Option { + context + .get::<_, serde_json::Value>("apollo::demand_control::result_by_subgraph") + .expect("can't deserialize") +} + +fn get_actual_cost(context: &Context) -> Option { + context + .get::<_, f64>("apollo::demand_control::actual_cost") + .expect("can't deserialize") +} + +fn get_estimated_cost(context: &Context) -> Option { + context + .get::<_, f64>("apollo::demand_control::estimated_cost") + .expect("can't deserialize") +} + +fn get_estimated_cost_by_subgraph(context: &Context) -> Option { + context + .get::<_, serde_json::Value>("apollo::demand_control::estimated_cost_by_subgraph") + .expect("can't deserialize") +} + +fn get_subgraph_call_count(context: &Context) -> Option { + context + .get::<_, serde_json::Value>("apollo::experimental_mock_subgraphs::subgraph_call_count") + .expect("can't deserialize") +} + +fn estimated_too_expensive(error: &&graphql::Error) -> bool { + error + .extensions + .get("code") + .map_or(false, |code| code == CODE_TOO_EXPENSIVE) +} + +fn subgraph_estimated_too_expensive(error: &&graphql::Error) -> bool { + error + .extensions + .get("code") + .map_or(false, |code| code == CODE_SUBGRAPH_TOO_EXPENSIVE) +} + +mod basic_fragments_tests { + use apollo_router::TestHarness; + use apollo_router::services::supergraph; + use tokio_stream::StreamExt; + use tower::BoxError; + use tower::ServiceExt; + + use super::CODE_OK; + use super::CODE_SUBGRAPH_TOO_EXPENSIVE; + use super::CODE_TOO_EXPENSIVE; + use super::estimated_too_expensive; + use super::get_estimated_cost; + use super::get_estimated_cost_by_subgraph; + use super::get_result; + use super::get_result_by_subgraph; + use super::get_strategy; + use super::get_subgraph_call_count; + use super::subgraph_estimated_too_expensive; + + fn schema() -> &'static str { + include_str!( + "../../src/plugins/demand_control/cost_calculator/fixtures/basic_supergraph_schema.graphql" + ) + } + + fn query() -> &'static str { + include_str!( + "../../src/plugins/demand_control/cost_calculator/fixtures/basic_fragments_query.graphql" + ) + } + + fn subgraphs() -> serde_json::Value { + serde_json::json!({ + "products": { + "query": { + "interfaceInstance1": {"__typename": "SecondObjectType", "field1": null, "field2": "hello"}, + "someUnion": {"__typename": "FirstObjectType", "innerList": []} + }, + } + }) + } + + async fn supergraph_service( + demand_control: serde_json::Value, + ) -> Result { + TestHarness::builder() + .schema(schema()) + .configuration_json(serde_json::json!({ + "include_subgraph_errors": {"all": true}, + "demand_control": demand_control, + "experimental_mock_subgraphs": subgraphs(), + }))? + .build_supergraph() + .await + } + + async fn query_supergraph_service( + demand_control: serde_json::Value, + ) -> Result { + let service = supergraph_service(demand_control).await?; + let request = supergraph::Request::fake_builder().query(query()).build()?; + service.oneshot(request).await + } + + #[tokio::test(flavor = "multi_thread")] + #[rstest::rstest] + async fn requests_within_max_are_accepted( + #[values(12.0, 15.0)] max_cost: f64, + ) -> Result<(), BoxError> { + // query total cost is 12.0; max_cost >= 12.0 should result in query being accepted + let demand_control = serde_json::json!({ + "enabled": true, + "mode": "enforce", + "strategy": { + "static_estimated": { + "list_size": 10, + "max": max_cost + } + } + }); + + let response = query_supergraph_service(demand_control).await?; + + let context = response.context; + let body = response.response.into_body().next().await.unwrap(); + + // estimates + assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!(&get_result(&context).unwrap(), CODE_OK); + assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); + + assert_eq!( + get_result_by_subgraph(&context).unwrap(), + serde_json::json!({ "products": CODE_OK }) + ); + assert_eq!( + get_estimated_cost_by_subgraph(&context).unwrap(), + serde_json::json!({ "products": 12.0 }) + ); + + // actuals + assert!(body.data.is_some()); + assert!(body.errors.is_empty()); + + let subgraph_call_count = get_subgraph_call_count(&context).unwrap(); + assert_eq!(subgraph_call_count["products"], 1); + + // TODO: check actuals, once we figure out how to handle by-subgraph actuals + + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + #[rstest::rstest] + async fn requests_exceeding_max_are_rejected( + #[values(5.0, 10.0)] max_cost: f64, + ) -> Result<(), BoxError> { + // query total cost is 12.0; all `max_cost` values are less than this, so the response should + // be an error + let demand_control = serde_json::json!({ + "enabled": true, + "mode": "enforce", + "strategy": { + "static_estimated": { + "list_size": 10, + "max": max_cost + } + } + }); + + let response = query_supergraph_service(demand_control).await?; + + let context = response.context; + let body = response.response.into_body().next().await.unwrap(); + + // estimates + assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!(&get_result(&context).unwrap(), CODE_TOO_EXPENSIVE); + assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); + + assert_eq!( + get_estimated_cost_by_subgraph(&context).unwrap(), + serde_json::json!({ "products": 12.0 }) + ); + + // actuals + assert!(body.data.is_none()); + + let error = body.errors.iter().find(estimated_too_expensive).unwrap(); + assert_eq!(error.extensions["cost.estimated"], 12.0); + assert_eq!(error.extensions["cost.max"], max_cost); + + assert!(get_subgraph_call_count(&context).is_none()); + + Ok(()) + } + + #[tokio::test] + async fn requests_which_exceed_subgraph_limit_are_partially_accepted() -> Result<(), BoxError> { + // query checks products once; query should be accepted based on max but products subgraph + // should not be called. + let demand_control = serde_json::json!({ + "enabled": true, + "mode": "enforce", + "strategy": { + "static_estimated": { + "list_size": 10, + "max": 15, + "subgraphs": { + "all": {}, + "subgraphs": { + "products": { + "max": 10 + } + } + } + } + } + }); + + let response = query_supergraph_service(demand_control).await?; + + let context = response.context; + let body = response.response.into_body().next().await.unwrap(); + + // estimates + assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!(&get_result(&context).unwrap(), CODE_OK); + assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); + assert_eq!( + get_result_by_subgraph(&context).unwrap(), + serde_json::json!({ "products": CODE_SUBGRAPH_TOO_EXPENSIVE }) + ); + assert_eq!( + get_estimated_cost_by_subgraph(&context).unwrap(), + serde_json::json!({ "products": 12.0}) + ); + + // actuals + assert!(body.data.is_some()); + assert!(body.errors.iter().find(estimated_too_expensive).is_none()); + + let error = body + .errors + .iter() + .find(subgraph_estimated_too_expensive) + .unwrap(); + assert_eq!(error.extensions["cost.subgraph.estimated"], 12.0); + assert_eq!(error.extensions["cost.subgraph.max"], 10.0); + + let subgraph_call_count = get_subgraph_call_count(&context).unwrap_or_default(); + assert!(subgraph_call_count.get("products").is_none()); + + // TODO: check actuals, once we figure out how to handle by-subgraph actuals + + Ok(()) + } +} + +mod federated_ships_tests { + use apollo_router::TestHarness; + use apollo_router::services::supergraph; + use tokio_stream::StreamExt; + use tower::BoxError; + use tower::ServiceExt; + + use super::CODE_OK; + use super::CODE_SUBGRAPH_TOO_EXPENSIVE; + use super::CODE_TOO_EXPENSIVE; + use super::estimated_too_expensive; + use super::get_estimated_cost; + use super::get_estimated_cost_by_subgraph; + use super::get_result; + use super::get_result_by_subgraph; + use super::get_strategy; + use super::get_subgraph_call_count; + use super::subgraph_estimated_too_expensive; + + fn schema() -> &'static str { + include_str!( + "../../src/plugins/demand_control/cost_calculator/fixtures/federated_ships_schema.graphql" + ) + } + + fn query() -> &'static str { + include_str!( + "../../src/plugins/demand_control/cost_calculator/fixtures/federated_ships_required_query.graphql" + ) + } + + fn subgraphs() -> serde_json::Value { + serde_json::json!({ + "vehicles": { + "query": { + "ships": [ + {"__typename": "Ship", "id": 1, "name": "Ship1", "owner": {"__typename": "User", "licenseNumber": 10},}, + {"__typename": "Ship", "id": 2, "name": "Ship2", "owner": {"__typename": "User", "licenseNumber": 11},}, + {"__typename": "Ship", "id": 3, "name": "Ship3", "owner": {"__typename": "User", "licenseNumber": 12},}, + ], + }, + "entities": [ + {"__typename": "Ship", "id": 1, "owner": {"addresses": [{"zipCode": 18263}]}, "registrationFee": 129.2}, + {"__typename": "Ship", "id": 2, "owner": {"addresses": [{"zipCode": 61027}]}, "registrationFee": 14.0}, + {"__typename": "Ship", "id": 3, "owner": {"addresses": [{"zipCode": 86204}]}, "registrationFee": 97.15}, + {"__typename": "Ship", "id": 1, "owner": null, "registrationFee": null}, + {"__typename": "Ship", "id": 2, "owner": null, "registrationFee": null}, + {"__typename": "Ship", "id": 3, "owner": null, "registrationFee": null}, + ] + }, + "users": { + "entities": [ + {"__typename": "User", "licenseNumber": 10, "addresses": [{"zipCode": 18263}]}, + {"__typename": "User", "licenseNumber": 11, "addresses": [{"zipCode": 61027}]}, + {"__typename": "User", "licenseNumber": 12, "addresses": [{"zipCode": 86204}]}, + ], + } + }) + } + + async fn supergraph_service( + demand_control: serde_json::Value, + ) -> Result { + TestHarness::builder() + .schema(schema()) + .configuration_json(serde_json::json!({ + "include_subgraph_errors": {"all": true}, + "demand_control": demand_control, + "experimental_mock_subgraphs": subgraphs(), + }))? + .build_supergraph() + .await + } + + async fn query_supergraph_service( + demand_control: serde_json::Value, + ) -> Result { + let service = supergraph_service(demand_control).await?; + let request = supergraph::Request::fake_builder().query(query()).build()?; + service.oneshot(request).await + } + + #[tokio::test(flavor = "multi_thread")] + #[rstest::rstest] + async fn requests_within_max_are_accepted( + #[values(10400.0, 10500.0)] max_cost: f64, + ) -> Result<(), BoxError> { + // query total cost is 10400 for list_size = 100; all `max_cost` values are geq than this, + // so the response should be OK + let demand_control = serde_json::json!({ + "enabled": true, + "mode": "enforce", + "strategy": { + "static_estimated": { + "list_size": 100, + "max": max_cost + } + } + }); + + let response = query_supergraph_service(demand_control).await?; + + let context = response.context; + let body = response.response.into_body().next().await.unwrap(); + + // estimates + assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!(&get_result(&context).unwrap(), CODE_OK); + assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); + + assert_eq!( + get_result_by_subgraph(&context).unwrap(), + serde_json::json!({ "users": CODE_OK, "vehicles": CODE_OK }) + ); + assert_eq!( + get_estimated_cost_by_subgraph(&context).unwrap(), + serde_json::json!({ "users": 10100.0, "vehicles": 300.0 }) + ); + + // actuals + assert!(body.data.is_some()); + assert!(body.errors.is_empty()); + + let subgraph_call_count = get_subgraph_call_count(&context).unwrap(); + assert_eq!(subgraph_call_count["users"], 1); + assert_eq!(subgraph_call_count["vehicles"], 2); + + // TODO: check actuals, once we figure out how to handle by-subgraph actuals + + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + #[rstest::rstest] + async fn requests_exceeding_max_are_rejected( + #[values(100.0, 10000.0)] max_cost: f64, + ) -> Result<(), BoxError> { + // query total cost is 10400 for list_size = 100; all `max_cost` values are less than this, + // so the response should be an error + let demand_control = serde_json::json!({ + "enabled": true, + "mode": "enforce", + "strategy": { + "static_estimated": { + "list_size": 100, + "max": max_cost + } + } + }); + + let response = query_supergraph_service(demand_control).await?; + + let context = response.context; + let body = response.response.into_body().next().await.unwrap(); + + // estimates + assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!(&get_result(&context).unwrap(), CODE_TOO_EXPENSIVE); + assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); + + assert_eq!( + get_estimated_cost_by_subgraph(&context).unwrap(), + serde_json::json!({ "users": 10100.0, "vehicles": 300.0 }) + ); + + // actuals + assert!(body.data.is_none()); + + let error = body.errors.iter().find(estimated_too_expensive).unwrap(); + assert_eq!(error.extensions["cost.estimated"], 10400.0); + assert_eq!(error.extensions["cost.max"], max_cost); + + assert!(get_subgraph_call_count(&context).is_none()); + + Ok(()) + } + + #[tokio::test] + async fn requests_which_exceed_subgraph_limit_are_partially_accepted() -> Result<(), BoxError> { + // query checks vehicles, then users, then vehicles. + // interrupting the users check via a demand control limit should still permit both vehicles + // checks. + let demand_control = serde_json::json!({ + "enabled": true, + "mode": "enforce", + "strategy": { + "static_estimated": { + "list_size": 100, + "max": 15000.0, + "subgraphs": { + "all": {}, + "subgraphs": { + "users": { + "max": 0 + } + } + } + } + } + }); + + let response = query_supergraph_service(demand_control).await?; + + let context = response.context; + let body = response.response.into_body().next().await.unwrap(); + + // estimates + assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!(&get_result(&context).unwrap(), CODE_OK); + assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); + assert_eq!( + get_result_by_subgraph(&context).unwrap(), + serde_json::json!({ "users": CODE_SUBGRAPH_TOO_EXPENSIVE, "vehicles": CODE_OK }) + ); + assert_eq!( + get_estimated_cost_by_subgraph(&context).unwrap(), + serde_json::json!({ "users": 10100.0, "vehicles": 300.0 }) + ); + + // actuals + assert!(body.data.is_some()); + assert!(body.errors.iter().find(estimated_too_expensive).is_none()); + + let error = body + .errors + .iter() + .find(subgraph_estimated_too_expensive) + .unwrap(); + assert_eq!(error.extensions["cost.subgraph.estimated"], 10100.0); + assert_eq!(error.extensions["cost.subgraph.max"], 0.0); + + let subgraph_call_count = get_subgraph_call_count(&context).unwrap(); + assert!(subgraph_call_count.get("users").is_none()); + assert_eq!(subgraph_call_count["vehicles"], 2); + + // TODO: check actuals, once we figure out how to handle by-subgraph actuals + + Ok(()) + } +} diff --git a/apollo-router/tests/integration/mod.rs b/apollo-router/tests/integration/mod.rs index dc554f1e6f..aa81b9bdc5 100644 --- a/apollo-router/tests/integration/mod.rs +++ b/apollo-router/tests/integration/mod.rs @@ -10,6 +10,7 @@ pub(crate) mod redis_monitor; mod allowed_features; mod connectors; mod coprocessor; +mod demand_control; mod docs; // In the CI environment we only install Redis on x86_64 Linux mod directives; From c9be0e3ba55b39a9b86a4b10a1ee43485a845881 Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Wed, 14 Jan 2026 11:27:55 -0500 Subject: [PATCH 12/17] chore: migrate features to new strategy --- .../cost_calculator/static_cost.rs | 12 +- .../fixtures/invalid_per_subgraph.yaml | 2 +- .../fixtures/per_subgraph_inheritance.yaml | 2 +- .../fixtures/per_subgraph_no_inheritance.yaml | 2 +- .../src/plugins/demand_control/mod.rs | 24 +- .../plugins/demand_control/strategy/mod.rs | 26 ++- .../strategy/static_estimated.rs | 139 +----------- .../strategy/static_estimated_by_subgraph.rs | 208 ++++++++++++++++++ .../tests/integration/demand_control.rs | 42 +++- 9 files changed, 295 insertions(+), 162 deletions(-) create mode 100644 apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs 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 4e32cc8b25..ec760eb206 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 @@ -206,10 +206,8 @@ impl StaticCostCalculator { } } - fn subgraph_list_size(&self, subgraph_name: &str) -> u32 { - self.subgraph_list_sizes - .get(subgraph_name) - .unwrap_or(self.list_size) + fn subgraph_list_size(&self, subgraph_name: &str) -> Option { + *self.subgraph_list_sizes.get(subgraph_name) } /// Scores a field within a GraphQL operation, handling some expected cases where @@ -269,8 +267,10 @@ impl StaticCostCalculator { .and_then(|dir| dir.expected_size) { expected_size - } else if let Some(subgraph) = subgraph { - self.subgraph_list_size(subgraph) as i32 + } else if let Some(subgraph) = subgraph + && let Some(subgraph_list_size) = self.subgraph_list_size(subgraph) + { + subgraph_list_size as i32 } else { self.list_size as i32 }; diff --git a/apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml b/apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml index 8ddabcb984..eca6789aac 100644 --- a/apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml +++ b/apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml @@ -2,7 +2,7 @@ demand_control: enabled: true mode: enforce strategy: - static_estimated: + static_estimated_by_subgraph: list_size: 1 max: 10 subgraphs: diff --git a/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml index 57d4acc553..2bcb7d7de7 100644 --- a/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml +++ b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml @@ -2,7 +2,7 @@ demand_control: enabled: true mode: enforce strategy: - static_estimated: + static_estimated_by_subgraph: list_size: 1 max: 10 subgraphs: diff --git a/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml index bebb4fbf7c..682f78e86b 100644 --- a/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml +++ b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml @@ -2,7 +2,7 @@ demand_control: enabled: true mode: enforce strategy: - static_estimated: + static_estimated_by_subgraph: list_size: 1 max: 10 subgraphs: diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index 49875f53a5..00d8d31258 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -83,6 +83,26 @@ pub(crate) enum StrategyConfig { list_size: u32, /// The maximum cost of a query max: f64, + }, + + /// A simple, statically-defined cost mapping for operations and types. + /// + /// Operation costs: + /// - Mutation: 10 + /// - Query: 0 + /// - Subscription 0 + /// + /// Type costs: + /// - Object: 1 + /// - Interface: 1 + /// - Union: 1 + /// - Scalar: 0 + /// - Enum: 0 + StaticEstimatedBySubgraph { + /// The assumed length of lists returned by the operation. + list_size: u32, + /// The maximum cost of a query + max: f64, /// Cost control by subgraph #[serde(default)] subgraphs: SubgraphConfiguration, @@ -113,8 +133,8 @@ impl StrategyConfig { } #[allow(irrefutable_let_patterns)] - // need to destructure StrategyConfig and ignore StrategyConfig::Test - let StrategyConfig::StaticEstimated { subgraphs, .. } = self else { + // need to destructure StrategyConfig::StaticEstimatedBySubgraph and ignore the others + let StrategyConfig::StaticEstimatedBySubgraph { subgraphs, .. } = self else { return Ok(()); }; diff --git a/apollo-router/src/plugins/demand_control/strategy/mod.rs b/apollo-router/src/plugins/demand_control/strategy/mod.rs index 774a44594e..e6a636e17e 100644 --- a/apollo-router/src/plugins/demand_control/strategy/mod.rs +++ b/apollo-router/src/plugins/demand_control/strategy/mod.rs @@ -14,10 +14,12 @@ use crate::plugins::demand_control::SubgraphStrategyLimit; use crate::plugins::demand_control::cost_calculator::schema::DemandControlledSchema; use crate::plugins::demand_control::cost_calculator::static_cost::StaticCostCalculator; use crate::plugins::demand_control::strategy::static_estimated::StaticEstimated; +use crate::plugins::demand_control::strategy::static_estimated_by_subgraph::StaticEstimatedBySubgraph; use crate::services::execution; use crate::services::subgraph; mod static_estimated; +mod static_estimated_by_subgraph; #[cfg(test)] mod test; @@ -99,12 +101,12 @@ impl StrategyFactory { // Function extracted for use in tests - allows us to build a `StaticEstimated` directly rather // than a `impl StrategyImpl` - fn create_static_estimated_strategy( + fn create_static_estimated_by_subgraph_strategy( &self, list_size: u32, max: f64, subgraphs: &SubgraphConfiguration, - ) -> StaticEstimated { + ) -> StaticEstimatedBySubgraph { let subgraph_list_sizes = Arc::new(subgraphs.extract(|strategy| strategy.list_size)); let subgraph_maxes = Arc::new(subgraphs.extract(|strategy| strategy.max)); let cost_calculator = StaticCostCalculator::new( @@ -113,7 +115,7 @@ impl StrategyFactory { list_size, subgraph_list_sizes, ); - StaticEstimated { + StaticEstimatedBySubgraph { max, subgraph_maxes, cost_calculator, @@ -122,11 +124,25 @@ impl StrategyFactory { pub(crate) fn create(&self) -> Strategy { let strategy: Arc = match &self.config.strategy { - StrategyConfig::StaticEstimated { + StrategyConfig::StaticEstimated { list_size, max } => { + let cost_calculator = StaticCostCalculator::new( + self.supergraph_schema.clone(), + self.subgraph_schemas.clone(), + *list_size, + Arc::new(SubgraphConfiguration::default()), + ); + Arc::new(StaticEstimated { + max: *max, + cost_calculator, + }) + } + StrategyConfig::StaticEstimatedBySubgraph { list_size, max, subgraphs, - } => Arc::new(self.create_static_estimated_strategy(*list_size, *max, subgraphs)), + } => Arc::new( + self.create_static_estimated_by_subgraph_strategy(*list_size, *max, subgraphs), + ), #[cfg(test)] StrategyConfig::Test { stage, error } => Arc::new(test::Test { stage: stage.clone(), diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index 6e5a077a55..b7d034f764 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -1,11 +1,7 @@ -use std::sync::Arc; - use apollo_compiler::ExecutableDocument; -use crate::configuration::subgraph::SubgraphConfiguration; use crate::graphql; use crate::plugins::demand_control::DemandControlError; -use crate::plugins::demand_control::cost_calculator::static_cost::CostBySubgraph; use crate::plugins::demand_control::cost_calculator::static_cost::StaticCostCalculator; use crate::plugins::demand_control::strategy::StrategyImpl; use crate::services::execution; @@ -15,16 +11,9 @@ use crate::services::subgraph; pub(crate) struct StaticEstimated { // The estimated value of the demand pub(crate) max: f64, - pub(crate) subgraph_maxes: Arc>>, pub(crate) cost_calculator: StaticCostCalculator, } -impl StaticEstimated { - fn subgraph_max(&self, subgraph_name: &str) -> Option { - *self.subgraph_maxes.get(subgraph_name) - } -} - impl StrategyImpl for StaticEstimated { fn on_execution_request(&self, request: &execution::Request) -> Result<(), DemandControlError> { self.cost_calculator @@ -38,9 +27,6 @@ impl StrategyImpl for StaticEstimated { .context .insert_cost_strategy("static_estimated".to_string())?; request.context.insert_estimated_cost(cost)?; - request - .context - .insert_estimated_cost_by_subgraph(cost_by_subgraph)?; if cost > self.max { let error = DemandControlError::EstimatedCostTooExpensive { @@ -58,59 +44,16 @@ impl StrategyImpl for StaticEstimated { }) } - /// Reject subgraph requests when the total subgraph cost exceeds the subgraph max. - fn on_subgraph_request(&self, request: &subgraph::Request) -> Result<(), DemandControlError> { - let subgraph_name = request.subgraph_name.clone(); - let subgraph_cost = request - .context - .get_estimated_cost_for_subgraph(&subgraph_name)?; - - if let Some(subgraph_cost) = subgraph_cost - && let Some(subgraph_max) = self.subgraph_max(&subgraph_name) - && subgraph_cost > subgraph_max - { - let error = DemandControlError::EstimatedSubgraphCostTooExpensive { - subgraph: subgraph_name.clone(), - estimated_cost: subgraph_cost, - max_cost: subgraph_max, - }; - request - .context - .insert_cost_by_subgraph_result(subgraph_name, error.code().to_string())?; - Err(error) - } else { - request - .context - .insert_cost_by_subgraph_result(subgraph_name, "COST_OK".to_string())?; - Ok(()) - } + fn on_subgraph_request(&self, _request: &subgraph::Request) -> Result<(), DemandControlError> { + Ok(()) } - /// Parse response to measure actual cost for subgraph. - /// - /// TODO: make this configurable; don't want to incur penalty if we don't want to fn on_subgraph_response( &self, - subgraph_name: String, - request: &ExecutableDocument, - response: &subgraph::Response, + _subgraph_name: String, + _request: &ExecutableDocument, + _response: &subgraph::Response, ) -> Result<(), DemandControlError> { - let subgraph_response_body = response.response.body(); - // TODO: cost calculator cannot currently calculate the cost of an entities fetch - let cost = self.cost_calculator.actual( - request, - subgraph_response_body, - &response - .context - .extensions() - .with_lock(|lock| lock.get().cloned()) - .unwrap_or_default(), - )?; - - response - .context - .update_actual_cost_by_subgraph(CostBySubgraph::new(subgraph_name, cost))?; - Ok(()) } @@ -134,75 +77,3 @@ impl StrategyImpl for StaticEstimated { Ok(()) } } - -#[cfg(test)] -mod tests { - use tower::BoxError; - - use super::StaticEstimated; - use crate::plugins::demand_control::DemandControl; - use crate::plugins::demand_control::StrategyConfig; - use crate::plugins::test::PluginTestHarness; - - async fn load_config_and_extract_strategy( - config: &'static str, - ) -> Result { - let schema_str = - include_str!("../cost_calculator/fixtures/basic_supergraph_schema.graphql"); - let plugin = PluginTestHarness::::builder() - .config(config) - .schema(schema_str) - .build() - .await?; - - let StrategyConfig::StaticEstimated { - list_size, - max, - ref subgraphs, - } = plugin.config.strategy - else { - panic!("must provide static_estimated config"); - }; - let strategy = plugin - .strategy_factory - .create_static_estimated_strategy(list_size, max, subgraphs); - Ok(strategy) - } - - #[tokio::test] - async fn test_per_subgraph_configuration_inheritance() { - let config = include_str!("../fixtures/per_subgraph_inheritance.yaml"); - - let strategy = load_config_and_extract_strategy(config).await.unwrap(); - assert_eq!(strategy.subgraph_max("reviews").unwrap(), 2.0); - assert_eq!(strategy.subgraph_max("products").unwrap(), 5.0); - assert_eq!(strategy.subgraph_max("users").unwrap(), 5.0); - } - - #[tokio::test] - async fn test_per_subgraph_configuration_no_inheritance() { - let config = include_str!("../fixtures/per_subgraph_no_inheritance.yaml"); - - let strategy = load_config_and_extract_strategy(config).await.unwrap(); - assert_eq!(strategy.subgraph_max("reviews").unwrap(), 2.0); - assert!(strategy.subgraph_max("products").is_none()); - assert!(strategy.subgraph_max("users").is_none()); - } - - #[tokio::test] - async fn test_invalid_per_subgraph_configuration() { - let config = include_str!("../fixtures/invalid_per_subgraph.yaml"); - let strategy_result = load_config_and_extract_strategy(config).await; - - match strategy_result { - Ok(strategy) => { - eprintln!("{:?}", strategy.subgraph_maxes); - panic!("Expected error") - } - Err(err) => assert_eq!( - &err.to_string(), - "Maximum per-subgraph query cost for `products` is negative" - ), - }; - } -} diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs new file mode 100644 index 0000000000..df9e1764fc --- /dev/null +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs @@ -0,0 +1,208 @@ +use std::sync::Arc; + +use apollo_compiler::ExecutableDocument; + +use crate::configuration::subgraph::SubgraphConfiguration; +use crate::graphql; +use crate::plugins::demand_control::DemandControlError; +use crate::plugins::demand_control::cost_calculator::static_cost::CostBySubgraph; +use crate::plugins::demand_control::cost_calculator::static_cost::StaticCostCalculator; +use crate::plugins::demand_control::strategy::StrategyImpl; +use crate::services::execution; +use crate::services::subgraph; + +/// This strategy will reject requests if the estimated cost of the request exceeds the maximum cost. +pub(crate) struct StaticEstimatedBySubgraph { + // The estimated value of the demand + pub(crate) max: f64, + pub(crate) subgraph_maxes: Arc>>, + pub(crate) cost_calculator: StaticCostCalculator, +} + +impl StaticEstimatedBySubgraph { + fn subgraph_max(&self, subgraph_name: &str) -> Option { + *self.subgraph_maxes.get(subgraph_name) + } +} + +impl StrategyImpl for StaticEstimatedBySubgraph { + fn on_execution_request(&self, request: &execution::Request) -> Result<(), DemandControlError> { + self.cost_calculator + .planned( + &request.query_plan, + &request.supergraph_request.body().variables, + ) + .and_then(|cost_by_subgraph| { + let cost = cost_by_subgraph.total(); + request + .context + .insert_cost_strategy("static_estimated_by_subgraph".to_string())?; + request.context.insert_estimated_cost(cost)?; + request + .context + .insert_estimated_cost_by_subgraph(cost_by_subgraph)?; + + if cost > self.max { + let error = DemandControlError::EstimatedCostTooExpensive { + estimated_cost: cost, + max_cost: self.max, + }; + request + .context + .insert_cost_result(error.code().to_string())?; + Err(error) + } else { + request.context.insert_cost_result("COST_OK".to_string())?; + Ok(()) + } + }) + } + + /// Reject subgraph requests when the total subgraph cost exceeds the subgraph max. + fn on_subgraph_request(&self, request: &subgraph::Request) -> Result<(), DemandControlError> { + let subgraph_name = request.subgraph_name.clone(); + let subgraph_cost = request + .context + .get_estimated_cost_for_subgraph(&subgraph_name)?; + + if let Some(subgraph_cost) = subgraph_cost + && let Some(subgraph_max) = self.subgraph_max(&subgraph_name) + && subgraph_cost > subgraph_max + { + let error = DemandControlError::EstimatedSubgraphCostTooExpensive { + subgraph: subgraph_name.clone(), + estimated_cost: subgraph_cost, + max_cost: subgraph_max, + }; + request + .context + .insert_cost_by_subgraph_result(subgraph_name, error.code().to_string())?; + Err(error) + } else { + request + .context + .insert_cost_by_subgraph_result(subgraph_name, "COST_OK".to_string())?; + Ok(()) + } + } + + /// Parse response to measure actual cost for subgraph. + /// + /// TODO: make this configurable; don't want to incur penalty if we don't want to + fn on_subgraph_response( + &self, + subgraph_name: String, + request: &ExecutableDocument, + response: &subgraph::Response, + ) -> Result<(), DemandControlError> { + let subgraph_response_body = response.response.body(); + // TODO: cost calculator cannot currently calculate the cost of an entities fetch + let cost = self.cost_calculator.actual( + request, + subgraph_response_body, + &response + .context + .extensions() + .with_lock(|lock| lock.get().cloned()) + .unwrap_or_default(), + )?; + + response + .context + .update_actual_cost_by_subgraph(CostBySubgraph::new(subgraph_name, cost))?; + + Ok(()) + } + + fn on_execution_response( + &self, + context: &crate::Context, + request: &ExecutableDocument, + response: &graphql::Response, + ) -> Result<(), DemandControlError> { + if response.data.is_some() { + let cost = self.cost_calculator.actual( + request, + response, + &context + .extensions() + .with_lock(|lock| lock.get().cloned()) + .unwrap_or_default(), + )?; + context.insert_actual_cost(cost)?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use tower::BoxError; + + use super::StaticEstimatedBySubgraph; + use crate::plugins::demand_control::DemandControl; + use crate::plugins::demand_control::StrategyConfig; + use crate::plugins::test::PluginTestHarness; + + async fn load_config_and_extract_strategy( + config: &'static str, + ) -> Result { + let schema_str = + include_str!("../cost_calculator/fixtures/basic_supergraph_schema.graphql"); + let plugin = PluginTestHarness::::builder() + .config(config) + .schema(schema_str) + .build() + .await?; + + let StrategyConfig::StaticEstimatedBySubgraph { + list_size, + max, + ref subgraphs, + } = plugin.config.strategy + else { + panic!("must provide static_estimated_by_subgraph config"); + }; + let strategy = plugin + .strategy_factory + .create_static_estimated_by_subgraph_strategy(list_size, max, subgraphs); + Ok(strategy) + } + + #[tokio::test] + async fn test_per_subgraph_configuration_inheritance() { + let config = include_str!("../fixtures/per_subgraph_inheritance.yaml"); + + let strategy = load_config_and_extract_strategy(config).await.unwrap(); + assert_eq!(strategy.subgraph_max("reviews").unwrap(), 2.0); + assert_eq!(strategy.subgraph_max("products").unwrap(), 5.0); + assert_eq!(strategy.subgraph_max("users").unwrap(), 5.0); + } + + #[tokio::test] + async fn test_per_subgraph_configuration_no_inheritance() { + let config = include_str!("../fixtures/per_subgraph_no_inheritance.yaml"); + + let strategy = load_config_and_extract_strategy(config).await.unwrap(); + assert_eq!(strategy.subgraph_max("reviews").unwrap(), 2.0); + assert!(strategy.subgraph_max("products").is_none()); + assert!(strategy.subgraph_max("users").is_none()); + } + + #[tokio::test] + async fn test_invalid_per_subgraph_configuration() { + let config = include_str!("../fixtures/invalid_per_subgraph.yaml"); + let strategy_result = load_config_and_extract_strategy(config).await; + + match strategy_result { + Ok(strategy) => { + eprintln!("{:?}", strategy.subgraph_maxes); + panic!("Expected error") + } + Err(err) => assert_eq!( + &err.to_string(), + "Maximum per-subgraph query cost for `products` is negative" + ), + }; + } +} diff --git a/apollo-router/tests/integration/demand_control.rs b/apollo-router/tests/integration/demand_control.rs index 0941b4802e..9e7df330e7 100644 --- a/apollo-router/tests/integration/demand_control.rs +++ b/apollo-router/tests/integration/demand_control.rs @@ -135,7 +135,7 @@ mod basic_fragments_tests { "enabled": true, "mode": "enforce", "strategy": { - "static_estimated": { + "static_estimated_by_subgraph": { "list_size": 10, "max": max_cost } @@ -148,7 +148,10 @@ mod basic_fragments_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!( + &get_strategy(&context).unwrap(), + "static_estimated_by_subgraph" + ); assert_eq!(&get_result(&context).unwrap(), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); @@ -184,7 +187,7 @@ mod basic_fragments_tests { "enabled": true, "mode": "enforce", "strategy": { - "static_estimated": { + "static_estimated_by_subgraph": { "list_size": 10, "max": max_cost } @@ -197,7 +200,10 @@ mod basic_fragments_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!( + &get_strategy(&context).unwrap(), + "static_estimated_by_subgraph" + ); assert_eq!(&get_result(&context).unwrap(), CODE_TOO_EXPENSIVE); assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); @@ -226,7 +232,7 @@ mod basic_fragments_tests { "enabled": true, "mode": "enforce", "strategy": { - "static_estimated": { + "static_estimated_by_subgraph": { "list_size": 10, "max": 15, "subgraphs": { @@ -247,7 +253,10 @@ mod basic_fragments_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!( + &get_strategy(&context).unwrap(), + "static_estimated_by_subgraph" + ); assert_eq!(&get_result(&context).unwrap(), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); assert_eq!( @@ -373,7 +382,7 @@ mod federated_ships_tests { "enabled": true, "mode": "enforce", "strategy": { - "static_estimated": { + "static_estimated_by_subgraph": { "list_size": 100, "max": max_cost } @@ -386,7 +395,10 @@ mod federated_ships_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!( + &get_strategy(&context).unwrap(), + "static_estimated_by_subgraph" + ); assert_eq!(&get_result(&context).unwrap(), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); @@ -423,7 +435,7 @@ mod federated_ships_tests { "enabled": true, "mode": "enforce", "strategy": { - "static_estimated": { + "static_estimated_by_subgraph": { "list_size": 100, "max": max_cost } @@ -436,7 +448,10 @@ mod federated_ships_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!( + &get_strategy(&context).unwrap(), + "static_estimated_by_subgraph" + ); assert_eq!(&get_result(&context).unwrap(), CODE_TOO_EXPENSIVE); assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); @@ -466,7 +481,7 @@ mod federated_ships_tests { "enabled": true, "mode": "enforce", "strategy": { - "static_estimated": { + "static_estimated_by_subgraph": { "list_size": 100, "max": 15000.0, "subgraphs": { @@ -487,7 +502,10 @@ mod federated_ships_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context).unwrap(), "static_estimated"); + assert_eq!( + &get_strategy(&context).unwrap(), + "static_estimated_by_subgraph" + ); assert_eq!(&get_result(&context).unwrap(), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); assert_eq!( From 4f723409792504f7a025318d676a61d6acdf04f1 Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Wed, 14 Jan 2026 11:31:55 -0500 Subject: [PATCH 13/17] test: strategy and result should always be present --- .../tests/integration/demand_control.rs | 54 +++++++------------ 1 file changed, 20 insertions(+), 34 deletions(-) diff --git a/apollo-router/tests/integration/demand_control.rs b/apollo-router/tests/integration/demand_control.rs index 9e7df330e7..5a649b571a 100644 --- a/apollo-router/tests/integration/demand_control.rs +++ b/apollo-router/tests/integration/demand_control.rs @@ -5,16 +5,20 @@ const CODE_OK: &str = "COST_OK"; const CODE_TOO_EXPENSIVE: &str = "COST_ESTIMATED_TOO_EXPENSIVE"; const CODE_SUBGRAPH_TOO_EXPENSIVE: &str = "SUBGRAPH_COST_ESTIMATED_TOO_EXPENSIVE"; -fn get_strategy(context: &Context) -> Option { +fn get_strategy(context: &Context) -> String { + let field = "apollo::demand_control::strategy"; context - .get::<_, String>("apollo::demand_control::strategy") + .get::<_, String>(field) .expect("can't deserialize") + .expect(&format!("context missing {field}")) } -fn get_result(context: &Context) -> Option { +fn get_result(context: &Context) -> String { + let field = "apollo::demand_control::result"; context - .get::<_, String>("apollo::demand_control::result") + .get::<_, String>(field) .expect("can't deserialize") + .expect(&format!("context missing {field}")) } fn get_result_by_subgraph(context: &Context) -> Option { @@ -148,11 +152,8 @@ mod basic_fragments_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!( - &get_strategy(&context).unwrap(), - "static_estimated_by_subgraph" - ); - assert_eq!(&get_result(&context).unwrap(), CODE_OK); + assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_result(&context), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); assert_eq!( @@ -200,11 +201,8 @@ mod basic_fragments_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!( - &get_strategy(&context).unwrap(), - "static_estimated_by_subgraph" - ); - assert_eq!(&get_result(&context).unwrap(), CODE_TOO_EXPENSIVE); + assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_result(&context), CODE_TOO_EXPENSIVE); assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); assert_eq!( @@ -253,11 +251,8 @@ mod basic_fragments_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!( - &get_strategy(&context).unwrap(), - "static_estimated_by_subgraph" - ); - assert_eq!(&get_result(&context).unwrap(), CODE_OK); + assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_result(&context), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); assert_eq!( get_result_by_subgraph(&context).unwrap(), @@ -395,11 +390,8 @@ mod federated_ships_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!( - &get_strategy(&context).unwrap(), - "static_estimated_by_subgraph" - ); - assert_eq!(&get_result(&context).unwrap(), CODE_OK); + assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_result(&context), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); assert_eq!( @@ -448,11 +440,8 @@ mod federated_ships_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!( - &get_strategy(&context).unwrap(), - "static_estimated_by_subgraph" - ); - assert_eq!(&get_result(&context).unwrap(), CODE_TOO_EXPENSIVE); + assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_result(&context), CODE_TOO_EXPENSIVE); assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); assert_eq!( @@ -502,11 +491,8 @@ mod federated_ships_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!( - &get_strategy(&context).unwrap(), - "static_estimated_by_subgraph" - ); - assert_eq!(&get_result(&context).unwrap(), CODE_OK); + assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_result(&context), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); assert_eq!( get_result_by_subgraph(&context).unwrap(), From 901bea2fd8de5b980b6e8f4d35b0444e09b9de83 Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Tue, 20 Jan 2026 16:04:54 -0500 Subject: [PATCH 14/17] feat: accurately measure costs per subgraph return object --- .../cost_calculator/static_cost.rs | 104 ++++++++++-------- .../strategy/static_estimated.rs | 1 + .../strategy/static_estimated_by_subgraph.rs | 18 +-- .../tests/integration/demand_control.rs | 87 ++++++++++----- 4 files changed, 127 insertions(+), 83 deletions(-) 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 ec760eb206..97fdaff64a 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 @@ -607,8 +607,9 @@ impl StaticCostCalculator { request: &ExecutableDocument, response: &Response, variables: &Object, + include_entities: bool, ) -> Result { - let mut visitor = ResponseCostCalculator::new(&self.supergraph_schema); + let mut visitor = ResponseCostCalculator::new(&self.supergraph_schema, include_entities); visitor.visit(request, response, variables); Ok(visitor.cost) } @@ -617,11 +618,16 @@ impl StaticCostCalculator { pub(crate) struct ResponseCostCalculator<'a> { pub(crate) cost: f64, schema: &'a DemandControlledSchema, + include_entities: bool, } impl<'schema> ResponseCostCalculator<'schema> { - pub(crate) fn new(schema: &'schema DemandControlledSchema) -> Self { - Self { cost: 0.0, schema } + pub(crate) fn new(schema: &'schema DemandControlledSchema, include_entities: bool) -> Self { + Self { + cost: 0.0, + schema, + include_entities, + } } fn score_response_field( @@ -637,53 +643,61 @@ impl<'schema> ResponseCostCalculator<'schema> { if field.name == TYPENAME { return; } - 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.output_field_definition(parent_ty, &field.name); + + // if the definition is None, that means one of two things: + // (1) the field is missing from the schema, or + // (2) the query is an `_entities` query. + // If the field _should_ be there and isn't, or we don't want to score entities, return now. + let is_entities_query = parent_ty == "Query" && field.name == "_entities"; + if definition.is_none() && !(is_entities_query && self.include_entities) { + tracing::debug!( + "Failed to get schema definition for field {}.{}. The resulting response cost will be a partial result.", + parent_ty, + field.name, + ); + return; + } + + match value { + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => { + self.cost += definition + .and_then(|d| d.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 + .and_then(|d| d.cost_directive()) + .map_or(1.0, |cost| cost.weight()); + self.visit_selections(request, variables, &field.selection_set, children); + } + } - 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; - } + if include_argument_score && let Some(definition) = definition { + 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::debug!( - "Failed to get schema definition for argument {}.{}({}:). The resulting response cost will be a partial result.", - parent_ty, - field.name, - argument.name, - ) + eprintln!("argument score is none"); } + } else { + tracing::debug!( + "Failed to get schema definition for argument {}.{}({}:). The resulting response cost will be a partial result.", + parent_ty, + field.name, + argument.name, + ) } } - } else { - tracing::debug!( - "Failed to get schema definition for field {}.{}. The resulting response cost will be a partial result.", - parent_ty, - field.name, - ) } } } @@ -936,7 +950,7 @@ mod tests { 100, Default::default(), ) - .actual(&query.executable, &response, &variables) + .actual(&query.executable, &response, &variables, false) .unwrap() } @@ -969,7 +983,7 @@ mod tests { 100, Default::default(), ) - .actual(&query, &response, &variables) + .actual(&query, &response, &variables, false) .unwrap() } diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index b7d034f764..72226c3dff 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -71,6 +71,7 @@ impl StrategyImpl for StaticEstimated { .extensions() .with_lock(|lock| lock.get().cloned()) .unwrap_or_default(), + false, )?; context.insert_actual_cost(cost)?; } diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs index df9e1764fc..2ce7c11bc0 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs @@ -87,8 +87,6 @@ impl StrategyImpl for StaticEstimatedBySubgraph { } /// Parse response to measure actual cost for subgraph. - /// - /// TODO: make this configurable; don't want to incur penalty if we don't want to fn on_subgraph_response( &self, subgraph_name: String, @@ -96,7 +94,6 @@ impl StrategyImpl for StaticEstimatedBySubgraph { response: &subgraph::Response, ) -> Result<(), DemandControlError> { let subgraph_response_body = response.response.body(); - // TODO: cost calculator cannot currently calculate the cost of an entities fetch let cost = self.cost_calculator.actual( request, subgraph_response_body, @@ -105,6 +102,7 @@ impl StrategyImpl for StaticEstimatedBySubgraph { .extensions() .with_lock(|lock| lock.get().cloned()) .unwrap_or_default(), + true, )?; response @@ -117,18 +115,14 @@ impl StrategyImpl for StaticEstimatedBySubgraph { fn on_execution_response( &self, context: &crate::Context, - request: &ExecutableDocument, + _request: &ExecutableDocument, response: &graphql::Response, ) -> Result<(), DemandControlError> { + // sum up values from subgraph responses if response.data.is_some() { - let cost = self.cost_calculator.actual( - request, - response, - &context - .extensions() - .with_lock(|lock| lock.get().cloned()) - .unwrap_or_default(), - )?; + let cost = context + .get_actual_cost_by_subgraph()? + .map_or(0.0, |cost| cost.total()); context.insert_actual_cost(cost)?; } Ok(()) diff --git a/apollo-router/tests/integration/demand_control.rs b/apollo-router/tests/integration/demand_control.rs index 5a649b571a..d044e63b45 100644 --- a/apollo-router/tests/integration/demand_control.rs +++ b/apollo-router/tests/integration/demand_control.rs @@ -10,7 +10,7 @@ fn get_strategy(context: &Context) -> String { context .get::<_, String>(field) .expect("can't deserialize") - .expect(&format!("context missing {field}")) + .unwrap_or_else(|| panic!("context missing {field}")) } fn get_result(context: &Context) -> String { @@ -18,7 +18,7 @@ fn get_result(context: &Context) -> String { context .get::<_, String>(field) .expect("can't deserialize") - .expect(&format!("context missing {field}")) + .unwrap_or_else(|| panic!("context missing {field}")) } fn get_result_by_subgraph(context: &Context) -> Option { @@ -33,6 +33,12 @@ fn get_actual_cost(context: &Context) -> Option { .expect("can't deserialize") } +fn get_actual_cost_by_subgraph(context: &Context) -> Option { + context + .get::<_, serde_json::Value>("apollo::demand_control::actual_cost_by_subgraph") + .expect("can't deserialize") +} + fn get_estimated_cost(context: &Context) -> Option { context .get::<_, f64>("apollo::demand_control::estimated_cost") @@ -55,16 +61,19 @@ fn estimated_too_expensive(error: &&graphql::Error) -> bool { error .extensions .get("code") - .map_or(false, |code| code == CODE_TOO_EXPENSIVE) + .is_some_and(|code| code == CODE_TOO_EXPENSIVE) } fn subgraph_estimated_too_expensive(error: &&graphql::Error) -> bool { error .extensions .get("code") - .map_or(false, |code| code == CODE_SUBGRAPH_TOO_EXPENSIVE) + .is_some_and(|code| code == CODE_SUBGRAPH_TOO_EXPENSIVE) } +// TODO: add tests for static_estimated as well. would be good to have tests showing the difference +// between actual costs, and asserting that estimated costs are the same + mod basic_fragments_tests { use apollo_router::TestHarness; use apollo_router::services::supergraph; @@ -76,6 +85,8 @@ mod basic_fragments_tests { use super::CODE_SUBGRAPH_TOO_EXPENSIVE; use super::CODE_TOO_EXPENSIVE; use super::estimated_too_expensive; + use super::get_actual_cost; + use super::get_actual_cost_by_subgraph; use super::get_estimated_cost; use super::get_estimated_cost_by_subgraph; use super::get_result; @@ -131,7 +142,7 @@ mod basic_fragments_tests { #[tokio::test(flavor = "multi_thread")] #[rstest::rstest] - async fn requests_within_max_are_accepted( + async fn requests_within_max_are_accepted_with_static_estimated_by_subgraph_strategy( #[values(12.0, 15.0)] max_cost: f64, ) -> Result<(), BoxError> { // query total cost is 12.0; max_cost >= 12.0 should result in query being accepted @@ -172,18 +183,22 @@ mod basic_fragments_tests { let subgraph_call_count = get_subgraph_call_count(&context).unwrap(); assert_eq!(subgraph_call_count["products"], 1); - // TODO: check actuals, once we figure out how to handle by-subgraph actuals + assert_eq!(get_actual_cost(&context).unwrap(), 2.0); + assert_eq!( + get_actual_cost_by_subgraph(&context).unwrap(), + serde_json::json!({ "products": 2.0 }) + ); Ok(()) } #[tokio::test(flavor = "multi_thread")] - #[rstest::rstest] - async fn requests_exceeding_max_are_rejected( - #[values(5.0, 10.0)] max_cost: f64, - ) -> Result<(), BoxError> { - // query total cost is 12.0; all `max_cost` values are less than this, so the response should - // be an error + async fn requests_exceeding_max_are_rejected_with_static_estimated_by_subgraph_strategy() + -> Result<(), BoxError> { + let max_cost = 10.0; + + // query total cost is 12.0 for list_size = 10; since `max_cost` value is less than this, + // the response should be an error let demand_control = serde_json::json!({ "enabled": true, "mode": "enforce", @@ -219,11 +234,15 @@ mod basic_fragments_tests { assert!(get_subgraph_call_count(&context).is_none()); + assert!(get_actual_cost(&context).is_none()); + assert!(get_actual_cost_by_subgraph(&context).is_none()); + Ok(()) } #[tokio::test] - async fn requests_which_exceed_subgraph_limit_are_partially_accepted() -> Result<(), BoxError> { + async fn requests_which_exceed_subgraph_limit_are_partially_accepted_with_static_estimated_by_subgraph_strategy() + -> Result<(), BoxError> { // query checks products once; query should be accepted based on max but products subgraph // should not be called. let demand_control = serde_json::json!({ @@ -265,7 +284,7 @@ mod basic_fragments_tests { // actuals assert!(body.data.is_some()); - assert!(body.errors.iter().find(estimated_too_expensive).is_none()); + assert!(!body.errors.iter().any(|e| estimated_too_expensive(&e))); let error = body .errors @@ -278,7 +297,8 @@ mod basic_fragments_tests { let subgraph_call_count = get_subgraph_call_count(&context).unwrap_or_default(); assert!(subgraph_call_count.get("products").is_none()); - // TODO: check actuals, once we figure out how to handle by-subgraph actuals + assert_eq!(get_actual_cost(&context).unwrap(), 0.0); + assert!(get_actual_cost_by_subgraph(&context).is_none()); Ok(()) } @@ -295,6 +315,8 @@ mod federated_ships_tests { use super::CODE_SUBGRAPH_TOO_EXPENSIVE; use super::CODE_TOO_EXPENSIVE; use super::estimated_too_expensive; + use super::get_actual_cost; + use super::get_actual_cost_by_subgraph; use super::get_estimated_cost; use super::get_estimated_cost_by_subgraph; use super::get_result; @@ -368,7 +390,7 @@ mod federated_ships_tests { #[tokio::test(flavor = "multi_thread")] #[rstest::rstest] - async fn requests_within_max_are_accepted( + async fn requests_within_max_are_accepted_with_static_estimated_by_subgraph_strategy( #[values(10400.0, 10500.0)] max_cost: f64, ) -> Result<(), BoxError> { // query total cost is 10400 for list_size = 100; all `max_cost` values are geq than this, @@ -411,18 +433,22 @@ mod federated_ships_tests { assert_eq!(subgraph_call_count["users"], 1); assert_eq!(subgraph_call_count["vehicles"], 2); - // TODO: check actuals, once we figure out how to handle by-subgraph actuals + assert_eq!(get_actual_cost(&context).unwrap(), 15.0); + assert_eq!( + get_actual_cost_by_subgraph(&context).unwrap(), + serde_json::json!({ "users": 6.0, "vehicles": 9.0 }) + ); Ok(()) } #[tokio::test(flavor = "multi_thread")] - #[rstest::rstest] - async fn requests_exceeding_max_are_rejected( - #[values(100.0, 10000.0)] max_cost: f64, - ) -> Result<(), BoxError> { - // query total cost is 10400 for list_size = 100; all `max_cost` values are less than this, - // so the response should be an error + async fn requests_exceeding_max_are_rejected_with_static_estimated_by_subgraph_strategy() + -> Result<(), BoxError> { + let max_cost = 10000.0; + + // query total cost is 10400 for list_size = 100; since `max_cost` value is less than this, + // the response should be an error let demand_control = serde_json::json!({ "enabled": true, "mode": "enforce", @@ -458,11 +484,15 @@ mod federated_ships_tests { assert!(get_subgraph_call_count(&context).is_none()); + assert!(get_actual_cost(&context).is_none()); + assert!(get_actual_cost_by_subgraph(&context).is_none()); + Ok(()) } #[tokio::test] - async fn requests_which_exceed_subgraph_limit_are_partially_accepted() -> Result<(), BoxError> { + async fn requests_which_exceed_subgraph_limit_are_partially_accepted_with_static_estimated_by_subgraph_strategy() + -> Result<(), BoxError> { // query checks vehicles, then users, then vehicles. // interrupting the users check via a demand control limit should still permit both vehicles // checks. @@ -505,7 +535,7 @@ mod federated_ships_tests { // actuals assert!(body.data.is_some()); - assert!(body.errors.iter().find(estimated_too_expensive).is_none()); + assert!(!body.errors.iter().any(|e| estimated_too_expensive(&e))); let error = body .errors @@ -519,7 +549,12 @@ mod federated_ships_tests { assert!(subgraph_call_count.get("users").is_none()); assert_eq!(subgraph_call_count["vehicles"], 2); - // TODO: check actuals, once we figure out how to handle by-subgraph actuals + assert_eq!(get_actual_cost(&context).unwrap(), 9.0); + assert_eq!( + get_actual_cost_by_subgraph(&context).unwrap(), + serde_json::json!({ "vehicles": 9.0 }) + ); + Ok(()) } From 15b19627ae5c9453f1e38d1e1f5fb8d4888c5242 Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Tue, 20 Jan 2026 16:05:56 -0500 Subject: [PATCH 15/17] test: add tests for old static_estimated strategy --- .../tests/integration/demand_control.rs | 187 ++++++++++++++++++ 1 file changed, 187 insertions(+) diff --git a/apollo-router/tests/integration/demand_control.rs b/apollo-router/tests/integration/demand_control.rs index d044e63b45..6a07050f08 100644 --- a/apollo-router/tests/integration/demand_control.rs +++ b/apollo-router/tests/integration/demand_control.rs @@ -302,6 +302,98 @@ mod basic_fragments_tests { Ok(()) } + + // Duplicates `requests_within_max_are_accepted_with_static_estimated_by_subgraph_strategy` but + // uses old `static_estimated` strategy. + #[tokio::test(flavor = "multi_thread")] + #[rstest::rstest] + async fn requests_within_max_are_accepted_with_static_estimated_strategy( + #[values(12.0, 15.0)] max_cost: f64, + ) -> Result<(), BoxError> { + // query total cost is 12.0; max_cost >= 12.0 should result in query being accepted + let demand_control = serde_json::json!({ + "enabled": true, + "mode": "enforce", + "strategy": { + "static_estimated": { + "list_size": 10, + "max": max_cost + } + } + }); + + let response = query_supergraph_service(demand_control).await?; + + let context = response.context; + let body = response.response.into_body().next().await.unwrap(); + + // estimates + assert_eq!(&get_strategy(&context), "static_estimated"); + assert_eq!(&get_result(&context), CODE_OK); + assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); + + assert!(get_result_by_subgraph(&context).is_none()); + assert!(get_estimated_cost_by_subgraph(&context).is_none()); + + // actuals + assert!(body.data.is_some()); + assert!(body.errors.is_empty()); + + let subgraph_call_count = get_subgraph_call_count(&context).unwrap(); + assert_eq!(subgraph_call_count["products"], 1); + + assert_eq!(get_actual_cost(&context).unwrap(), 2.0); + assert!(get_actual_cost_by_subgraph(&context).is_none(),); + + Ok(()) + } + + // Duplicates `requests_exceeding_max_are_rejected_with_static_estimated_by_subgraph_strategy` but + // uses old `static_estimated` strategy. + #[tokio::test(flavor = "multi_thread")] + async fn requests_exceeding_max_are_rejected_with_static_estimated_strategy() + -> Result<(), BoxError> { + let max_cost = 10.0; + + // query total cost is 12.0 for list_size = 10; since `max_cost` value is less than this, + // the response should be an error + let demand_control = serde_json::json!({ + "enabled": true, + "mode": "enforce", + "strategy": { + "static_estimated": { + "list_size": 10, + "max": max_cost + } + } + }); + + let response = query_supergraph_service(demand_control).await?; + + let context = response.context; + let body = response.response.into_body().next().await.unwrap(); + + // estimates + assert_eq!(&get_strategy(&context), "static_estimated"); + assert_eq!(&get_result(&context), CODE_TOO_EXPENSIVE); + assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); + + assert!(get_estimated_cost_by_subgraph(&context).is_none()); + + // actuals + assert!(body.data.is_none()); + + let error = body.errors.iter().find(estimated_too_expensive).unwrap(); + assert_eq!(error.extensions["cost.estimated"], 12.0); + assert_eq!(error.extensions["cost.max"], max_cost); + + assert!(get_subgraph_call_count(&context).is_none()); + + assert!(get_actual_cost(&context).is_none()); + assert!(get_actual_cost_by_subgraph(&context).is_none()); + + Ok(()) + } } mod federated_ships_tests { @@ -555,6 +647,101 @@ mod federated_ships_tests { serde_json::json!({ "vehicles": 9.0 }) ); + Ok(()) + } + + // Duplicates `requests_within_max_are_accepted_with_static_estimated_by_subgraph_strategy` but + // uses old `static_estimated` strategy. + #[tokio::test(flavor = "multi_thread")] + #[rstest::rstest] + async fn requests_within_max_are_accepted_with_static_estimated_strategy( + #[values(10400.0, 10500.0)] max_cost: f64, + ) -> Result<(), BoxError> { + // query total cost is 10400 for list_size = 100; all `max_cost` values are geq than this, + // so the response should be OK + let demand_control = serde_json::json!({ + "enabled": true, + "mode": "enforce", + "strategy": { + "static_estimated": { + "list_size": 100, + "max": max_cost + } + } + }); + + let response = query_supergraph_service(demand_control).await?; + + let context = response.context; + let body = response.response.into_body().next().await.unwrap(); + + // estimates + assert_eq!(&get_strategy(&context), "static_estimated"); + assert_eq!(&get_result(&context), CODE_OK); + assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); + + assert!(get_result_by_subgraph(&context).is_none()); + assert!(get_estimated_cost_by_subgraph(&context).is_none()); + + // actuals + assert!(body.data.is_some()); + assert!(body.errors.is_empty()); + + let subgraph_call_count = get_subgraph_call_count(&context).unwrap(); + assert_eq!(subgraph_call_count["users"], 1); + assert_eq!(subgraph_call_count["vehicles"], 2); + + // NOTE: this differs from the value in the static_estimated_by_subgraph strategy because + // this strategy sums up the result structure, not all the work from the partial responses + assert_eq!(get_actual_cost(&context).unwrap(), 3.0); + assert!(get_actual_cost_by_subgraph(&context).is_none()); + + Ok(()) + } + + // Duplicates `requests_exceeded_max_are_rejected_with_static_estimated_by_subgraph_strategy` but + // uses old `static_estimated` strategy. + #[tokio::test(flavor = "multi_thread")] + async fn requests_exceeding_max_are_rejected_with_static_estimated_strategy() + -> Result<(), BoxError> { + let max_cost = 10000.0; + + // query total cost is 10400 for list_size = 100; since `max_cost` value is less than this, + // the response should be an error + let demand_control = serde_json::json!({ + "enabled": true, + "mode": "enforce", + "strategy": { + "static_estimated": { + "list_size": 100, + "max": max_cost + } + } + }); + + let response = query_supergraph_service(demand_control).await?; + + let context = response.context; + let body = response.response.into_body().next().await.unwrap(); + + // estimates + assert_eq!(&get_strategy(&context), "static_estimated"); + assert_eq!(&get_result(&context), CODE_TOO_EXPENSIVE); + assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); + + assert!(get_estimated_cost_by_subgraph(&context).is_none()); + + // actuals + assert!(body.data.is_none()); + + let error = body.errors.iter().find(estimated_too_expensive).unwrap(); + assert_eq!(error.extensions["cost.estimated"], 10400.0); + assert_eq!(error.extensions["cost.max"], max_cost); + + assert!(get_subgraph_call_count(&context).is_none()); + + assert!(get_actual_cost(&context).is_none()); + assert!(get_actual_cost_by_subgraph(&context).is_none()); Ok(()) } From 1a0b10ab997e09440c5489bd3dc10a5dba9a89ff Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Wed, 21 Jan 2026 14:19:03 -0500 Subject: [PATCH 16/17] fix: make actuals more accurate --- ...nfiguration__tests__schema_generation.snap | 35 ++- .../fixtures/invalid_per_subgraph.yaml | 2 +- .../fixtures/per_subgraph_inheritance.yaml | 2 +- .../fixtures/per_subgraph_no_inheritance.yaml | 2 +- .../src/plugins/demand_control/mod.rs | 60 ++++-- .../plugins/demand_control/strategy/mod.rs | 35 ++- .../strategy/static_estimated.rs | 170 ++++++++++++++- .../strategy/static_estimated_by_subgraph.rs | 202 ----------------- .../telemetry/metrics/apollo/studio.rs | 9 +- apollo-router/src/plugins/telemetry/mod.rs | 4 + .../tests/integration/demand_control.rs | 204 ++++++------------ .../apollo_reports__demand_control_stats.snap | 4 +- ...pollo_reports__demand_control_trace-2.snap | 3 +- .../apollo_reports__demand_control_trace.snap | 3 +- ...ports__demand_control_trace_batched-2.snap | 5 +- ...reports__demand_control_trace_batched.snap | 5 +- 16 files changed, 330 insertions(+), 415 deletions(-) delete mode 100644 apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs diff --git a/apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap b/apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap index 4738826d69..41691dbb59 100644 --- a/apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap +++ b/apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap @@ -115,6 +115,21 @@ expression: "&schema" } ] }, + "ActualCostComputationMode": { + "oneOf": [ + { + "enum": [ + "by_subgraph" + ], + "type": "string" + }, + { + "const": "legacy", + "deprecated": true, + "type": "string" + } + ] + }, "All": { "enum": [ "all" @@ -9128,6 +9143,14 @@ expression: "&schema" "static_estimated": { "additionalProperties": false, "properties": { + "actual_cost_computation_mode": { + "allOf": [ + { + "$ref": "#/definitions/ActualCostComputationMode" + } + ], + "description": "The strategy used to calculate the actual cost incurred by an operation.\n\n* `by_subgraph` (default) computes the cost of each subgraph response and sums them\n to get the total query cost.\n* `legacy` computes the cost based on the final structure of the composed response, not\n including any interim structures from subgraph responses that did not make it to the\n composed response." + }, "list_size": { "description": "The assumed length of lists returned by the operation.", "format": "uint32", @@ -9145,13 +9168,19 @@ expression: "&schema" "$ref": "#/definitions/SubgraphSubgraphStrategyLimitConfiguration" } ], - "description": "Per-subgraph cost control" + "default": { + "all": { + "list_size": null, + "max": null + }, + "subgraphs": {} + }, + "description": "Cost control by subgraph" } }, "required": [ "list_size", - "max", - "subgraphs" + "max" ], "type": "object" } diff --git a/apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml b/apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml index eca6789aac..8ddabcb984 100644 --- a/apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml +++ b/apollo-router/src/plugins/demand_control/fixtures/invalid_per_subgraph.yaml @@ -2,7 +2,7 @@ demand_control: enabled: true mode: enforce strategy: - static_estimated_by_subgraph: + static_estimated: list_size: 1 max: 10 subgraphs: diff --git a/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml index 2bcb7d7de7..57d4acc553 100644 --- a/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml +++ b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_inheritance.yaml @@ -2,7 +2,7 @@ demand_control: enabled: true mode: enforce strategy: - static_estimated_by_subgraph: + static_estimated: list_size: 1 max: 10 subgraphs: diff --git a/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml index 682f78e86b..bebb4fbf7c 100644 --- a/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml +++ b/apollo-router/src/plugins/demand_control/fixtures/per_subgraph_no_inheritance.yaml @@ -2,7 +2,7 @@ demand_control: enabled: true mode: enforce strategy: - static_estimated_by_subgraph: + static_estimated: list_size: 1 max: 10 subgraphs: diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index 00d8d31258..58b03f70cc 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -83,26 +83,17 @@ pub(crate) enum StrategyConfig { list_size: u32, /// The maximum cost of a query max: f64, - }, - /// A simple, statically-defined cost mapping for operations and types. - /// - /// Operation costs: - /// - Mutation: 10 - /// - Query: 0 - /// - Subscription 0 - /// - /// Type costs: - /// - Object: 1 - /// - Interface: 1 - /// - Union: 1 - /// - Scalar: 0 - /// - Enum: 0 - StaticEstimatedBySubgraph { - /// The assumed length of lists returned by the operation. - list_size: u32, - /// The maximum cost of a query - max: f64, + /// The strategy used to calculate the actual cost incurred by an operation. + /// + /// * `by_subgraph` (default) computes the cost of each subgraph response and sums them + /// to get the total query cost. + /// * `legacy` computes the cost based on the final structure of the composed response, not + /// including any interim structures from subgraph responses that did not make it to the + /// composed response. + #[serde(default)] + actual_cost_computation_mode: ActualCostComputationMode, + /// Cost control by subgraph #[serde(default)] subgraphs: SubgraphConfiguration, @@ -115,6 +106,17 @@ pub(crate) enum StrategyConfig { }, } +#[derive(Copy, Clone, Debug, Default, Deserialize, Serialize, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub(crate) enum ActualCostComputationMode { + #[default] + BySubgraph, + + #[deprecated(since = "TBD", note = "use `BySubgraph` instead")] + #[warn(deprecated_in_future)] + Legacy, +} + #[derive(Clone, Default, Debug, Serialize, Deserialize, JsonSchema)] pub(crate) struct SubgraphStrategyLimit { /// The assumed length of lists returned by the operation for this subgraph. @@ -133,11 +135,27 @@ impl StrategyConfig { } #[allow(irrefutable_let_patterns)] - // need to destructure StrategyConfig::StaticEstimatedBySubgraph and ignore the others - let StrategyConfig::StaticEstimatedBySubgraph { subgraphs, .. } = self else { + // need to destructure StrategyConfig::StaticEstimated and ignore StrategyConfig::Test + let StrategyConfig::StaticEstimated { + subgraphs, + actual_cost_computation_mode, + .. + } = self + else { return Ok(()); }; + #[allow(deprecated_in_future)] + if matches!( + actual_cost_computation_mode, + ActualCostComputationMode::Legacy + ) { + tracing::warn!( + "Actual cost computation mode `{}` will be deprecated in the future; migrate to `{}` when possible", + actual_cost_computation_mode + ); + } + if subgraphs.all.max.is_some_and(|s| s < 0.0) { return Err(Error::NegativeQueryCost("all".to_string()).into()); } diff --git a/apollo-router/src/plugins/demand_control/strategy/mod.rs b/apollo-router/src/plugins/demand_control/strategy/mod.rs index e6a636e17e..f9892eeaea 100644 --- a/apollo-router/src/plugins/demand_control/strategy/mod.rs +++ b/apollo-router/src/plugins/demand_control/strategy/mod.rs @@ -6,6 +6,7 @@ use apollo_compiler::ExecutableDocument; use crate::Context; use crate::configuration::subgraph::SubgraphConfiguration; use crate::graphql; +use crate::plugins::demand_control::ActualCostComputationMode; use crate::plugins::demand_control::DemandControlConfig; use crate::plugins::demand_control::DemandControlError; use crate::plugins::demand_control::Mode; @@ -14,12 +15,10 @@ use crate::plugins::demand_control::SubgraphStrategyLimit; use crate::plugins::demand_control::cost_calculator::schema::DemandControlledSchema; use crate::plugins::demand_control::cost_calculator::static_cost::StaticCostCalculator; use crate::plugins::demand_control::strategy::static_estimated::StaticEstimated; -use crate::plugins::demand_control::strategy::static_estimated_by_subgraph::StaticEstimatedBySubgraph; use crate::services::execution; use crate::services::subgraph; mod static_estimated; -mod static_estimated_by_subgraph; #[cfg(test)] mod test; @@ -101,12 +100,13 @@ impl StrategyFactory { // Function extracted for use in tests - allows us to build a `StaticEstimated` directly rather // than a `impl StrategyImpl` - fn create_static_estimated_by_subgraph_strategy( + fn create_static_estimated_strategy( &self, list_size: u32, max: f64, + actual_cost_computation_mode: ActualCostComputationMode, subgraphs: &SubgraphConfiguration, - ) -> StaticEstimatedBySubgraph { + ) -> StaticEstimated { let subgraph_list_sizes = Arc::new(subgraphs.extract(|strategy| strategy.list_size)); let subgraph_maxes = Arc::new(subgraphs.extract(|strategy| strategy.max)); let cost_calculator = StaticCostCalculator::new( @@ -115,34 +115,27 @@ impl StrategyFactory { list_size, subgraph_list_sizes, ); - StaticEstimatedBySubgraph { + StaticEstimated { max, subgraph_maxes, + actual_cost_computation_mode, cost_calculator, } } pub(crate) fn create(&self) -> Strategy { let strategy: Arc = match &self.config.strategy { - StrategyConfig::StaticEstimated { list_size, max } => { - let cost_calculator = StaticCostCalculator::new( - self.supergraph_schema.clone(), - self.subgraph_schemas.clone(), - *list_size, - Arc::new(SubgraphConfiguration::default()), - ); - Arc::new(StaticEstimated { - max: *max, - cost_calculator, - }) - } - StrategyConfig::StaticEstimatedBySubgraph { + StrategyConfig::StaticEstimated { list_size, max, + actual_cost_computation_mode, subgraphs, - } => Arc::new( - self.create_static_estimated_by_subgraph_strategy(*list_size, *max, subgraphs), - ), + } => Arc::new(self.create_static_estimated_strategy( + *list_size, + *max, + *actual_cost_computation_mode, + subgraphs, + )), #[cfg(test)] StrategyConfig::Test { stage, error } => Arc::new(test::Test { stage: stage.clone(), diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index 72226c3dff..3d22ed85bc 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -1,7 +1,12 @@ +use std::sync::Arc; + use apollo_compiler::ExecutableDocument; +use crate::configuration::subgraph::SubgraphConfiguration; use crate::graphql; +use crate::plugins::demand_control::ActualCostComputationMode; use crate::plugins::demand_control::DemandControlError; +use crate::plugins::demand_control::cost_calculator::static_cost::CostBySubgraph; use crate::plugins::demand_control::cost_calculator::static_cost::StaticCostCalculator; use crate::plugins::demand_control::strategy::StrategyImpl; use crate::services::execution; @@ -11,9 +16,17 @@ use crate::services::subgraph; pub(crate) struct StaticEstimated { // The estimated value of the demand pub(crate) max: f64, + pub(crate) subgraph_maxes: Arc>>, + pub(crate) actual_cost_computation_mode: ActualCostComputationMode, pub(crate) cost_calculator: StaticCostCalculator, } +impl StaticEstimated { + fn subgraph_max(&self, subgraph_name: &str) -> Option { + *self.subgraph_maxes.get(subgraph_name) + } +} + impl StrategyImpl for StaticEstimated { fn on_execution_request(&self, request: &execution::Request) -> Result<(), DemandControlError> { self.cost_calculator @@ -27,6 +40,9 @@ impl StrategyImpl for StaticEstimated { .context .insert_cost_strategy("static_estimated".to_string())?; request.context.insert_estimated_cost(cost)?; + request + .context + .insert_estimated_cost_by_subgraph(cost_by_subgraph)?; if cost > self.max { let error = DemandControlError::EstimatedCostTooExpensive { @@ -44,27 +60,84 @@ impl StrategyImpl for StaticEstimated { }) } - fn on_subgraph_request(&self, _request: &subgraph::Request) -> Result<(), DemandControlError> { - Ok(()) + /// Reject subgraph requests when the total subgraph cost exceeds the subgraph max. + fn on_subgraph_request(&self, request: &subgraph::Request) -> Result<(), DemandControlError> { + let subgraph_name = request.subgraph_name.clone(); + let subgraph_cost = request + .context + .get_estimated_cost_for_subgraph(&subgraph_name)?; + + if let Some(subgraph_cost) = subgraph_cost + && let Some(subgraph_max) = self.subgraph_max(&subgraph_name) + && subgraph_cost > subgraph_max + { + let error = DemandControlError::EstimatedSubgraphCostTooExpensive { + subgraph: subgraph_name.clone(), + estimated_cost: subgraph_cost, + max_cost: subgraph_max, + }; + request + .context + .insert_cost_by_subgraph_result(subgraph_name, error.code().to_string())?; + Err(error) + } else { + request + .context + .insert_cost_by_subgraph_result(subgraph_name, "COST_OK".to_string())?; + Ok(()) + } } + /// Determine actual cost for this specific subgraph request fn on_subgraph_response( &self, - _subgraph_name: String, - _request: &ExecutableDocument, - _response: &subgraph::Response, + subgraph_name: String, + request: &ExecutableDocument, + response: &subgraph::Response, ) -> Result<(), DemandControlError> { + if !matches!( + self.actual_cost_computation_mode, + ActualCostComputationMode::BySubgraph + ) { + return Ok(()); + } + + let subgraph_response_body = response.response.body(); + let cost = self.cost_calculator.actual( + request, + subgraph_response_body, + &response + .context + .extensions() + .with_lock(|lock| lock.get().cloned()) + .unwrap_or_default(), + true, + )?; + + response + .context + .update_actual_cost_by_subgraph(CostBySubgraph::new(subgraph_name, cost))?; + Ok(()) } + /// Sum up each subgraph response cost to determine the total cost fn on_execution_response( &self, context: &crate::Context, request: &ExecutableDocument, response: &graphql::Response, ) -> Result<(), DemandControlError> { - if response.data.is_some() { - let cost = self.cost_calculator.actual( + if !response.data.is_some() { + return Ok(()); + } + + let cost = match self.actual_cost_computation_mode { + ActualCostComputationMode::BySubgraph => context + .get_actual_cost_by_subgraph()? + .map_or(0.0, |cost| cost.total()), + #[allow(deprecated_in_future)] + ActualCostComputationMode::Legacy => self.cost_calculator.actual( request, response, &context @@ -72,9 +145,86 @@ impl StrategyImpl for StaticEstimated { .with_lock(|lock| lock.get().cloned()) .unwrap_or_default(), false, - )?; - context.insert_actual_cost(cost)?; - } + )?, + }; + + context.insert_actual_cost(cost)?; Ok(()) } } + +#[cfg(test)] +mod tests { + use tower::BoxError; + + use super::StaticEstimated; + use crate::plugins::demand_control::DemandControl; + use crate::plugins::demand_control::StrategyConfig; + use crate::plugins::test::PluginTestHarness; + + async fn load_config_and_extract_strategy( + config: &'static str, + ) -> Result { + let schema_str = + include_str!("../cost_calculator/fixtures/basic_supergraph_schema.graphql"); + let plugin = PluginTestHarness::::builder() + .config(config) + .schema(schema_str) + .build() + .await?; + + let StrategyConfig::StaticEstimated { + list_size, + max, + actual_cost_computation_mode, + ref subgraphs, + } = plugin.config.strategy + else { + panic!("must provide static_estimated config"); + }; + let strategy = plugin.strategy_factory.create_static_estimated_strategy( + list_size, + max, + actual_cost_computation_mode, + subgraphs, + ); + Ok(strategy) + } + + #[tokio::test] + async fn test_per_subgraph_configuration_inheritance() { + let config = include_str!("../fixtures/per_subgraph_inheritance.yaml"); + + let strategy = load_config_and_extract_strategy(config).await.unwrap(); + assert_eq!(strategy.subgraph_max("reviews").unwrap(), 2.0); + assert_eq!(strategy.subgraph_max("products").unwrap(), 5.0); + assert_eq!(strategy.subgraph_max("users").unwrap(), 5.0); + } + + #[tokio::test] + async fn test_per_subgraph_configuration_no_inheritance() { + let config = include_str!("../fixtures/per_subgraph_no_inheritance.yaml"); + + let strategy = load_config_and_extract_strategy(config).await.unwrap(); + assert_eq!(strategy.subgraph_max("reviews").unwrap(), 2.0); + assert!(strategy.subgraph_max("products").is_none()); + assert!(strategy.subgraph_max("users").is_none()); + } + + #[tokio::test] + async fn test_invalid_per_subgraph_configuration() { + let config = include_str!("../fixtures/invalid_per_subgraph.yaml"); + let strategy_result = load_config_and_extract_strategy(config).await; + + match strategy_result { + Ok(strategy) => { + eprintln!("{:?}", strategy.subgraph_maxes); + panic!("Expected error") + } + Err(err) => assert_eq!( + &err.to_string(), + "Maximum per-subgraph query cost for `products` is negative" + ), + }; + } +} diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs deleted file mode 100644 index 2ce7c11bc0..0000000000 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated_by_subgraph.rs +++ /dev/null @@ -1,202 +0,0 @@ -use std::sync::Arc; - -use apollo_compiler::ExecutableDocument; - -use crate::configuration::subgraph::SubgraphConfiguration; -use crate::graphql; -use crate::plugins::demand_control::DemandControlError; -use crate::plugins::demand_control::cost_calculator::static_cost::CostBySubgraph; -use crate::plugins::demand_control::cost_calculator::static_cost::StaticCostCalculator; -use crate::plugins::demand_control::strategy::StrategyImpl; -use crate::services::execution; -use crate::services::subgraph; - -/// This strategy will reject requests if the estimated cost of the request exceeds the maximum cost. -pub(crate) struct StaticEstimatedBySubgraph { - // The estimated value of the demand - pub(crate) max: f64, - pub(crate) subgraph_maxes: Arc>>, - pub(crate) cost_calculator: StaticCostCalculator, -} - -impl StaticEstimatedBySubgraph { - fn subgraph_max(&self, subgraph_name: &str) -> Option { - *self.subgraph_maxes.get(subgraph_name) - } -} - -impl StrategyImpl for StaticEstimatedBySubgraph { - fn on_execution_request(&self, request: &execution::Request) -> Result<(), DemandControlError> { - self.cost_calculator - .planned( - &request.query_plan, - &request.supergraph_request.body().variables, - ) - .and_then(|cost_by_subgraph| { - let cost = cost_by_subgraph.total(); - request - .context - .insert_cost_strategy("static_estimated_by_subgraph".to_string())?; - request.context.insert_estimated_cost(cost)?; - request - .context - .insert_estimated_cost_by_subgraph(cost_by_subgraph)?; - - if cost > self.max { - let error = DemandControlError::EstimatedCostTooExpensive { - estimated_cost: cost, - max_cost: self.max, - }; - request - .context - .insert_cost_result(error.code().to_string())?; - Err(error) - } else { - request.context.insert_cost_result("COST_OK".to_string())?; - Ok(()) - } - }) - } - - /// Reject subgraph requests when the total subgraph cost exceeds the subgraph max. - fn on_subgraph_request(&self, request: &subgraph::Request) -> Result<(), DemandControlError> { - let subgraph_name = request.subgraph_name.clone(); - let subgraph_cost = request - .context - .get_estimated_cost_for_subgraph(&subgraph_name)?; - - if let Some(subgraph_cost) = subgraph_cost - && let Some(subgraph_max) = self.subgraph_max(&subgraph_name) - && subgraph_cost > subgraph_max - { - let error = DemandControlError::EstimatedSubgraphCostTooExpensive { - subgraph: subgraph_name.clone(), - estimated_cost: subgraph_cost, - max_cost: subgraph_max, - }; - request - .context - .insert_cost_by_subgraph_result(subgraph_name, error.code().to_string())?; - Err(error) - } else { - request - .context - .insert_cost_by_subgraph_result(subgraph_name, "COST_OK".to_string())?; - Ok(()) - } - } - - /// Parse response to measure actual cost for subgraph. - fn on_subgraph_response( - &self, - subgraph_name: String, - request: &ExecutableDocument, - response: &subgraph::Response, - ) -> Result<(), DemandControlError> { - let subgraph_response_body = response.response.body(); - let cost = self.cost_calculator.actual( - request, - subgraph_response_body, - &response - .context - .extensions() - .with_lock(|lock| lock.get().cloned()) - .unwrap_or_default(), - true, - )?; - - response - .context - .update_actual_cost_by_subgraph(CostBySubgraph::new(subgraph_name, cost))?; - - Ok(()) - } - - fn on_execution_response( - &self, - context: &crate::Context, - _request: &ExecutableDocument, - response: &graphql::Response, - ) -> Result<(), DemandControlError> { - // sum up values from subgraph responses - if response.data.is_some() { - let cost = context - .get_actual_cost_by_subgraph()? - .map_or(0.0, |cost| cost.total()); - context.insert_actual_cost(cost)?; - } - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use tower::BoxError; - - use super::StaticEstimatedBySubgraph; - use crate::plugins::demand_control::DemandControl; - use crate::plugins::demand_control::StrategyConfig; - use crate::plugins::test::PluginTestHarness; - - async fn load_config_and_extract_strategy( - config: &'static str, - ) -> Result { - let schema_str = - include_str!("../cost_calculator/fixtures/basic_supergraph_schema.graphql"); - let plugin = PluginTestHarness::::builder() - .config(config) - .schema(schema_str) - .build() - .await?; - - let StrategyConfig::StaticEstimatedBySubgraph { - list_size, - max, - ref subgraphs, - } = plugin.config.strategy - else { - panic!("must provide static_estimated_by_subgraph config"); - }; - let strategy = plugin - .strategy_factory - .create_static_estimated_by_subgraph_strategy(list_size, max, subgraphs); - Ok(strategy) - } - - #[tokio::test] - async fn test_per_subgraph_configuration_inheritance() { - let config = include_str!("../fixtures/per_subgraph_inheritance.yaml"); - - let strategy = load_config_and_extract_strategy(config).await.unwrap(); - assert_eq!(strategy.subgraph_max("reviews").unwrap(), 2.0); - assert_eq!(strategy.subgraph_max("products").unwrap(), 5.0); - assert_eq!(strategy.subgraph_max("users").unwrap(), 5.0); - } - - #[tokio::test] - async fn test_per_subgraph_configuration_no_inheritance() { - let config = include_str!("../fixtures/per_subgraph_no_inheritance.yaml"); - - let strategy = load_config_and_extract_strategy(config).await.unwrap(); - assert_eq!(strategy.subgraph_max("reviews").unwrap(), 2.0); - assert!(strategy.subgraph_max("products").is_none()); - assert!(strategy.subgraph_max("users").is_none()); - } - - #[tokio::test] - async fn test_invalid_per_subgraph_configuration() { - let config = include_str!("../fixtures/invalid_per_subgraph.yaml"); - let strategy_result = load_config_and_extract_strategy(config).await; - - match strategy_result { - Ok(strategy) => { - eprintln!("{:?}", strategy.subgraph_maxes); - panic!("Expected error") - } - Err(err) => assert_eq!( - &err.to_string(), - "Maximum per-subgraph query cost for `products` is negative" - ), - }; - } -} diff --git a/apollo-router/src/plugins/telemetry/metrics/apollo/studio.rs b/apollo-router/src/plugins/telemetry/metrics/apollo/studio.rs index eb4769e46c..f8dd27f33e 100644 --- a/apollo-router/src/plugins/telemetry/metrics/apollo/studio.rs +++ b/apollo-router/src/plugins/telemetry/metrics/apollo/studio.rs @@ -374,8 +374,9 @@ impl From for LimitsStats { pub(crate) struct SingleLimitsStats { pub(crate) strategy: Option, pub(crate) cost_estimated: Option, - pub(crate) cost_by_subgraph_estimated: Option, pub(crate) cost_actual: Option, + pub(crate) cost_by_subgraph_estimated: Option, + pub(crate) cost_by_subgraph_actual: Option, pub(crate) depth: u64, pub(crate) height: u64, pub(crate) alias_count: u64, @@ -623,11 +624,15 @@ mod test { limits_stats: SingleLimitsStats { strategy: Some("test".to_string()), cost_estimated: Some(10.0), + cost_actual: Some(7.0), cost_by_subgraph_estimated: Some(CostBySubgraph::new( "products".to_string(), 10.0, )), - cost_actual: Some(7.0), + cost_by_subgraph_actual: Some(CostBySubgraph::new( + "products".to_string(), + 7.0, + )), depth: 2, height: 4, alias_count: 0, diff --git a/apollo-router/src/plugins/telemetry/mod.rs b/apollo-router/src/plugins/telemetry/mod.rs index 73e39ee634..05da9767fb 100644 --- a/apollo-router/src/plugins/telemetry/mod.rs +++ b/apollo-router/src/plugins/telemetry/mod.rs @@ -1479,6 +1479,10 @@ impl Telemetry { .get_estimated_cost_by_subgraph() .ok() .flatten(), + cost_by_subgraph_actual: context + .get_actual_cost_by_subgraph() + .ok() + .flatten(), // These limits are related to the Traffic Shaping feature, unrelated to the Demand Control plugin depth: query_limits.map_or(0, |ql| ql.depth as u64), diff --git a/apollo-router/tests/integration/demand_control.rs b/apollo-router/tests/integration/demand_control.rs index 6a07050f08..5c4b0a919b 100644 --- a/apollo-router/tests/integration/demand_control.rs +++ b/apollo-router/tests/integration/demand_control.rs @@ -142,7 +142,7 @@ mod basic_fragments_tests { #[tokio::test(flavor = "multi_thread")] #[rstest::rstest] - async fn requests_within_max_are_accepted_with_static_estimated_by_subgraph_strategy( + async fn requests_within_max_are_accepted( #[values(12.0, 15.0)] max_cost: f64, ) -> Result<(), BoxError> { // query total cost is 12.0; max_cost >= 12.0 should result in query being accepted @@ -150,7 +150,7 @@ mod basic_fragments_tests { "enabled": true, "mode": "enforce", "strategy": { - "static_estimated_by_subgraph": { + "static_estimated": { "list_size": 10, "max": max_cost } @@ -163,7 +163,7 @@ mod basic_fragments_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_strategy(&context), "static_estimated"); assert_eq!(&get_result(&context), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); @@ -193,8 +193,7 @@ mod basic_fragments_tests { } #[tokio::test(flavor = "multi_thread")] - async fn requests_exceeding_max_are_rejected_with_static_estimated_by_subgraph_strategy() - -> Result<(), BoxError> { + async fn requests_exceeding_max_are_rejected() -> Result<(), BoxError> { let max_cost = 10.0; // query total cost is 12.0 for list_size = 10; since `max_cost` value is less than this, @@ -203,7 +202,7 @@ mod basic_fragments_tests { "enabled": true, "mode": "enforce", "strategy": { - "static_estimated_by_subgraph": { + "static_estimated": { "list_size": 10, "max": max_cost } @@ -216,7 +215,7 @@ mod basic_fragments_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_strategy(&context), "static_estimated"); assert_eq!(&get_result(&context), CODE_TOO_EXPENSIVE); assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); @@ -241,15 +240,14 @@ mod basic_fragments_tests { } #[tokio::test] - async fn requests_which_exceed_subgraph_limit_are_partially_accepted_with_static_estimated_by_subgraph_strategy() - -> Result<(), BoxError> { + async fn requests_which_exceed_subgraph_limit_are_partially_accepted() -> Result<(), BoxError> { // query checks products once; query should be accepted based on max but products subgraph // should not be called. let demand_control = serde_json::json!({ "enabled": true, "mode": "enforce", "strategy": { - "static_estimated_by_subgraph": { + "static_estimated": { "list_size": 10, "max": 15, "subgraphs": { @@ -270,7 +268,7 @@ mod basic_fragments_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_strategy(&context), "static_estimated"); assert_eq!(&get_result(&context), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); assert_eq!( @@ -303,12 +301,14 @@ mod basic_fragments_tests { Ok(()) } - // Duplicates `requests_within_max_are_accepted_with_static_estimated_by_subgraph_strategy` but - // uses old `static_estimated` strategy. #[tokio::test(flavor = "multi_thread")] #[rstest::rstest] - async fn requests_within_max_are_accepted_with_static_estimated_strategy( - #[values(12.0, 15.0)] max_cost: f64, + #[case::new_cost_mode("by_subgraph", 2.0, true)] + #[case::legacy_cost_mode("legacy", 2.0, false)] + async fn actual_cost_computation_can_vary_based_on_mode( + #[case] computation_mode: &str, + #[case] expected_actual_cost: f64, + #[case] should_have_per_subgraph_actuals: bool, ) -> Result<(), BoxError> { // query total cost is 12.0; max_cost >= 12.0 should result in query being accepted let demand_control = serde_json::json!({ @@ -317,7 +317,8 @@ mod basic_fragments_tests { "strategy": { "static_estimated": { "list_size": 10, - "max": max_cost + "actual_cost_computation_mode": computation_mode, + "max": 20.0 } } }); @@ -332,8 +333,14 @@ mod basic_fragments_tests { assert_eq!(&get_result(&context), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); - assert!(get_result_by_subgraph(&context).is_none()); - assert!(get_estimated_cost_by_subgraph(&context).is_none()); + assert_eq!( + get_result_by_subgraph(&context).unwrap(), + serde_json::json!({ "products": CODE_OK }) + ); + assert_eq!( + get_estimated_cost_by_subgraph(&context).unwrap(), + serde_json::json!({ "products": 12.0 }) + ); // actuals assert!(body.data.is_some()); @@ -342,55 +349,11 @@ mod basic_fragments_tests { let subgraph_call_count = get_subgraph_call_count(&context).unwrap(); assert_eq!(subgraph_call_count["products"], 1); - assert_eq!(get_actual_cost(&context).unwrap(), 2.0); - assert!(get_actual_cost_by_subgraph(&context).is_none(),); - - Ok(()) - } - - // Duplicates `requests_exceeding_max_are_rejected_with_static_estimated_by_subgraph_strategy` but - // uses old `static_estimated` strategy. - #[tokio::test(flavor = "multi_thread")] - async fn requests_exceeding_max_are_rejected_with_static_estimated_strategy() - -> Result<(), BoxError> { - let max_cost = 10.0; - - // query total cost is 12.0 for list_size = 10; since `max_cost` value is less than this, - // the response should be an error - let demand_control = serde_json::json!({ - "enabled": true, - "mode": "enforce", - "strategy": { - "static_estimated": { - "list_size": 10, - "max": max_cost - } - } - }); - - let response = query_supergraph_service(demand_control).await?; - - let context = response.context; - let body = response.response.into_body().next().await.unwrap(); - - // estimates - assert_eq!(&get_strategy(&context), "static_estimated"); - assert_eq!(&get_result(&context), CODE_TOO_EXPENSIVE); - assert_eq!(get_estimated_cost(&context).unwrap(), 12.0); - - assert!(get_estimated_cost_by_subgraph(&context).is_none()); - - // actuals - assert!(body.data.is_none()); - - let error = body.errors.iter().find(estimated_too_expensive).unwrap(); - assert_eq!(error.extensions["cost.estimated"], 12.0); - assert_eq!(error.extensions["cost.max"], max_cost); - - assert!(get_subgraph_call_count(&context).is_none()); - - assert!(get_actual_cost(&context).is_none()); - assert!(get_actual_cost_by_subgraph(&context).is_none()); + assert_eq!(get_actual_cost(&context).unwrap(), expected_actual_cost); + assert_eq!( + get_actual_cost_by_subgraph(&context).is_some(), + should_have_per_subgraph_actuals + ); Ok(()) } @@ -482,7 +445,7 @@ mod federated_ships_tests { #[tokio::test(flavor = "multi_thread")] #[rstest::rstest] - async fn requests_within_max_are_accepted_with_static_estimated_by_subgraph_strategy( + async fn requests_within_max_are_accepted( #[values(10400.0, 10500.0)] max_cost: f64, ) -> Result<(), BoxError> { // query total cost is 10400 for list_size = 100; all `max_cost` values are geq than this, @@ -491,7 +454,7 @@ mod federated_ships_tests { "enabled": true, "mode": "enforce", "strategy": { - "static_estimated_by_subgraph": { + "static_estimated": { "list_size": 100, "max": max_cost } @@ -504,7 +467,7 @@ mod federated_ships_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_strategy(&context), "static_estimated"); assert_eq!(&get_result(&context), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); @@ -535,8 +498,7 @@ mod federated_ships_tests { } #[tokio::test(flavor = "multi_thread")] - async fn requests_exceeding_max_are_rejected_with_static_estimated_by_subgraph_strategy() - -> Result<(), BoxError> { + async fn requests_exceeding_max_are_rejected() -> Result<(), BoxError> { let max_cost = 10000.0; // query total cost is 10400 for list_size = 100; since `max_cost` value is less than this, @@ -545,7 +507,7 @@ mod federated_ships_tests { "enabled": true, "mode": "enforce", "strategy": { - "static_estimated_by_subgraph": { + "static_estimated": { "list_size": 100, "max": max_cost } @@ -558,7 +520,7 @@ mod federated_ships_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_strategy(&context), "static_estimated"); assert_eq!(&get_result(&context), CODE_TOO_EXPENSIVE); assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); @@ -583,8 +545,7 @@ mod federated_ships_tests { } #[tokio::test] - async fn requests_which_exceed_subgraph_limit_are_partially_accepted_with_static_estimated_by_subgraph_strategy() - -> Result<(), BoxError> { + async fn requests_which_exceed_subgraph_limit_are_partially_accepted() -> Result<(), BoxError> { // query checks vehicles, then users, then vehicles. // interrupting the users check via a demand control limit should still permit both vehicles // checks. @@ -592,7 +553,7 @@ mod federated_ships_tests { "enabled": true, "mode": "enforce", "strategy": { - "static_estimated_by_subgraph": { + "static_estimated": { "list_size": 100, "max": 15000.0, "subgraphs": { @@ -613,7 +574,7 @@ mod federated_ships_tests { let body = response.response.into_body().next().await.unwrap(); // estimates - assert_eq!(&get_strategy(&context), "static_estimated_by_subgraph"); + assert_eq!(&get_strategy(&context), "static_estimated"); assert_eq!(&get_result(&context), CODE_OK); assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); assert_eq!( @@ -650,22 +611,25 @@ mod federated_ships_tests { Ok(()) } - // Duplicates `requests_within_max_are_accepted_with_static_estimated_by_subgraph_strategy` but - // uses old `static_estimated` strategy. #[tokio::test(flavor = "multi_thread")] #[rstest::rstest] - async fn requests_within_max_are_accepted_with_static_estimated_strategy( - #[values(10400.0, 10500.0)] max_cost: f64, + #[case::new_cost_mode("by_subgraph", 15.0, true)] + #[case::legacy_cost_mode("legacy", 3.0, false)] + async fn actual_cost_computation_can_vary_based_on_mode( + #[case] computation_mode: &str, + #[case] expected_actual_cost: f64, + #[case] should_have_per_subgraph_actuals: bool, ) -> Result<(), BoxError> { - // query total cost is 10400 for list_size = 100; all `max_cost` values are geq than this, - // so the response should be OK + // estimated cost is 140 because list_size is 10, in contrast to tests above which use + // list_size 100 and therefore have estimated cost 10400 let demand_control = serde_json::json!({ "enabled": true, "mode": "enforce", "strategy": { "static_estimated": { - "list_size": 100, - "max": max_cost + "list_size": 10, + "actual_cost_computation_mode": computation_mode, + "max": 150.0 } } }); @@ -678,10 +642,16 @@ mod federated_ships_tests { // estimates assert_eq!(&get_strategy(&context), "static_estimated"); assert_eq!(&get_result(&context), CODE_OK); - assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); + assert_eq!(get_estimated_cost(&context).unwrap(), 140.0); - assert!(get_result_by_subgraph(&context).is_none()); - assert!(get_estimated_cost_by_subgraph(&context).is_none()); + assert_eq!( + get_result_by_subgraph(&context).unwrap(), + serde_json::json!({ "users": CODE_OK, "vehicles": CODE_OK }) + ); + assert_eq!( + get_estimated_cost_by_subgraph(&context).unwrap(), + serde_json::json!({ "users": 110.0, "vehicles": 30.0 }) + ); // actuals assert!(body.data.is_some()); @@ -691,57 +661,11 @@ mod federated_ships_tests { assert_eq!(subgraph_call_count["users"], 1); assert_eq!(subgraph_call_count["vehicles"], 2); - // NOTE: this differs from the value in the static_estimated_by_subgraph strategy because - // this strategy sums up the result structure, not all the work from the partial responses - assert_eq!(get_actual_cost(&context).unwrap(), 3.0); - assert!(get_actual_cost_by_subgraph(&context).is_none()); - - Ok(()) - } - - // Duplicates `requests_exceeded_max_are_rejected_with_static_estimated_by_subgraph_strategy` but - // uses old `static_estimated` strategy. - #[tokio::test(flavor = "multi_thread")] - async fn requests_exceeding_max_are_rejected_with_static_estimated_strategy() - -> Result<(), BoxError> { - let max_cost = 10000.0; - - // query total cost is 10400 for list_size = 100; since `max_cost` value is less than this, - // the response should be an error - let demand_control = serde_json::json!({ - "enabled": true, - "mode": "enforce", - "strategy": { - "static_estimated": { - "list_size": 100, - "max": max_cost - } - } - }); - - let response = query_supergraph_service(demand_control).await?; - - let context = response.context; - let body = response.response.into_body().next().await.unwrap(); - - // estimates - assert_eq!(&get_strategy(&context), "static_estimated"); - assert_eq!(&get_result(&context), CODE_TOO_EXPENSIVE); - assert_eq!(get_estimated_cost(&context).unwrap(), 10400.0); - - assert!(get_estimated_cost_by_subgraph(&context).is_none()); - - // actuals - assert!(body.data.is_none()); - - let error = body.errors.iter().find(estimated_too_expensive).unwrap(); - assert_eq!(error.extensions["cost.estimated"], 10400.0); - assert_eq!(error.extensions["cost.max"], max_cost); - - assert!(get_subgraph_call_count(&context).is_none()); - - assert!(get_actual_cost(&context).is_none()); - assert!(get_actual_cost_by_subgraph(&context).is_none()); + assert_eq!(get_actual_cost(&context).unwrap(), expected_actual_cost); + assert_eq!( + get_actual_cost_by_subgraph(&context).is_some(), + should_have_per_subgraph_actuals + ); Ok(()) } diff --git a/apollo-router/tests/snapshots/apollo_reports__demand_control_stats.snap b/apollo-router/tests/snapshots/apollo_reports__demand_control_stats.snap index a581f9baf7..628bf2eba5 100644 --- a/apollo-router/tests/snapshots/apollo_reports__demand_control_stats.snap +++ b/apollo-router/tests/snapshots/apollo_reports__demand_control_stats.snap @@ -1,7 +1,6 @@ --- source: apollo-router/tests/apollo_reports.rs expression: report -snapshot_kind: text --- header: graph_ref: test @@ -203,9 +202,8 @@ traces_per_query: - 0 - 0 - 0 - - 0 - 1 - max_cost_actual: 20 + max_cost_actual: 18 depth: 4 height: 7 alias_count: 0 diff --git a/apollo-router/tests/snapshots/apollo_reports__demand_control_trace-2.snap b/apollo-router/tests/snapshots/apollo_reports__demand_control_trace-2.snap index 2419a5d911..c83e2c5b02 100644 --- a/apollo-router/tests/snapshots/apollo_reports__demand_control_trace-2.snap +++ b/apollo-router/tests/snapshots/apollo_reports__demand_control_trace-2.snap @@ -1,7 +1,6 @@ --- source: apollo-router/tests/apollo_reports.rs expression: report -snapshot_kind: text --- header: graph_ref: test @@ -605,7 +604,7 @@ traces_per_query: result: COST_OK strategy: static_estimated cost_estimated: 230 - cost_actual: 20 + cost_actual: 18 depth: 4 height: 7 alias_count: 0 diff --git a/apollo-router/tests/snapshots/apollo_reports__demand_control_trace.snap b/apollo-router/tests/snapshots/apollo_reports__demand_control_trace.snap index 2419a5d911..c83e2c5b02 100644 --- a/apollo-router/tests/snapshots/apollo_reports__demand_control_trace.snap +++ b/apollo-router/tests/snapshots/apollo_reports__demand_control_trace.snap @@ -1,7 +1,6 @@ --- source: apollo-router/tests/apollo_reports.rs expression: report -snapshot_kind: text --- header: graph_ref: test @@ -605,7 +604,7 @@ traces_per_query: result: COST_OK strategy: static_estimated cost_estimated: 230 - cost_actual: 20 + cost_actual: 18 depth: 4 height: 7 alias_count: 0 diff --git a/apollo-router/tests/snapshots/apollo_reports__demand_control_trace_batched-2.snap b/apollo-router/tests/snapshots/apollo_reports__demand_control_trace_batched-2.snap index 9eaa200d72..ee75edaba1 100644 --- a/apollo-router/tests/snapshots/apollo_reports__demand_control_trace_batched-2.snap +++ b/apollo-router/tests/snapshots/apollo_reports__demand_control_trace_batched-2.snap @@ -1,7 +1,6 @@ --- source: apollo-router/tests/apollo_reports.rs expression: report -snapshot_kind: text --- header: graph_ref: test @@ -605,7 +604,7 @@ traces_per_query: result: COST_OK strategy: static_estimated cost_estimated: 230 - cost_actual: 20 + cost_actual: 18 depth: 4 height: 7 alias_count: 0 @@ -1205,7 +1204,7 @@ traces_per_query: result: COST_OK strategy: static_estimated cost_estimated: 230 - cost_actual: 20 + cost_actual: 18 depth: 4 height: 7 alias_count: 0 diff --git a/apollo-router/tests/snapshots/apollo_reports__demand_control_trace_batched.snap b/apollo-router/tests/snapshots/apollo_reports__demand_control_trace_batched.snap index 9eaa200d72..ee75edaba1 100644 --- a/apollo-router/tests/snapshots/apollo_reports__demand_control_trace_batched.snap +++ b/apollo-router/tests/snapshots/apollo_reports__demand_control_trace_batched.snap @@ -1,7 +1,6 @@ --- source: apollo-router/tests/apollo_reports.rs expression: report -snapshot_kind: text --- header: graph_ref: test @@ -605,7 +604,7 @@ traces_per_query: result: COST_OK strategy: static_estimated cost_estimated: 230 - cost_actual: 20 + cost_actual: 18 depth: 4 height: 7 alias_count: 0 @@ -1205,7 +1204,7 @@ traces_per_query: result: COST_OK strategy: static_estimated cost_estimated: 230 - cost_actual: 20 + cost_actual: 18 depth: 4 height: 7 alias_count: 0 From 4c2f10255a8eed4312cd97856c7defceeb236fd9 Mon Sep 17 00:00:00 2001 From: carodewig <16093297+carodewig@users.noreply.github.com> Date: Wed, 21 Jan 2026 14:43:53 -0500 Subject: [PATCH 17/17] lint: fix errors from early commit --- apollo-router/src/plugins/demand_control/mod.rs | 5 ++--- .../src/plugins/demand_control/strategy/static_estimated.rs | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/apollo-router/src/plugins/demand_control/mod.rs b/apollo-router/src/plugins/demand_control/mod.rs index 58b03f70cc..2a2899b1a6 100644 --- a/apollo-router/src/plugins/demand_control/mod.rs +++ b/apollo-router/src/plugins/demand_control/mod.rs @@ -106,7 +106,7 @@ pub(crate) enum StrategyConfig { }, } -#[derive(Copy, Clone, Debug, Default, Deserialize, Serialize, JsonSchema)] +#[derive(Copy, Clone, Debug, Default, Deserialize, JsonSchema)] #[serde(rename_all = "snake_case")] pub(crate) enum ActualCostComputationMode { #[default] @@ -151,8 +151,7 @@ impl StrategyConfig { ActualCostComputationMode::Legacy ) { tracing::warn!( - "Actual cost computation mode `{}` will be deprecated in the future; migrate to `{}` when possible", - actual_cost_computation_mode + "Actual cost computation mode `legacy` will be deprecated in the future; migrate to `by_subgraph` when possible", ); } diff --git a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs index 3d22ed85bc..c3412cfafe 100644 --- a/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs +++ b/apollo-router/src/plugins/demand_control/strategy/static_estimated.rs @@ -128,7 +128,7 @@ impl StrategyImpl for StaticEstimated { request: &ExecutableDocument, response: &graphql::Response, ) -> Result<(), DemandControlError> { - if !response.data.is_some() { + if response.data.is_none() { return Ok(()); }