fix: demand control actual costs should consider each subgraph fetch (#8827)

Author: carodewig

.changesets/fix_caroline_demand_control_actuals.md

@@ -0,0 +1,28 @@
+### Demand control actual costs should consider each subgraph fetch ([PR #8827](https://github.com/apollographql/router/pull/8827))
+
+The demand control feature estimates query costs by summing together the cost of each subgraph operation. This allows it
+to capture any intermediate work that must be completed to return a complete response.
+
+Prior to this version, the actual query cost computation only considered the final response shape; it did not include
+any of the intermediate work done in its total.
+
+This version fixes that behavior to compute the actual query cost as the sum of all subgraph response costs. This more
+accurately reflects the work done per operation and allows a more meaningful comparison
+between actual and estimated costs.
+
+Note: if you would like to disable the new actual cost computation behavior, you should set the router configuration
+option `demand_control.strategy.static_estimated.actual_cost_mode` to `response_shape`.
+
+```yaml
+demand_control:
+  enabled: true
+  mode: enforce
+  strategy:
+    static_estimated:
+      max: 10
+      list_size: 10
+      actual_cost_mode: by_subgraph # the default value
+      # actual_cost_mode: response_shape # revert to prior actual cost computation mode
+```
+
+By [@carodewig](https://github.com/carodewig) in https://github.com/apollographql/router/pull/8827

.config/nextest.toml

@@ -67,6 +67,7 @@ or ( binary_id(=apollo-router::integration_tests) & test(=integration::connector
 or ( binary_id(=apollo-router::integration_tests) & test(=integration::connectors::authentication::test_aws_sig_v4_signing) )
 or ( binary_id(=apollo-router::integration_tests) & test(=integration::coprocessor::test_coprocessor_response_handling) )
 or ( binary_id(=apollo-router::integration_tests) & test(=integration::coprocessor::test_error_not_propagated_to_client) )
+or ( binary_id(=apollo-router::integration_tests) & test(=integration::entity_cache::test_cache_metrics) )
 or ( binary_id(=apollo-router::integration_tests) & test(=integration::file_upload::it_fails_incompatible_query_order) )
 or ( binary_id(=apollo-router::integration_tests) & test(=integration::file_upload::it_fails_invalid_file_order) )
 or ( binary_id(=apollo-router::integration_tests) & test(=integration::file_upload::it_fails_invalid_multipart_order) )

apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap

@@ -115,6 +115,21 @@ expression: "&schema"
         }
       ]
     },
+    "ActualCostMode": {
+      "oneOf": [
+        {
+          "const": "by_subgraph",
+          "description": "Computes the cost of each subgraph response and sums them to get the total query cost.",
+          "type": "string"
+        },
+        {
+          "const": "by_response_shape",
+          "deprecated": true,
+          "description": "Computes the cost based on the final structure of the composed response, not including any\ninterim structures from subgraph responses that did not make it to the composed response.",
+          "type": "string"
+        }
+      ]
+    },
     "All": {
       "enum": [
         "all"
@@ -9137,6 +9152,14 @@ expression: "&schema"
             "static_estimated": {
               "additionalProperties": false,
               "properties": {
+                "actual_cost_mode": {
+                  "allOf": [
+                    {
+                      "$ref": "#/definitions/ActualCostMode"
+                    }
+                  ],
+                  "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* `by_response_shape` computes the cost based on the final structure of the composed\n  response, not including any interim structures from subgraph responses that did not\n  make it to the composed response."
+                },
                 "list_size": {
                   "description": "The assumed length of lists returned by the operation.",
                   "format": "uint32",

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

@@ -2,4 +2,22 @@ mod directives;
 pub(in crate::plugins::demand_control) mod schema;
 pub(crate) mod static_cost;
 
+use std::collections::HashMap;
+
 use crate::plugins::demand_control::DemandControlError;
+
+#[derive(Clone, Default, Debug, serde::Serialize, serde::Deserialize)]
+pub(crate) struct CostBySubgraph(HashMap<String, f64>);
+impl CostBySubgraph {
+    pub(crate) fn add_or_insert(&mut self, subgraph: &str, value: f64) {
+        if let Some(subgraph_cost) = self.0.get_mut(subgraph) {
+            *subgraph_cost += value;
+        } else {
+            self.0.insert(subgraph.to_string(), value);
+        }
+    }
+
+    pub(crate) fn total(&self) -> f64 {
+        self.0.values().sum()
+    }
+}

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

@@ -553,53 +553,57 @@ 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);
+
+        // We need to have a field definition for later processing, unless the query is an
+        // `_entities` query. If the field should be there and isn't, return now.
+        let is_entities_query = parent_ty == "Query" && field.name == "_entities";
+        if definition.is_none() && !is_entities_query {
+            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;
-                        }
-                    } 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,
-                        )
+        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,
+                    )
                 }
             }
-        } else {
-            tracing::debug!(
-                "Failed to get schema definition for field {}.{}. The resulting response cost will be a partial result.",
-                parent_ty,
-                field.name,
-            )
         }
     }
 }

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

@@ -34,6 +34,7 @@ use crate::json_ext::Object;
 use crate::layers::ServiceBuilderExt;
 use crate::plugin::Plugin;
 use crate::plugin::PluginInit;
+use crate::plugins::demand_control::cost_calculator::CostBySubgraph;
 use crate::plugins::demand_control::cost_calculator::schema::DemandControlledSchema;
 use crate::plugins::demand_control::strategy::Strategy;
 use crate::plugins::demand_control::strategy::StrategyFactory;
@@ -51,6 +52,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_ACTUAL_KEY: &str =
+    "apollo::demand_control::actual_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")]
@@ -73,6 +77,16 @@ pub(crate) enum StrategyConfig {
         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.
+        /// * `by_response_shape` 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_mode: ActualCostMode,
     },
 
     #[cfg(test)]
@@ -82,6 +96,41 @@ pub(crate) enum StrategyConfig {
     },
 }
 
+#[derive(Copy, Clone, Debug, Default, Deserialize, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub(crate) enum ActualCostMode {
+    /// Computes the cost of each subgraph response and sums them to get the total query cost.
+    #[default]
+    BySubgraph,
+
+    /// 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.
+    #[deprecated(since = "TBD", note = "use `BySubgraph` instead")]
+    #[warn(deprecated_in_future)]
+    ByResponseShape,
+}
+
+impl StrategyConfig {
+    fn validate(&self) -> Result<(), BoxError> {
+        let actual_cost_mode = match self {
+            StrategyConfig::StaticEstimated {
+                actual_cost_mode, ..
+            } => actual_cost_mode,
+            #[cfg(test)]
+            StrategyConfig::Test { .. } => return Ok(()),
+        };
+
+        #[allow(deprecated_in_future)]
+        if matches!(actual_cost_mode, ActualCostMode::ByResponseShape) {
+            tracing::warn!(
+                "Actual cost computation mode `by_response_shape` will be deprecated in the future; migrate to `by_subgraph` when possible",
+            );
+        }
+
+        Ok(())
+    }
+}
+
 #[derive(Copy, Clone, Debug, Serialize, Deserialize, JsonSchema, Eq, PartialEq)]
 #[serde(deny_unknown_fields, rename_all = "snake_case")]
 pub(crate) enum Mode {
@@ -268,6 +317,30 @@ impl Context {
         Ok(estimated.zip(actual).map(|(est, act)| est - act))
     }
 
+    pub(crate) fn get_actual_cost_by_subgraph(
+        &self,
+    ) -> Result<Option<CostBySubgraph>, DemandControlError> {
+        self.get::<&str, CostBySubgraph>(COST_BY_SUBGRAPH_ACTUAL_KEY)
+            .map_err(|e| DemandControlError::ContextSerializationError(e.to_string()))
+    }
+
+    pub(crate) fn update_actual_cost_by_subgraph(
+        &self,
+        subgraph: &str,
+        cost: f64,
+    ) -> 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.add_or_insert(subgraph, cost);
+                existing_cost
+            },
+        )
+        .map_err(|e| DemandControlError::ContextSerializationError(e.to_string()))?;
+        Ok(())
+    }
+
     pub(crate) fn insert_cost_result(&self, result: String) -> Result<(), DemandControlError> {
         self.insert(COST_RESULT_KEY, result)
             .map_err(|e| DemandControlError::ContextSerializationError(e.to_string()))?;
@@ -348,6 +421,8 @@ impl Plugin for DemandControl {
                 .insert(subgraph_name.clone(), demand_controlled_subgraph_schema);
         }
 
+        init.config.strategy.validate()?;
+
         Ok(DemandControl {
             strategy_factory: StrategyFactory::new(
                 init.config.clone(),
@@ -501,7 +576,7 @@ impl Plugin for DemandControl {
                     |(subgraph_name, req): (String, Arc<Valid<ExecutableDocument>>), 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(req.as_ref(), &resp, &subgraph_name) {
                             Ok(_) => resp,
                             Err(err) => subgraph::Response::builder()
                                 .errors(

apollo-router/src/plugins/demand_control/strategy/mod.rs

@@ -52,8 +52,12 @@ impl Strategy {
         &self,
         request: &ExecutableDocument,
         response: &subgraph::Response,
+        subgraph_name: &str,
     ) -> Result<(), DemandControlError> {
-        match self.inner.on_subgraph_response(request, response) {
+        match self
+            .inner
+            .on_subgraph_response(request, response, subgraph_name)
+        {
             Err(e) if self.mode == Mode::Enforce => Err(e),
             _ => Ok(()),
         }
@@ -93,8 +97,13 @@ impl StrategyFactory {
 
     pub(crate) fn create(&self) -> Strategy {
         let strategy: Arc<dyn StrategyImpl> = match &self.config.strategy {
-            StrategyConfig::StaticEstimated { list_size, max } => Arc::new(StaticEstimated {
+            StrategyConfig::StaticEstimated {
+                list_size,
+                max,
+                actual_cost_mode,
+            } => Arc::new(StaticEstimated {
                 max: *max,
+                actual_cost_mode: *actual_cost_mode,
                 cost_calculator: StaticCostCalculator::new(
                     self.supergraph_schema.clone(),
                     self.subgraph_schemas.clone(),
@@ -122,6 +131,7 @@ pub(crate) trait StrategyImpl: Send + Sync {
         &self,
         request: &ExecutableDocument,
         response: &subgraph::Response,
+        subgraph_name: &str,
     ) -> Result<(), DemandControlError>;
     fn on_execution_response(
         &self,