test: demand control non-clobbery

.changesets/feat_caroline_demand_control_by_subgraph.md

@@ -0,0 +1,62 @@
+### Support subgraph-level demand control ([PR #8829](https://github.com/apollographql/router/pull/8829))
+
+Subgraph-level demand control lets you enforce per-subgraph query cost limits in Apollo Router, in addition to the
+existing global cost limit for the whole supergraph. This helps you protect specific backend services that have
+different
+capacity or cost profiles from being overwhelmed by expensive operations.
+
+When a subgraph‑specific cost limit is exceeded, the router:
+
+* Still runs the rest of the operation, including other subgraphs whose cost is within limits.
+* Skips calls to only the over‑budget subgraph, and composes the response as if that subgraph had returned null, instead
+  of rejecting the entire query.
+
+Per‑subgraph limits apply to the total work for that subgraph in a single operation. For each request, the router tracks
+the aggregate estimated cost per subgraph across the entire query plan. If the same subgraph is fetched multiple times
+(for example, through entity lookups, nested fetches, or conditional branches), those costs are summed together and the
+subgraph’s limit is enforced against that total.
+
+#### Configuration
+
+```yaml
+demand_control:
+  enabled: true
+  mode: enforce
+  strategy:
+    static_estimated:
+      max: 10
+      list_size: 10
+      actual_cost_mode: by_subgraph
+      subgraphs: # <---- everything from here down is new (all fields optional)
+        all:
+          max: 8
+          list_size: 10
+        subgraphs:
+          products:
+            max: 6
+            # list_size omitted, 10 implied because of all.list_size
+          reviews:
+            list_size: 50
+            # max omitted, 8 implied because of all.max
+```
+
+#### Example
+
+Consider a topProducts query, which fetches a list of products from a products subgraph and then performs an entity
+lookup for each product in a reviews subgraph. Assume that the products cost is 10 and the reviews cost is 5, leading to
+a total estimated cost of 15 (10 + 5).
+
+Previously, you would only be able to restrict that query via `demand_control.static_estimated.max`:
+
+* If you set it <= 15, the query would execute
+* If you set it >15, the query would be rejected
+*
+
+This feature allows much more granular control. In addition to `demand_control.static_estimated.max`, which operates as
+before, there are also per subgraph maxes.
+
+For example, if you set `max = 20` and `reviews.max = 2`, the query will 'pass' the aggregate check (15 < 20) and will
+execute on the products subgraph (no limit specified), but will not execute against the reviews subgraph (5 > 2). The
+result will be composed as if the reviews subgraph had returned null.
+
+By [@carodewig](https://github.com/carodewig) in https://github.com/apollographql/router/pull/8829

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

@@ -9170,6 +9170,21 @@ expression: "&schema"
                   "description": "The maximum cost of a query",
                   "format": "double",
                   "type": "number"
+                },
+                "subgraph": {
+                  "allOf": [
+                    {
+                      "$ref": "#/definitions/SubgraphSubgraphStrategyConfigConfiguration"
+                    }
+                  ],
+                  "default": {
+                    "all": {
+                      "list_size": null,
+                      "max": null
+                    },
+                    "subgraphs": {}
+                  },
+                  "description": "Cost control by subgraph"
                 }
               },
               "required": [
@@ -10441,6 +10456,28 @@ expression: "&schema"
       },
       "type": "object"
     },
+    "SubgraphStrategyConfig": {
+      "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": {
@@ -10524,6 +10561,32 @@ expression: "&schema"
       },
       "type": "object"
     },
+    "SubgraphSubgraphStrategyConfigConfiguration": {
+      "description": "Configuration options pertaining to the subgraph server component.",
+      "properties": {
+        "all": {
+          "allOf": [
+            {
+              "$ref": "#/definitions/SubgraphStrategyConfig"
+            }
+          ],
+          "default": {
+            "list_size": null,
+            "max": null
+          },
+          "description": "options applying to all subgraphs"
+        },
+        "subgraphs": {
+          "additionalProperties": {
+            "$ref": "#/definitions/SubgraphStrategyConfig"
+          },
+          "default": {},
+          "description": "per subgraph options",
+          "type": "object"
+        }
+      },
+      "type": "object"
+    },
     "SubgraphTlsClientConfiguration": {
       "description": "Configuration options pertaining to the subgraph server component.",
       "properties": {

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<V>` by extracting a value `V` from `&T`
+    pub(crate) fn extract<V: Default + Serialize + JsonSchema>(
+        &self,
+        extract_fn: fn(&T) -> V,
+    ) -> SubgraphConfiguration<V> {
+        SubgraphConfiguration {
+            all: extract_fn(&self.all),
+            subgraphs: self
+                .subgraphs
+                .iter()
+                .map(|(k, v)| (k.clone(), extract_fn(v)))
+                .collect(),
+        }
+    }
 }
 
 impl<T> Debug for SubgraphConfiguration<T>

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

@@ -3,12 +3,17 @@ pub(in crate::plugins::demand_control) mod schema;
 pub(crate) mod static_cost;
 
 use std::collections::HashMap;
+use std::ops::AddAssign;
 
 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 new(subgraph: &str, value: f64) -> Self {
+        Self(HashMap::from([(subgraph.to_string(), value)]))
+    }
+
     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;
@@ -17,7 +22,58 @@ impl CostBySubgraph {
         }
     }
 
+    pub(crate) fn get(&self, subgraph: &str) -> Option<f64> {
+        self.0.get(subgraph).copied()
+    }
+
     pub(crate) fn total(&self) -> f64 {
         self.0.values().sum()
     }
+
+    /// Creates a new `CostBySubgraph` where each value in the map is the maximum of its value
+    /// in the two input `CostBySubgraph`s.
+    ///
+    /// ```rust
+    /// let cost1 = CostBySubgraph::new("hello", 1.0);
+    /// let mut cost2 = CostBySubgraph::new("hello", 2.0);
+    /// cost2.add_or_insert("world", 1.0);
+    ///
+    /// let max = CostBySubgraph::maximum(cost1, cost2);
+    /// assert_eq!(max.0.get("hello"), Some(2.0));
+    /// assert_eq!(max.0.get("world"), Some(1.0));
+    /// ```
+    pub(crate) fn maximum(mut cost1: Self, cost2: Self) -> Self {
+        for (subgraph, value) in cost2.0.into_iter() {
+            if let Some(subgraph_cost) = cost1.0.get_mut(&subgraph) {
+                *subgraph_cost = subgraph_cost.max(value);
+            } else {
+                cost1.0.insert(subgraph, value);
+            }
+        }
+
+        cost1
+    }
+}
+
+impl AddAssign for CostBySubgraph {
+    fn add_assign(&mut self, rhs: Self) {
+        for (subgraph, value) in rhs.0.into_iter() {
+            if let Some(subgraph_cost) = self.0.get_mut(&subgraph) {
+                *subgraph_cost += value;
+            } else {
+                self.0.insert(subgraph, value);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+impl From<&[(&str, f64)]> for CostBySubgraph {
+    fn from(values: &[(&str, f64)]) -> Self {
+        let mut cost = Self(HashMap::default());
+        for (subgraph, value) in values {
+            cost.add_or_insert(subgraph, *value);
+        }
+        cost
+    }
 }