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::{fetch::OperationKind, http_compat, 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/headers.rs

@@ -271,6 +271,7 @@ mod test {
     use tower::{BoxError, Layer};
     use tower::{Service, ServiceExt};
 
+    use crate::fetch::OperationKind;
     use crate::headers::{
         InsertConfig, InsertLayer, PropagateConfig, PropagateLayer, RemoveConfig, RemoveLayer,
     };
@@ -471,6 +472,7 @@ mod test {
                 .body(Request::builder().query("query").build())
                 .unwrap()
                 .into(),
+            operation_kind: OperationKind::Query,
             context: example_originating_request(),
         }
     }

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

@@ -1,3 +1,4 @@
 pub mod cache;
+pub mod deduplication;
 pub mod forbid_http_get_mutations;
 pub mod headers;

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

@@ -265,9 +265,9 @@ pub(crate) mod fetch {
         operation_kind: OperationKind,
     }
 
-    #[derive(Debug, PartialEq, Deserialize)]
+    #[derive(Copy, Clone, Debug, PartialEq, Deserialize)]
     #[serde(rename_all = "camelCase")]
-    pub(crate) enum OperationKind {
+    pub enum OperationKind {
         Query,
         Mutation,
         Subscription,
@@ -349,6 +349,7 @@ pub(crate) mod fetch {
         ) -> Result<Value, FetchError> {
             let FetchNode {
                 operation,
+                operation_kind,
                 service_name,
                 ..
             } = self;
@@ -375,6 +376,7 @@ pub(crate) 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__tests__query_plan_from_json.snap

@@ -1,6 +1,6 @@
 ---
 source: apollo-router-core/src/query_planner/mod.rs
-assertion_line: 473
+assertion_line: 467
 expression: query_plan
 
 ---

apollo-router-core/src/request.rs

@@ -10,7 +10,7 @@ use typed_builder::TypedBuilder;
 #[derive(Clone, Derivative, Serialize, Deserialize, TypedBuilder, Default)]
 #[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 }

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

@@ -1,5 +1,6 @@
 pub use self::execution_service::*;
 pub use self::router_service::*;
+use crate::fetch::OperationKind;
 use crate::layers::cache::CachingLayer;
 use crate::prelude::graphql::*;
 use moka::sync::Cache;
@@ -84,9 +85,12 @@ pub struct SubgraphRequest {
     pub http_request: http_compat::Request<Request>,
 
     pub context: Context,
+
+    pub operation_kind: OperationKind,
 }
 
 assert_impl_all!(SubgraphResponse: Send);
+#[derive(Clone, Debug)]
 pub struct SubgraphResponse {
     pub response: http_compat::Response<Response>,
 

apollo-router/src/reqwest_subgraph_service.rs

@@ -2,6 +2,7 @@ use apollo_router_core::prelude::*;
 use futures::future::BoxFuture;
 use std::sync::Arc;
 use std::task::Poll;
+use tower::BoxError;
 use tracing::Instrument;
 use typed_builder::TypedBuilder;
 
@@ -40,21 +41,21 @@ impl ReqwestSubgraphService {
 
 impl tower::Service<graphql::SubgraphRequest> for ReqwestSubgraphService {
     type Response = graphql::SubgraphResponse;
-    type Error = tower::BoxError;
+    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>> {
         //TODO backpressure
         Poll::Ready(Ok(()))
     }
 
-    fn call(
-        &mut self,
-        graphql::SubgraphRequest {
+    fn call(&mut self, request: graphql::SubgraphRequest) -> Self::Future {
+        let graphql::SubgraphRequest {
             http_request,
             context,
-        }: graphql::SubgraphRequest,
-    ) -> Self::Future {
+            ..
+        } = request;
+
         let http_client = self.http_client.clone();
         let target_url = if http_request.uri() == "/" {
             self.url.clone()
@@ -82,7 +83,14 @@ impl tower::Service<graphql::SubgraphRequest> for ReqwestSubgraphService {
             request.headers_mut().extend(headers.into_iter());
             *request.version_mut() = version;
 
-            let response = http_client.execute(request).await?;
+            let response = http_client.execute(request).await.map_err(|err| {
+                tracing::error!(fetch_error = format!("{:?}", err).as_str());
+
+                graphql::FetchError::SubrequestHttpError {
+                    service: service_name.clone(),
+                    reason: err.to_string(),
+                }
+            })?;
             let body = response
                 .bytes()
                 .instrument(tracing::debug_span!("aggregate_response_data"))

apollo-router/src/router_factory.rs

@@ -1,5 +1,6 @@
 use crate::configuration::{Configuration, ConfigurationError};
 use crate::reqwest_subgraph_service::ReqwestSubgraphService;
+use apollo_router_core::deduplication::QueryDeduplicationLayer;
 use apollo_router_core::{
     http_compat::{Request, Response},
     PluggableRouterServiceBuilder, ResponseBody, RouterRequest, Schema,
@@ -70,9 +71,9 @@ impl RouterServiceFactory for YamlRouterServiceFactory {
         let mut builder = PluggableRouterServiceBuilder::new(schema, buffer, dispatcher.clone());
 
         for (name, subgraph) in &configuration.subgraphs {
-            let mut subgraph_service = BoxService::new(ReqwestSubgraphService::new(
-                name.to_string(),
-                subgraph.routing_url.clone(),
+            let dedup_layer = QueryDeduplicationLayer;
+            let mut subgraph_service = BoxService::new(dedup_layer.layer(
+                ReqwestSubgraphService::new(name.to_string(), subgraph.routing_url.clone()),
             ));
 
             for layer in &subgraph.layers {