Query deduplication

apollo-router-core/src/context.rs

@@ -4,7 +4,7 @@ use futures::Future;
 use std::sync::Arc;
 use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct Context<T = Arc<http_compat::Request<Request>>> {
     /// Original request to the Router.
     pub request: T,

apollo-router-core/src/error.rs

@@ -12,7 +12,7 @@ use tracing::level_filters::LevelFilter;
 ///
 /// Note that these are not actually returned to the client, but are instead converted to JSON for
 /// [`struct@Error`].
-#[derive(Error, Display, Debug, Serialize, Deserialize)]
+#[derive(Error, Display, Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "type")]
 #[ignore_extra_doc_attributes]
 pub enum FetchError {

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

@@ -0,0 +1,130 @@
+use crate::{http_compat, OperationKind, Request, SubgraphRequest, SubgraphResponse};
+use futures::{future::BoxFuture, lock::Mutex};
+use std::{collections::HashMap, sync::Arc, task::Poll};
+use tokio::sync::broadcast::{self, Sender};
+use tower::{BoxError, Layer, ServiceExt};
+
+pub struct QueryDeduplicationLayer;
+
+impl<S> Layer<S> for QueryDeduplicationLayer
+where
+    S: tower::Service<SubgraphRequest, Response = SubgraphResponse, Error = BoxError> + Clone,
+{
+    type Service = QueryDeduplicationService<S>;
+
+    fn layer(&self, service: S) -> Self::Service {
+        QueryDeduplicationService::new(service)
+    }
+}
+
+type WaitMap =
+    Arc<Mutex<HashMap<http_compat::Request<Request>, Sender<Result<SubgraphResponse, String>>>>>;
+
+pub struct QueryDeduplicationService<S> {
+    service: S,
+    wait_map: WaitMap,
+}
+
+impl<S> QueryDeduplicationService<S>
+where
+    S: tower::Service<SubgraphRequest, Response = SubgraphResponse, Error = BoxError> + Clone,
+{
+    fn new(service: S) -> Self {
+        QueryDeduplicationService {
+            service,
+            wait_map: Arc::new(Mutex::new(HashMap::new())),
+        }
+    }
+
+    async fn dedup(
+        service: S,
+        wait_map: WaitMap,
+        request: SubgraphRequest,
+    ) -> Result<SubgraphResponse, BoxError> {
+        loop {
+            let mut locked_wait_map = wait_map.lock().await;
+            match locked_wait_map.get_mut(&request.http_request) {
+                Some(waiter) => {
+                    // Register interest in key
+                    let mut receiver = waiter.subscribe();
+                    drop(locked_wait_map);
+
+                    match receiver.recv().await {
+                        Ok(value) => {
+                            return value
+                                .map(|response| SubgraphResponse {
+                                    response: response.response,
+                                    context: request.context,
+                                })
+                                .map_err(|e| e.into())
+                        }
+                        // there was an issue with the broadcast channel, retry fetching
+                        Err(_) => continue,
+                    }
+                }
+                None => {
+                    let (tx, _rx) = broadcast::channel(1);
+                    locked_wait_map.insert(request.http_request.clone(), tx.clone());
+                    drop(locked_wait_map);
+
+                    let context = request.context.clone();
+                    let http_request = request.http_request.clone();
+                    let res = service.ready_oneshot().await?.call(request).await;
+
+                    {
+                        let mut locked_wait_map = wait_map.lock().await;
+                        locked_wait_map.remove(&http_request);
+                    }
+
+                    // Let our waiters know
+                    let broadcast_value = res
+                        .as_ref()
+                        .map(|response| response.clone())
+                        .map_err(|e| e.to_string());
+
+                    // Our use case is very specific, so we are sure that
+                    // we won't get any errors here.
+                    tokio::task::spawn_blocking(move || {
+                        tx.send(broadcast_value)
+                            .expect("there is always at least one receiver alive, the _rx guard; qed")
+                    }).await
+                    .expect("can only fail if the task is aborted or if the internal code panics, neither is possible here; qed");
+
+                    return res.map(|response| SubgraphResponse {
+                        response: response.response,
+                        context,
+                    });
+                }
+            }
+        }
+    }
+}
+
+impl<S> tower::Service<SubgraphRequest> for QueryDeduplicationService<S>
+where
+    S: tower::Service<SubgraphRequest, Response = SubgraphResponse, Error = BoxError>
+        + Clone
+        + Send
+        + 'static,
+    <S as tower::Service<SubgraphRequest>>::Future: Send + 'static,
+{
+    type Response = SubgraphResponse;
+    type Error = BoxError;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.service.poll_ready(cx)
+    }
+
+    fn call(&mut self, request: SubgraphRequest) -> Self::Future {
+        let mut service = self.service.clone();
+
+        if request.operation_kind == OperationKind::Query {
+            let wait_map = self.wait_map.clone();
+
+            Box::pin(async move { Self::dedup(service, wait_map, request).await })
+        } else {
+            Box::pin(async move { service.call(request).await })
+        }
+    }
+}

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

@@ -1,2 +1,3 @@
 pub mod cache;
+pub mod deduplication;
 pub mod header_manipulation;

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

@@ -212,7 +212,7 @@ impl PlanNode {
 mod fetch {
     use super::selection::{select_object, Selection};
     use crate::prelude::graphql::*;
-    use serde::Deserialize;
+    use serde::{Deserialize, Deserializer};
     use std::sync::Arc;
     use tower::ServiceExt;
     use tracing::{instrument, Instrument};
@@ -234,6 +234,30 @@ mod fetch {
 
         /// The GraphQL subquery that is used for the fetch.
         operation: String,
+
+        /// The kind of operation (query, mutation or subscription)
+        #[serde(
+            rename(deserialize = "operationKind"),
+            deserialize_with = "parse_operation"
+        )]
+        operation_kind: OperationKind,
+    }
+
+    fn parse_operation<'de, D>(deserializer: D) -> Result<OperationKind, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let buf = String::deserialize(deserializer)?;
+
+        match buf.as_str() {
+            "query" => Ok(OperationKind::Query),
+            "mutation" => Ok(OperationKind::Mutation),
+            "subscription" => Ok(OperationKind::Subscription),
+            s => Err(serde::de::Error::custom(format!(
+                "unknown operation kind: {}",
+                s
+            ))),
+        }
     }
 
     struct Variables {
@@ -311,6 +335,7 @@ mod fetch {
         ) -> Result<Value, FetchError> {
             let FetchNode {
                 operation,
+                operation_kind,
                 service_name,
                 ..
             } = self;
@@ -340,6 +365,7 @@ mod fetch {
                     .unwrap()
                     .into(),
                 context: context.clone(),
+                operation_kind: *operation_kind,
             };
 
             let service = service_registry

apollo-router-core/src/query_planner/snapshots/apollo_router_core__query_planner__router_bridge_query_planner__tests__plan.snap

@@ -1,6 +1,6 @@
 ---
 source: apollo-router-core/src/query_planner/router_bridge_query_planner.rs
-assertion_line: 97
+assertion_line: 133
 expression: result
 
 ---
@@ -11,6 +11,7 @@ QueryPlan {
             requires: [],
             variable_usages: [],
             operation: "{me{name{first last}}}",
+            operation_kind: Query,
         },
     ),
 }

apollo-router-core/src/query_planner/snapshots/apollo_router_core__query_planner__tests__query_plan_from_json.snap

@@ -1,5 +1,6 @@
 ---
 source: apollo-router-core/src/query_planner/mod.rs
+assertion_line: 467
 expression: query_plan
 
 ---
@@ -11,6 +12,7 @@ Sequence {
                 requires: [],
                 variable_usages: [],
                 operation: "{topProducts{__typename ...on Book{__typename isbn}...on Furniture{name}}product(upc:\"1\"){__typename ...on Book{__typename isbn}...on Furniture{name}}}",
+                operation_kind: Query,
             },
         ),
         Parallel {
@@ -59,6 +61,7 @@ Sequence {
                                             "test_variable",
                                         ],
                                         operation: "query($representations:[_Any!]!){_entities(representations:$representations){...on Book{__typename isbn title year}}}",
+                                        operation_kind: Query,
                                     },
                                 ),
                             },
@@ -117,6 +120,7 @@ Sequence {
                                         ],
                                         variable_usages: [],
                                         operation: "query($representations:[_Any!]!){_entities(representations:$representations){...on Book{name}}}",
+                                        operation_kind: Query,
                                     },
                                 ),
                             },
@@ -164,6 +168,7 @@ Sequence {
                                         ],
                                         variable_usages: [],
                                         operation: "query($representations:[_Any!]!){_entities(representations:$representations){...on Book{__typename isbn title year}}}",
+                                        operation_kind: Query,
                                     },
                                 ),
                             },
@@ -221,6 +226,7 @@ Sequence {
                                         ],
                                         variable_usages: [],
                                         operation: "query($representations:[_Any!]!){_entities(representations:$representations){...on Book{name}}}",
+                                        operation_kind: Query,
                                     },
                                 ),
                             },

apollo-router-core/src/query_planner/testdata/query_plan.json

@@ -5,6 +5,7 @@
       "kind": "Fetch",
       "serviceName": "product",
       "variableUsages": [],
+      "operationKind": "query",
       "operation": "{topProducts{__typename ...on Book{__typename isbn}...on Furniture{name}}product(upc:\"1\"){__typename ...on Book{__typename isbn}...on Furniture{name}}}"
     },
     {
@@ -30,6 +31,7 @@
                   }
                 ],
                 "variableUsages": ["test_variable"],
+                "operationKind": "query",
                 "operation": "query($representations:[_Any!]!){_entities(representations:$representations){...on Book{__typename isbn title year}}}"
               }
             },
@@ -52,6 +54,7 @@
                   }
                 ],
                 "variableUsages": [],
+                "operationKind": "query",
                 "operation": "query($representations:[_Any!]!){_entities(representations:$representations){...on Book{name}}}"
               }
             }
@@ -77,6 +80,7 @@
                   }
                 ],
                 "variableUsages": [],
+                "operationKind": "query",
                 "operation": "query($representations:[_Any!]!){_entities(representations:$representations){...on Book{__typename isbn title year}}}"
               }
             },
@@ -99,6 +103,7 @@
                   }
                 ],
                 "variableUsages": [],
+                "operationKind": "query",
                 "operation": "query($representations:[_Any!]!){_entities(representations:$representations){...on Book{name}}}"
               }
             }

apollo-router-core/src/request.rs

@@ -10,7 +10,7 @@ use typed_builder::TypedBuilder;
 #[derive(Clone, Derivative, Serialize, Deserialize, TypedBuilder)]
 #[serde(rename_all = "camelCase")]
 #[builder(field_defaults(setter(into)))]
-#[derivative(Debug, PartialEq)]
+#[derivative(Debug, PartialEq, Eq, Hash)]
 pub struct Request {
     /// The graphql query.
     pub query: String,

apollo-router-core/src/services/http_compat.rs

@@ -1,6 +1,10 @@
 //! wrapper typpes for Request and Response from the http crate to improve their usability
 
-use std::ops::{Deref, DerefMut};
+use std::{
+    cmp::PartialEq,
+    hash::Hash,
+    ops::{Deref, DerefMut},
+};
 
 #[derive(Debug, Default)]
 pub struct Request<T> {
@@ -61,6 +65,52 @@ impl<T> DerefMut for Request<T> {
     }
 }
 
+impl<T: Hash> Hash for Request<T> {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.inner.method().hash(state);
+        self.inner.version().hash(state);
+        self.inner.uri().hash(state);
+        // this assumes headers are in the same order
+        for (name, value) in self.inner.headers() {
+            name.hash(state);
+            value.hash(state);
+        }
+        self.inner.body().hash(state);
+    }
+}
+
+impl<T: PartialEq> PartialEq for Request<T> {
+    fn eq(&self, other: &Self) -> bool {
+        let mut res = self.inner.method().eq(other.inner.method())
+            && self.inner.version().eq(&other.inner.version())
+            && self.inner.uri().eq(other.inner.uri());
+
+        if !res {
+            return false;
+        }
+        if self.inner.headers().len() != other.inner.headers().len() {
+            return false;
+        }
+
+        // this assumes headers are in the same order
+        for ((name, value), (other_name, other_value)) in self
+            .inner
+            .headers()
+            .iter()
+            .zip(other.inner.headers().iter())
+        {
+            res = name.eq(other_name) && value.eq(other_value);
+            if !res {
+                return false;
+            }
+        }
+
+        self.inner.body().eq(other.inner.body())
+    }
+}
+
+impl<T: PartialEq> Eq for Request<T> {}
+
 impl<T> From<http::Request<T>> for Request<T> {
     fn from(inner: http::Request<T>) -> Self {
         Request { inner }