perf: improve performance for JIT (#3017)

Co-authored-by: Sandipsinh Rathod <[email protected]> Co-authored-by: Sandipsinh Dilipsinh Rathod <[email protected]> Co-authored-by: Tushar Mathur <[email protected]> Co-authored-by: Amit Singh <[email protected]>
tailcallhq · Oct 21, 2024 · 912aad1 · 912aad1 · github-actions · Oct 21, 2024
1 parent 8ac5f7a
commit 912aad1
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 20 deletions.
diff --git a/src/core/app_context.rs b/src/core/app_context.rs
@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
 use async_graphql::dynamic::{self, DynamicRequest};
+use async_graphql::Response;
 use async_graphql_value::ConstValue;
 use dashmap::DashMap;
 
@@ -29,7 +30,7 @@ pub struct AppContext {
     pub endpoints: EndpointSet<Checked>,
     pub auth_ctx: Arc<GlobalAuthContext>,
     pub dedupe_handler: Arc<DedupeResult<IoId, ConstValue, Error>>,
-    pub dedupe_operation_handler: DedupeResult<OperationId, Lift<async_graphql::Response>, Error>,
+    pub dedupe_operation_handler: DedupeResult<OperationId, Lift<Response>, Error>,
     pub operation_plans: DashMap<OPHash, OperationPlan<async_graphql_value::Value>>,
 }
 

diff --git a/src/core/http/request_handler.rs b/src/core/http/request_handler.rs
@@ -55,11 +55,8 @@ fn not_found() -> Result<Response<Body>> {
 }
 
 fn create_request_context(req: &Request<Body>, app_ctx: &AppContext) -> RequestContext {
-    let upstream = app_ctx.blueprint.upstream.clone();
-    let allowed = upstream.allowed_headers;
-    let allowed_headers = create_allowed_headers(req.headers(), &allowed);
-
-    let _allowed = app_ctx.blueprint.server.get_experimental_headers();
+    let allowed_headers =
+        create_allowed_headers(req.headers(), &app_ctx.blueprint.upstream.allowed_headers);
     RequestContext::from(app_ctx).allowed_headers(allowed_headers)
 }
 
@@ -132,17 +129,15 @@ pub async fn graphql_request<T: DeserializeOwned + GraphQLRequestLike>(
 async fn execute_query<T: DeserializeOwned + GraphQLRequestLike>(
     app_ctx: &Arc<AppContext>,
     req_ctx: &Arc<RequestContext>,
-    mut request: T,
+    request: T,
     req: Parts,
 ) -> anyhow::Result<Response<Body>> {
     let mut response = if app_ctx.blueprint.server.enable_jit {
-        let is_query = request.is_query();
         let operation_id = request.operation_id(&req.headers);
         request
             .execute(&JITExecutor::new(
                 app_ctx.clone(),
                 req_ctx.clone(),
-                is_query,
                 operation_id,
             ))
             .await
@@ -157,7 +152,7 @@ async fn execute_query<T: DeserializeOwned + GraphQLRequestLike>(
 }
 
 fn create_allowed_headers(headers: &HeaderMap, allowed: &BTreeSet<String>) -> HeaderMap {
-    let mut new_headers = HeaderMap::new();
+    let mut new_headers = HeaderMap::with_capacity(allowed.len());
     for (k, v) in headers.iter() {
         if allowed
             .iter()

diff --git a/src/core/jit/exec_const.rs b/src/core/jit/exec_const.rs
@@ -80,7 +80,7 @@ impl ConstValueExecutor {
         let response = exe.execute(synth).await;
 
         // Cache the response if we know the output is always the same
-        if is_const {
+        if is_const && self.response.is_none() {
             self.response = Some(response.clone());
         }
 

diff --git a/src/core/jit/graphql_executor.rs b/src/core/jit/graphql_executor.rs
@@ -19,19 +19,18 @@ use crate::core::merge_right::MergeRight;
 pub struct JITExecutor {
     app_ctx: Arc<AppContext>,
     req_ctx: Arc<RequestContext>,
-    is_query: bool,
     operation_id: OperationId,
 }
 
 impl JITExecutor {
     pub fn new(
         app_ctx: Arc<AppContext>,
         req_ctx: Arc<RequestContext>,
-        is_query: bool,
         operation_id: OperationId,
     ) -> Self {
-        Self { app_ctx, req_ctx, is_query, operation_id }
+        Self { app_ctx, req_ctx, operation_id }
     }
+
     #[inline(always)]
     async fn exec(
         &self,
@@ -40,12 +39,10 @@ impl JITExecutor {
     ) -> Response {
         let is_introspection_query = self.app_ctx.blueprint.server.get_enable_introspection()
             && exec.plan.is_introspection_query;
-
         let jit_resp = exec
             .execute(&self.req_ctx, &jit_request)
             .await
             .into_async_graphql();
-
         if is_introspection_query {
             let async_req = async_graphql::Request::from(jit_request).only_introspection();
             let async_resp = self.app_ctx.execute(async_req).await;
@@ -54,6 +51,7 @@ impl JITExecutor {
             jit_resp
         }
     }
+
     #[inline(always)]
     async fn dedupe_and_exec(
         &self,
@@ -117,7 +115,7 @@ impl Executor for JITExecutor {
 
         async move {
             let jit_request = jit::Request::from(request);
-            let exec = if let Some(op) = self.app_ctx.operation_plans.get(&hash) {
+            let mut exec = if let Some(op) = self.app_ctx.operation_plans.get(&hash) {
                 ConstValueExecutor::from(op.value().clone())
             } else {
                 let exec = match ConstValueExecutor::try_new(&jit_request, &self.app_ctx) {
@@ -128,9 +126,9 @@ impl Executor for JITExecutor {
                 exec
             };
 
-            if let Some(ref response) = exec.response {
-                response.clone().into_async_graphql()
-            } else if self.is_query && exec.plan.is_dedupe {
+            if let Some(response) = std::mem::take(&mut exec.response) {
+                response.into_async_graphql()
+            } else if exec.plan.is_query() && exec.plan.is_dedupe {
                 self.dedupe_and_exec(exec, jit_request).await
             } else {
                 self.exec(exec, jit_request).await
Thread Stats	Avg	Stdev	Max	+/- Stdev
Latency	9.99ms	4.18ms	143.81ms	92.78%
Req/Sec	2.55k	319.81	3.51k	84.92%