fix: dedupe operations that are cacheable (tailcallhq#3099)

Co-authored-by: Tushar Mathur <[email protected]>

src/core/jit/graphql_executor.rs

@@ -113,7 +113,7 @@ impl JITExecutor {
             let is_const = exec.plan.is_const;
             let is_protected = exec.plan.is_protected;
 
-            let response = if exec.plan.is_query() && (exec.plan.is_dedupe || exec.plan.is_const) {
+            let response = if exec.plan.can_dedupe() {
                 self.dedupe_and_exec(exec, jit_request).await
             } else {
                 self.exec(exec, jit_request).await

src/core/jit/model.rs

@@ -1,6 +1,7 @@
 use std::borrow::Cow;
 use std::collections::HashMap;
 use std::fmt::{Debug, Formatter};
+use std::num::NonZeroU64;
 use std::sync::Arc;
 
 use async_graphql::parser::types::{ConstDirective, OperationType};
@@ -287,6 +288,7 @@ pub struct OperationPlan<Input> {
     pub is_dedupe: bool,
     pub is_const: bool,
     pub is_protected: bool,
+    pub min_cache_ttl: Option<NonZeroU64>,
     pub selection: Vec<Field<Input>>,
 }
 
@@ -317,6 +319,7 @@ impl<Input> OperationPlan<Input> {
             is_introspection_query: self.is_introspection_query,
             is_dedupe: self.is_dedupe,
             is_const: self.is_const,
+            min_cache_ttl: self.min_cache_ttl,
             is_protected: self.is_protected,
         })
     }
@@ -342,6 +345,7 @@ impl<Input> OperationPlan<Input> {
             is_introspection_query,
             is_dedupe: false,
             is_const: false,
+            min_cache_ttl: None,
             is_protected: false,
         }
     }
@@ -409,6 +413,10 @@ impl<Input> OperationPlan<Input> {
             None => true,
         }
     }
+    /// returns true if plan is dedupable
+    pub fn can_dedupe(&self) -> bool {
+        self.is_query() && (self.is_dedupe || self.is_const || self.min_cache_ttl.is_some())
+    }
 }
 
 #[derive(Clone, Debug)]

src/core/jit/request.rs

@@ -48,6 +48,7 @@ impl Request<ConstValue> {
         transform::CheckConst::new()
             .pipe(transform::CheckDedupe::new())
             .pipe(transform::CheckProtected::new())
+            .pipe(transform::CheckCache::new())
             .transform(plan)
             .to_result()
             // both transformers are infallible right now

src/core/jit/transform/check_cache.rs

@@ -0,0 +1,59 @@
+use std::convert::Infallible;
+use std::num::NonZeroU64;
+
+use tailcall_valid::Valid;
+
+use crate::core::ir::model::IR;
+use crate::core::jit::OperationPlan;
+use crate::core::Transform;
+
+/// A transformer that sets the minimum cache TTL for the operation plan based
+/// on the IR.
+pub struct CheckCache<A>(std::marker::PhantomData<A>);
+impl<A> CheckCache<A> {
+    pub fn new() -> Self {
+        Self(std::marker::PhantomData)
+    }
+}
+
+#[inline]
+fn check_cache(ir: &IR) -> Option<NonZeroU64> {
+    match ir {
+        IR::IO(_) => None,
+        IR::Cache(cache) => Some(cache.max_age),
+        IR::Path(ir, _) => check_cache(ir),
+        IR::Protect(ir) => check_cache(ir),
+        IR::Pipe(ir, ir1) => match (check_cache(ir), check_cache(ir1)) {
+            (Some(age1), Some(age2)) => Some(age1.min(age2)),
+            _ => None,
+        },
+        IR::Discriminate(_, ir) => check_cache(ir),
+        IR::Entity(hash_map) => {
+            let mut ttl = Some(NonZeroU64::MAX);
+            for ir in hash_map.values() {
+                ttl = std::cmp::min(ttl, check_cache(ir));
+            }
+            ttl
+        }
+        IR::Dynamic(_) | IR::ContextPath(_) | IR::Map(_) | IR::Service(_) => None,
+    }
+}
+
+impl<A> Transform for CheckCache<A> {
+    type Value = OperationPlan<A>;
+    type Error = Infallible;
+
+    fn transform(&self, mut plan: Self::Value) -> Valid<Self::Value, Self::Error> {
+        let mut ttl = Some(NonZeroU64::MAX);
+
+        for field in plan.selection.iter() {
+            if let Some(ir) = field.ir.as_ref() {
+                ttl = std::cmp::min(ttl, check_cache(ir));
+            }
+        }
+
+        plan.min_cache_ttl = ttl;
+
+        Valid::succeed(plan)
+    }
+}

src/core/jit/transform/input_resolver.rs

@@ -76,6 +76,7 @@ where
             is_dedupe: self.plan.is_dedupe,
             is_const: self.plan.is_const,
             is_protected: self.plan.is_protected,
+            min_cache_ttl: self.plan.min_cache_ttl,
             selection,
         })
     }

src/core/jit/transform/mod.rs

@@ -1,9 +1,11 @@
+mod check_cache;
 mod check_const;
 mod check_dedupe;
 mod check_protected;
 mod input_resolver;
 mod skip;
 
+pub use check_cache::*;
 pub use check_const::*;
 pub use check_dedupe::*;
 pub use check_protected::*;