Variable deduplication in query planner (#872)

Signed-off-by: Benjamin Coenen <5719034+bnjjj@users.noreply.github.com>

NEXT_CHANGELOG.md

@@ -37,6 +37,15 @@ To upgrade, remove any dependency on the former in `Cargo.toml` files (keeping o
 ```
 
 ## 🚀 Features
+<<<<<<< HEAD
+### Add an experimental optimization to deduplicate variables in query planner [PR #872](https://github.com/apollographql/router/pull/872)
+Get rid of duplicated variables in requests and responses of the query planner. This optimization is disabled by default, if you want to enable it you just need override your configuration:
+
+```yaml title="router.yaml"
+plugins:
+  experimental.traffic_shaping:
+    variables_deduplication: true # Enable the variables deduplication optimization
+=======
 ### Add more customizable metrics ([PR #1159](https://github.com/apollographql/router/pull/1159))
 Added the ability to add custom attributes/labels on metrics via the configuration file.
 Example:
@@ -53,6 +62,7 @@ telemetry:
             rename: "payload_type"
             default: "application/json"
           - named: "x-custom-header-to-add"
+>>>>>>> ecb875a36b2fdf88025a1fb571dcd2fb5e009778
 ```
 
 ### Allow to set a custom health check path ([PR #1164](https://github.com/apollographql/router/pull/1164))

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

@@ -318,7 +318,8 @@ expression: "&schema"
             "all": {
               "type": "object",
               "properties": {
-                "dedup": {
+                "query_deduplication": {
+                  "description": "Enable query deduplication",
                   "type": "boolean",
                   "nullable": true
                 }
@@ -330,12 +331,18 @@ expression: "&schema"
               "additionalProperties": {
                 "type": "object",
                 "properties": {
-                  "dedup": {
+                  "query_deduplication": {
+                    "description": "Enable query deduplication",
                     "type": "boolean",
                     "nullable": true
                   }
                 }
               }
+            },
+            "variables_deduplication": {
+              "description": "Enable variable deduplication optimization when sending requests to subgraphs (https://github.com/apollographql/router/issues/87)",
+              "type": "boolean",
+              "nullable": true
             }
           }
         }

apollo-router/src/json_ext.rs

@@ -365,7 +365,7 @@ where
 /// A GraphQL path element that is composes of strings or numbers.
 /// e.g `/book/3/name`
 #[doc(hidden)]
-#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash)]
 #[serde(untagged)]
 pub enum PathElement {
     /// A path element that given an array will flatmap the content.
@@ -424,7 +424,7 @@ where
 ///
 /// This can be composed of strings and numbers
 #[doc(hidden)]
-#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Default)]
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Default, Hash)]
 #[serde(transparent)]
 pub struct Path(pub Vec<PathElement>);
 

apollo-router/src/plugins/include_subgraph_errors.rs

@@ -178,7 +178,6 @@ mod test {
         );
 
         let builder = PluggableRouterServiceBuilder::new(schema.clone());
-
         let builder = builder
             .with_dyn_plugin("experimental.include_subgraph_errors".to_string(), plugin)
             .with_subgraph_service("accounts", account_service.clone())

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

@@ -20,19 +20,23 @@ use tower::{BoxError, ServiceBuilder, ServiceExt};
 
 use crate::plugin::Plugin;
 use crate::plugins::traffic_shaping::deduplication::QueryDeduplicationLayer;
-use crate::{register_plugin, ServiceBuilderExt, SubgraphRequest, SubgraphResponse};
+use crate::{
+    register_plugin, QueryPlannerRequest, QueryPlannerResponse, ServiceBuilderExt, SubgraphRequest,
+    SubgraphResponse,
+};
 
 #[derive(PartialEq, Debug, Clone, Deserialize, JsonSchema)]
 struct Shaping {
-    dedup: Option<bool>,
+    /// Enable query deduplication
+    query_deduplication: Option<bool>,
 }
 
 impl Shaping {
     fn merge(&self, fallback: Option<&Shaping>) -> Shaping {
         match fallback {
             None => self.clone(),
             Some(fallback) => Shaping {
-                dedup: self.dedup.or(fallback.dedup),
+                query_deduplication: self.query_deduplication.or(fallback.query_deduplication),
             },
         }
     }
@@ -44,6 +48,8 @@ struct Config {
     all: Option<Shaping>,
     #[serde(default)]
     subgraphs: HashMap<String, Shaping>,
+    /// Enable variable deduplication optimization when sending requests to subgraphs (https://github.com/apollographql/router/issues/87)
+    variables_deduplication: Option<bool>,
 }
 
 struct TrafficShaping {
@@ -70,7 +76,7 @@ impl Plugin for TrafficShaping {
 
         if let Some(config) = final_config {
             ServiceBuilder::new()
-                .option_layer(config.dedup.unwrap_or_default().then(|| {
+                .option_layer(config.query_deduplication.unwrap_or_default().then(|| {
                     //Buffer is required because dedup layer requires a clone service.
                     ServiceBuilder::new()
                         .layer(QueryDeduplicationLayer::default())
@@ -82,6 +88,22 @@ impl Plugin for TrafficShaping {
             service
         }
     }
+
+    fn query_planning_service(
+        &mut self,
+        service: BoxService<QueryPlannerRequest, QueryPlannerResponse, BoxError>,
+    ) -> BoxService<QueryPlannerRequest, QueryPlannerResponse, BoxError> {
+        if matches!(self.config.variables_deduplication, Some(true)) {
+            service
+                .map_request(|mut req: QueryPlannerRequest| {
+                    req.query_plan_options.enable_variable_deduplication = true;
+                    req
+                })
+                .boxed()
+        } else {
+            service
+        }
+    }
 }
 
 impl TrafficShaping {
@@ -98,17 +120,144 @@ register_plugin!("experimental", "traffic_shaping", TrafficShaping);
 
 #[cfg(test)]
 mod test {
+    use std::sync::Arc;
+
+    use futures::stream::BoxStream;
+    use futures::StreamExt;
+    use once_cell::sync::Lazy;
+    use serde_json_bytes::{ByteString, Value};
+    use tower::{util::BoxCloneService, Service};
+
+    use crate::{
+        utils::test::mock::subgraph::MockSubgraph, DynPlugin, Object,
+        PluggableRouterServiceBuilder, ResponseBody, RouterRequest, RouterResponse, Schema,
+    };
+
     use super::*;
 
+    static EXPECTED_RESPONSE: Lazy<ResponseBody> = Lazy::new(|| {
+        ResponseBody::GraphQL(serde_json::from_str(r#"{"data":{"topProducts":[{"upc":"1","name":"Table","reviews":[{"id":"1","product":{"name":"Table"},"author":{"id":"1","name":"Ada Lovelace"}},{"id":"4","product":{"name":"Table"},"author":{"id":"2","name":"Alan Turing"}}]},{"upc":"2","name":"Couch","reviews":[{"id":"2","product":{"name":"Couch"},"author":{"id":"1","name":"Ada Lovelace"}}]}]}}"#).unwrap())
+    });
+
+    static VALID_QUERY: &str = r#"query TopProducts($first: Int) { topProducts(first: $first) { upc name reviews { id product { name } author { id name } } } }"#;
+
+    async fn execute_router_test(
+        query: &str,
+        body: &ResponseBody,
+        mut router_service: BoxCloneService<
+            RouterRequest,
+            BoxStream<'static, RouterResponse>,
+            BoxError,
+        >,
+    ) {
+        let request = RouterRequest::fake_builder()
+            .query(query.to_string())
+            .variable("first", 2usize)
+            .build()
+            .expect("expecting valid request");
+
+        let response = router_service
+            .ready()
+            .await
+            .unwrap()
+            .call(request)
+            .await
+            .unwrap()
+            .next()
+            .await
+            .unwrap();
+        assert_eq!(response.response.body(), body);
+    }
+
+    async fn build_mock_router_with_variable_dedup_optimization(
+        plugin: Box<dyn DynPlugin>,
+    ) -> BoxCloneService<RouterRequest, BoxStream<'static, RouterResponse>, BoxError> {
+        let mut extensions = Object::new();
+        extensions.insert("test", Value::String(ByteString::from("value")));
+
+        let account_mocks = vec![
+            (
+                r#"{"query":"query TopProducts__accounts__3($representations:[_Any!]!){_entities(representations:$representations){...on User{name}}}","operationName":"TopProducts__accounts__3","variables":{"representations":[{"__typename":"User","id":"1"},{"__typename":"User","id":"2"}]}}"#,
+                r#"{"data":{"_entities":[{"name":"Ada Lovelace"},{"name":"Alan Turing"}]}}"#
+            )
+        ].into_iter().map(|(query, response)| (serde_json::from_str(query).unwrap(), serde_json::from_str(response).unwrap())).collect();
+        let account_service = MockSubgraph::new(account_mocks);
+
+        let review_mocks = vec![
+            (
+                r#"{"query":"query TopProducts__reviews__1($representations:[_Any!]!){_entities(representations:$representations){...on Product{reviews{id product{__typename upc}author{__typename id}}}}}","operationName":"TopProducts__reviews__1","variables":{"representations":[{"__typename":"Product","upc":"1"},{"__typename":"Product","upc":"2"}]}}"#,
+                r#"{"data":{"_entities":[{"reviews":[{"id":"1","product":{"__typename":"Product","upc":"1"},"author":{"__typename":"User","id":"1"}},{"id":"4","product":{"__typename":"Product","upc":"1"},"author":{"__typename":"User","id":"2"}}]},{"reviews":[{"id":"2","product":{"__typename":"Product","upc":"2"},"author":{"__typename":"User","id":"1"}}]}]}}"#
+            )
+            ].into_iter().map(|(query, response)| (serde_json::from_str(query).unwrap(), serde_json::from_str(response).unwrap())).collect();
+        let review_service = MockSubgraph::new(review_mocks);
+
+        let product_mocks = vec![
+            (
+                r#"{"query":"query TopProducts__products__0($first:Int){topProducts(first:$first){__typename upc name}}","operationName":"TopProducts__products__0","variables":{"first":2}}"#,
+                r#"{"data":{"topProducts":[{"__typename":"Product","upc":"1","name":"Table"},{"__typename":"Product","upc":"2","name":"Couch"}]}}"#
+            ),
+            (
+                r#"{"query":"query TopProducts__products__2($representations:[_Any!]!){_entities(representations:$representations){...on Product{name}}}","operationName":"TopProducts__products__2","variables":{"representations":[{"__typename":"Product","upc":"1"},{"__typename":"Product","upc":"2"}]}}"#,
+                r#"{"data":{"_entities":[{"name":"Table"},{"name":"Couch"}]}}"#
+            )
+            ].into_iter().map(|(query, response)| (serde_json::from_str(query).unwrap(), serde_json::from_str(response).unwrap())).collect();
+
+        let product_service = MockSubgraph::new(product_mocks).with_extensions(extensions);
+
+        let schema: Arc<Schema> = Arc::new(
+            include_str!(
+                "../../../../apollo-router-benchmarks/benches/fixtures/supergraph.graphql"
+            )
+            .parse()
+            .unwrap(),
+        );
+
+        let builder = PluggableRouterServiceBuilder::new(schema.clone());
+
+        let builder = builder
+            .with_dyn_plugin("experimental.traffic_shaping".to_string(), plugin)
+            .with_subgraph_service("accounts", account_service.clone())
+            .with_subgraph_service("reviews", review_service.clone())
+            .with_subgraph_service("products", product_service.clone());
+
+        let (router, _) = builder.build().await.expect("should build");
+
+        router
+    }
+
+    async fn get_taffic_shaping_plugin(config: &serde_json::Value) -> Box<dyn DynPlugin> {
+        // Build a redacting plugin
+        crate::plugins()
+            .get("experimental.traffic_shaping")
+            .expect("Plugin not found")
+            .create_instance(config)
+            .await
+            .expect("Plugin not created")
+    }
+
+    #[tokio::test]
+    async fn it_returns_valid_response_for_deduplicated_variables() {
+        let config = serde_yaml::from_str::<serde_json::Value>(
+            r#"
+        variables_deduplication: true
+        "#,
+        )
+        .unwrap();
+        // Build a redacting plugin
+        let plugin = get_taffic_shaping_plugin(&config).await;
+        let router = build_mock_router_with_variable_dedup_optimization(plugin).await;
+        execute_router_test(VALID_QUERY, &*EXPECTED_RESPONSE, router).await;
+    }
+
     #[test]
     fn test_merge_config() {
         let config = serde_yaml::from_str::<Config>(
             r#"
         all:
-          dedup: true
+          query_deduplication: true
         subgraphs: 
           products:
-            dedup: false
+            query_deduplication: false
         "#,
         )
         .unwrap();

apollo-router/src/query_planner/bridge_query_planner.rs

@@ -53,7 +53,7 @@ impl Service<QueryPlannerRequest> for BridgeQueryPlanner {
                         "presence of a query has been checked by the RouterService before; qed",
                     ),
                     body.operation_name.to_owned(),
-                    Default::default(),
+                    req.query_plan_options,
                 )
                 .await
             {
@@ -73,7 +73,7 @@ impl QueryPlanner for BridgeQueryPlanner {
         &self,
         query: String,
         operation: Option<String>,
-        _options: QueryPlanOptions,
+        options: QueryPlanOptions,
     ) -> Result<Arc<query_planner::QueryPlan>, QueryPlannerError> {
         let planner_result = self
             .planner
@@ -90,6 +90,7 @@ impl QueryPlanner for BridgeQueryPlanner {
             } => Ok(Arc::new(query_planner::QueryPlan {
                 usage_reporting,
                 root: node,
+                options,
             })),
             PlanSuccess {
                 data: QueryPlan { node: None },

apollo-router/src/query_planner/caching_query_planner.rs

@@ -87,7 +87,7 @@ where
                 .clone()
                 .expect("presence of a query has been checked by the RouterService before; qed"),
             body.operation_name.to_owned(),
-            QueryPlanOptions::default(),
+            request.query_plan_options,
         );
         let cm = self.cm.clone();
         Box::pin(async move {