Move dedup to a traffic_shaping plugin (#753)

* Move dedup to a traffic_shaping plugin

* Only buffer if dedup layer is applied. We'll have to do something more intelligent when we add more functionality.

* Take in review comments

Co-authored-by: bryn <bryn@apollographql.com>

CHANGELOG.md

@@ -38,11 +38,16 @@ server:
 ```
 
 ## 🐛 Fixes
+- **Move query dedup to an experimental `traffic_shaping` plugin** ([PR #753](https://github.com/apollographql/router/pull/753))
+
+  The experimental `traffic_shaping` plugin will be a central location where we can add things such as rate limiting and retry.
+
 - **Remove `hasNext` from our response objects** ([PR #733](https://github.com/apollographql/router/pull/733))
 
   `hasNext` is a field in the response that may be used in future to support features such as defer and stream. However, we are some way off supporting this and including it now may break clients. It has been removed.
 
 - **Extend Apollo uplink configurability** ([PR #741](https://github.com/apollographql/router/pull/741))
+
   Uplink url and poll interval can now be configured via command line arg and env variable:
   ```bash
     --apollo-schema-config-delivery-endpoint <apollo-schema-config-delivery-endpoint>

apollo-router-core/src/layers/deduplication.rs

@@ -8,6 +8,7 @@ use std::{
 use tokio::sync::broadcast::{self, Sender};
 use tower::{BoxError, Layer, ServiceExt};
 
+#[derive(Default)]
 pub struct QueryDeduplicationLayer;
 
 impl<S> Layer<S> for QueryDeduplicationLayer

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

@@ -1,2 +1,3 @@
 mod forbid_mutations;
 mod headers;
+mod traffic_shaping;

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

@@ -0,0 +1,119 @@
+use std::collections::HashMap;
+
+use schemars::JsonSchema;
+use serde::Deserialize;
+use tower::util::BoxService;
+use tower::{BoxError, ServiceBuilder, ServiceExt};
+
+use crate::deduplication::QueryDeduplicationLayer;
+use crate::plugin::Plugin;
+use crate::{register_plugin, SubgraphRequest, SubgraphResponse};
+
+#[derive(PartialEq, Debug, Clone, Deserialize, JsonSchema)]
+struct Shaping {
+    dedup: 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),
+            },
+        }
+    }
+}
+
+#[derive(PartialEq, Debug, Clone, Deserialize, JsonSchema)]
+struct Config {
+    #[serde(default)]
+    all: Option<Shaping>,
+    #[serde(default)]
+    subgraphs: HashMap<String, Shaping>,
+}
+
+struct TrafficShaping {
+    config: Config,
+}
+
+#[async_trait::async_trait]
+impl Plugin for TrafficShaping {
+    type Config = Config;
+
+    fn new(config: Self::Config) -> Result<Self, BoxError> {
+        Ok(Self { config })
+    }
+
+    fn subgraph_service(
+        &mut self,
+        name: &str,
+        service: BoxService<SubgraphRequest, SubgraphResponse, BoxError>,
+    ) -> BoxService<SubgraphRequest, SubgraphResponse, BoxError> {
+        // Either we have the subgraph config and we merge it with the all config, or we just have the all config or we have nothing.
+        let all_config = self.config.all.as_ref();
+        let subgraph_config = self.config.subgraphs.get(name);
+        let final_config = Self::merge_config(all_config, subgraph_config);
+
+        if let Some(config) = final_config {
+            ServiceBuilder::new()
+                .option_layer(config.dedup.unwrap_or_default().then(|| {
+                    //Buffer is required because dedup layer requires a clone service.
+                    ServiceBuilder::new()
+                        .layer(QueryDeduplicationLayer::default())
+                        .buffer(20_000)
+                }))
+                .service(service)
+                .boxed()
+        } else {
+            service
+        }
+    }
+}
+
+impl TrafficShaping {
+    fn merge_config(
+        all_config: Option<&Shaping>,
+        subgraph_config: Option<&Shaping>,
+    ) -> Option<Shaping> {
+        let merged_subgraph_config = subgraph_config.map(|c| c.merge(all_config));
+        merged_subgraph_config.or_else(|| all_config.cloned())
+    }
+}
+
+register_plugin!("experimental", "traffic_shaping", TrafficShaping);
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_merge_config() {
+        let config = serde_yaml::from_str::<Config>(
+            r#"
+        all:
+          dedup: true
+        subgraphs: 
+          products:
+            dedup: false
+        "#,
+        )
+        .unwrap();
+
+        assert_eq!(TrafficShaping::merge_config(None, None), None);
+        assert_eq!(
+            TrafficShaping::merge_config(config.all.as_ref(), None),
+            config.all
+        );
+        assert_eq!(
+            TrafficShaping::merge_config(config.all.as_ref(), config.subgraphs.get("products"))
+                .as_ref(),
+            config.subgraphs.get("products")
+        );
+
+        assert_eq!(
+            TrafficShaping::merge_config(None, config.subgraphs.get("products")).as_ref(),
+            config.subgraphs.get("products")
+        );
+    }
+}

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

@@ -1,6 +1,5 @@
 ---
 source: apollo-router/src/configuration/mod.rs
-assertion_line: 439
 expression: "&schema"
 ---
 {
@@ -271,6 +270,33 @@ expression: "&schema"
             }
           },
           "additionalProperties": false
+        },
+        "experimental.traffic_shaping": {
+          "type": "object",
+          "properties": {
+            "all": {
+              "type": "object",
+              "properties": {
+                "dedup": {
+                  "type": "boolean",
+                  "nullable": true
+                }
+              },
+              "nullable": true
+            },
+            "subgraphs": {
+              "type": "object",
+              "additionalProperties": {
+                "type": "object",
+                "properties": {
+                  "dedup": {
+                    "type": "boolean",
+                    "nullable": true
+                  }
+                }
+              }
+            }
+          }
         }
       },
       "additionalProperties": false

apollo-router/src/router_factory.rs

@@ -1,5 +1,4 @@
 use crate::configuration::{Configuration, ConfigurationError};
-use apollo_router_core::deduplication::QueryDeduplicationLayer;
 use apollo_router_core::{
     http_compat::{Request, Response},
     PluggableRouterServiceBuilder, ResponseBody, RouterRequest, Schema,
@@ -9,7 +8,7 @@ use apollo_router_core::{DynPlugin, ReqwestSubgraphService};
 use std::sync::Arc;
 use tower::buffer::Buffer;
 use tower::util::{BoxCloneService, BoxService};
-use tower::{BoxError, Layer, ServiceBuilder, ServiceExt};
+use tower::{BoxError, ServiceBuilder, ServiceExt};
 use tower_service::Service;
 
 /// Factory for creating a RouterService
@@ -78,9 +77,7 @@ impl RouterServiceFactory for YamlRouterServiceFactory {
         }
 
         for (name, _) in schema.subgraphs() {
-            let dedup_layer = QueryDeduplicationLayer;
-            let subgraph_service =
-                BoxService::new(dedup_layer.layer(ReqwestSubgraphService::new(name.to_string())));
+            let subgraph_service = BoxService::new(ReqwestSubgraphService::new(name.to_string()));
 
             builder = builder.with_subgraph_service(name, subgraph_service);
         }

docs/source/config.json

@@ -14,6 +14,7 @@
       "Logging": "/configuration/logging",
       "Header propagation": "/configuration/header-propagation",
       "OpenTelemetry": "/configuration/opentelemetry",
+      "Traffic shaping": "/configuration/traffic-shaping",
       "Usage reporting": "/configuration/spaceport"
     },
     "Development workflow": {

docs/source/configuration/traffic-shaping.mdx

@@ -0,0 +1,35 @@
+---
+title: Traffic shaping
+description: Configuring traffic shaping
+---
+
+import { Link } from "gatsby";
+
+> ⚠️ Apollo Router support for traffic shaping is currently experimental.
+
+The Apollo Router provides experimental support for traffic shaping.
+
+Currently features are limited, but are expected to grow over time:
+
+* **Sub-query deduplication** - Identical, in-flight, non-mutation sub-queries are compressed into a single request.
+
+## Configuration
+To configure traffic shaping add the `traffic_shaping` plugin to `your router.yaml`:
+
+```yaml title="router.yaml"
+plugins:
+  experimental.traffic_shaping:
+    all:
+      dedup: true # Enable dedup for all subgraphs.
+    subgraphs: 
+      products:
+        dedup: false # Disable dedup for products.
+```
+
+Note that configuration in the `subgraphs` section will take precedence over that in the `all` section.
+
+### Sub-query deduplication
+
+Deduplication will cause any identical, in-flight, non-mutation sub-queries to be merged into a single request. This can reduce network bandwidth and CPU at your subgraph.
+
+Note that only in flight requests are deduplicated.

router.yaml

@@ -0,0 +1,7 @@
+plugins:
+  experimental.traffic_shaping:
+    all:
+      dedup: true
+    subgraphs:
+      products:
+        dedup: true