Better impl

Author: ardatan

lib/executor/src/execution/plan.rs

@@ -19,7 +19,7 @@ use crate::{
         rewrites::FetchRewriteExt,
     },
     executors::{
-        common::{HttpExecutionRequest, HttpExecutionResponse},
+        common::{HttpExecutionResponse, SubgraphExecutionRequest},
         map::SubgraphExecutorMap,
     },
     headers::{
@@ -700,7 +700,7 @@ impl<'exec, 'req> Executor<'exec, 'req> {
         let variable_refs =
             select_fetch_variables(self.variable_values, node.variable_usages.as_ref());
 
-        let mut subgraph_request = HttpExecutionRequest {
+        let mut subgraph_request = SubgraphExecutionRequest {
             query: node.operation.document_str.as_str(),
             dedupe: self.dedupe_subgraph_requests,
             operation_name: node.operation_name.as_deref(),

lib/executor/src/executors/common.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, sync::Arc};
+use std::{collections::HashMap, sync::Arc, time::Duration};
 
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -9,7 +9,8 @@ use sonic_rs::Value;
 pub trait SubgraphExecutor {
     async fn execute<'a>(
         &self,
-        execution_request: HttpExecutionRequest<'a>,
+        execution_request: SubgraphExecutionRequest<'a>,
+        timeout: Option<Duration>,
     ) -> HttpExecutionResponse;
 
     fn to_boxed_arc<'a>(self) -> Arc<Box<dyn SubgraphExecutor + Send + Sync + 'a>>
@@ -26,7 +27,7 @@ pub type SubgraphExecutorBoxedArc = Arc<Box<SubgraphExecutorType>>;
 
 pub type SubgraphRequestExtensions = HashMap<String, Value>;
 
-pub struct HttpExecutionRequest<'a> {
+pub struct SubgraphExecutionRequest<'a> {
     pub query: &'a str,
     pub dedupe: bool,
     pub operation_name: Option<&'a str>,
@@ -37,7 +38,7 @@ pub struct HttpExecutionRequest<'a> {
     pub extensions: Option<SubgraphRequestExtensions>,
 }
 
-impl HttpExecutionRequest<'_> {
+impl SubgraphExecutionRequest<'_> {
     pub fn add_request_extensions_field(&mut self, key: String, value: Value) {
         self.extensions
             .get_or_insert_with(HashMap::new)

lib/executor/src/executors/http.rs

@@ -1,8 +1,10 @@
 use std::sync::Arc;
+use std::time::Duration;
 
 use crate::executors::common::HttpExecutionResponse;
 use crate::executors::dedupe::{request_fingerprint, ABuildHasher, SharedResponse};
 use dashmap::DashMap;
+use futures::TryFutureExt;
 use tokio::sync::OnceCell;
 
 use async_trait::async_trait;
@@ -18,7 +20,7 @@ use hyper_util::client::legacy::{connect::HttpConnector, Client};
 use tokio::sync::Semaphore;
 use tracing::debug;
 
-use crate::executors::common::HttpExecutionRequest;
+use crate::executors::common::SubgraphExecutionRequest;
 use crate::executors::error::SubgraphExecutorError;
 use crate::response::graphql_error::GraphQLError;
 use crate::utils::consts::CLOSE_BRACE;
@@ -75,7 +77,7 @@ impl HTTPSubgraphExecutor {
 
     fn build_request_body(
         &self,
-        execution_request: &HttpExecutionRequest<'_>,
+        execution_request: &SubgraphExecutionRequest<'_>,
     ) -> Result<Vec<u8>, SubgraphExecutorError> {
         let mut body = Vec::with_capacity(4096);
         body.put(FIRST_QUOTE_STR);
@@ -136,6 +138,7 @@ impl HTTPSubgraphExecutor {
         &self,
         body: Vec<u8>,
         headers: HeaderMap,
+        timeout: Option<Duration>,
     ) -> Result<SharedResponse, SubgraphExecutorError> {
         let mut req = hyper::Request::builder()
             .method(http::Method::POST)
@@ -150,9 +153,22 @@ impl HTTPSubgraphExecutor {
 
         debug!("making http request to {}", self.endpoint.to_string());
 
-        let res = self.http_client.request(req).await.map_err(|e| {
+        let res_fut = self.http_client.request(req).map_err(|e| {
             SubgraphExecutorError::RequestFailure(self.endpoint.to_string(), e.to_string())
-        })?;
+        });
+
+        let res = if let Some(timeout_duration) = timeout {
+            tokio::time::timeout(timeout_duration, res_fut)
+                .await
+                .map_err(|_| {
+                    SubgraphExecutorError::RequestTimeout(
+                        self.endpoint.to_string(),
+                        timeout_duration.as_secs(),
+                    )
+                })?
+        } else {
+            res_fut.await
+        }?;
 
         debug!(
             "http request to {} completed, status: {}",
@@ -211,7 +227,8 @@ impl SubgraphExecutor for HTTPSubgraphExecutor {
     #[tracing::instrument(skip_all, fields(subgraph_name = %self.subgraph_name))]
     async fn execute<'a>(
         &self,
-        execution_request: HttpExecutionRequest<'a>,
+        execution_request: SubgraphExecutionRequest<'a>,
+        timeout: Option<Duration>,
     ) -> HttpExecutionResponse {
         let body = match self.build_request_body(&execution_request) {
             Ok(body) => body,
@@ -233,7 +250,7 @@ impl SubgraphExecutor for HTTPSubgraphExecutor {
             // This unwrap is safe because the semaphore is never closed during the application's lifecycle.
             // `acquire()` only fails if the semaphore is closed, so this will always return `Ok`.
             let _permit = self.semaphore.acquire().await.unwrap();
-            return match self._send_request(body, headers).await {
+            return match self._send_request(body, headers, timeout).await {
                 Ok(shared_response) => HttpExecutionResponse {
                     body: shared_response.body,
                     headers: shared_response.headers,
@@ -265,7 +282,7 @@ impl SubgraphExecutor for HTTPSubgraphExecutor {
                     // This unwrap is safe because the semaphore is never closed during the application's lifecycle.
                     // `acquire()` only fails if the semaphore is closed, so this will always return `Ok`.
                     let _permit = self.semaphore.acquire().await.unwrap();
-                    self._send_request(body, headers).await
+                    self._send_request(body, headers, timeout).await
                 };
                 // It's important to remove the entry from the map before returning the result.
                 // This ensures that once the OnceCell is set, no future requests can join it.

lib/executor/src/executors/map.rs

@@ -21,7 +21,8 @@ use crate::{
     execution::client_request_details::ClientRequestDetails,
     executors::{
         common::{
-            HttpExecutionRequest, HttpExecutionResponse, SubgraphExecutor, SubgraphExecutorBoxedArc,
+            HttpExecutionResponse, SubgraphExecutionRequest, SubgraphExecutor,
+            SubgraphExecutorBoxedArc,
         },
         dedupe::{ABuildHasher, SharedResponse},
         duration_or_prog::{compile_duration_expression, DurationOrProgram},
@@ -109,7 +110,7 @@ impl SubgraphExecutorMap {
     pub async fn execute<'a, 'req>(
         &self,
         subgraph_name: &str,
-        execution_request: HttpExecutionRequest<'a>,
+        execution_request: SubgraphExecutionRequest<'a>,
         client_request: &ClientRequestDetails<'a, 'req>,
     ) -> HttpExecutionResponse {
         match self.get_or_create_executor(subgraph_name, client_request) {
@@ -119,30 +120,15 @@ impl SubgraphExecutorMap {
                     .get(subgraph_name)
                     .map(|t| resolve_duration_prog(t.value(), subgraph_name, client_request));
                 match timeout {
-                    Some(Ok(dur)) => tokio::time::timeout(dur, executor.execute(execution_request))
-                        .await
-                        .unwrap_or_else(|_| {
-                            error!(
-                                "Request to subgraph '{}' timed out after {:?}",
-                                subgraph_name, dur,
-                            );
-                            self.internal_server_error_response(
-                                SubgraphExecutorError::RequestTimeout(
-                                    subgraph_name.to_string(),
-                                    dur.as_millis() as u64,
-                                )
-                                .into(),
-                                subgraph_name,
-                            )
-                        }),
+                    Some(Ok(dur)) => executor.execute(execution_request, Some(dur)).await,
                     Some(Err(err)) => {
                         error!(
-                            "Timeout expression resolution error for subgraph '{}': {}",
+                            "Failed to resolve timeout for subgraph '{}': {}",
                             subgraph_name, err,
                         );
                         self.internal_server_error_response(err.into(), subgraph_name)
                     }
-                    None => executor.execute(execution_request).await,
+                    None => executor.execute(execution_request, None).await,
                 }
             }
             Err(err) => {