chore: make cache_key return Option<u64> (tailcallhq#2009)

alankritdabral · May 23, 2024 · 4cca178 · 4cca178
1 parent c9851bd
commit 4cca178
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 25 deletions.
diff --git a/src/core/graphql/request_template.rs b/src/core/graphql/request_template.rs
@@ -127,11 +127,11 @@ impl RequestTemplate {
 }
 
 impl<Ctx: PathGraphql + HasHeaders + GraphQLOperationContext> CacheKey<Ctx> for RequestTemplate {
-    fn cache_key(&self, ctx: &Ctx) -> u64 {
+    fn cache_key(&self, ctx: &Ctx) -> Option<u64> {
         let mut hasher = TailcallHasher::default();
         let graphql_query = self.render_graphql_query(ctx);
         graphql_query.hash(&mut hasher);
-        hasher.finish()
+        Some(hasher.finish())
     }
 }
 

diff --git a/src/core/grpc/request_template.rs b/src/core/grpc/request_template.rs
@@ -107,11 +107,11 @@ impl RenderedRequestTemplate {
 }
 
 impl<Ctx: PathString + HasHeaders> CacheKey<Ctx> for RequestTemplate {
-    fn cache_key(&self, ctx: &Ctx) -> u64 {
+    fn cache_key(&self, ctx: &Ctx) -> Option<u64> {
         let mut hasher = TailcallHasher::default();
         let rendered_req = self.render(ctx).unwrap();
         rendered_req.hash(&mut hasher);
-        hasher.finish()
+        Some(hasher.finish())
     }
 }
 

diff --git a/src/core/http/request_template.rs b/src/core/http/request_template.rs
@@ -225,7 +225,7 @@ impl TryFrom<Endpoint> for RequestTemplate {
 }
 
 impl<Ctx: PathString + HasHeaders> CacheKey<Ctx> for RequestTemplate {
-    fn cache_key(&self, ctx: &Ctx) -> u64 {
+    fn cache_key(&self, ctx: &Ctx) -> Option<u64> {
         let mut hasher = TailcallHasher::default();
         let state = &mut hasher;
 
@@ -251,7 +251,7 @@ impl<Ctx: PathString + HasHeaders> CacheKey<Ctx> for RequestTemplate {
         let url = self.create_url(ctx).unwrap();
         url.hash(state);
 
-        hasher.finish()
+        Some(hasher.finish())
     }
 }
 
@@ -739,7 +739,7 @@ mod tests {
         use crate::core::lambda::CacheKey;
         use crate::core::mustache::Mustache;
 
-        fn assert_no_duplicate<const N: usize>(arr: [u64; N]) {
+        fn assert_no_duplicate<const N: usize>(arr: [Option<u64>; N]) {
             let set = HashSet::from(arr);
             assert_eq!(arr.len(), set.len());
         }

diff --git a/src/core/lambda/cache.rs b/src/core/lambda/cache.rs
@@ -8,7 +8,7 @@ use async_graphql_value::ConstValue;
 use super::{Eval, EvaluationContext, EvaluationError, Expression, ResolverContextLike};
 
 pub trait CacheKey<Ctx> {
-    fn cache_key(&self, ctx: &Ctx) -> u64;
+    fn cache_key(&self, ctx: &Ctx) -> Option<u64>;
 }
 
 #[derive(Clone, Debug)]
@@ -41,17 +41,20 @@ impl Eval for Cache {
         Box::pin(async move {
             if let Expression::IO(io) = self.expr.deref() {
                 let key = io.cache_key(&ctx);
-
-                if let Some(val) = ctx.request_ctx.runtime.cache.get(&key).await? {
-                    Ok(val)
+                if let Some(key) = key {
+                    if let Some(val) = ctx.request_ctx.runtime.cache.get(&key).await? {
+                        Ok(val)
+                    } else {
+                        let val = self.expr.eval(ctx.clone()).await?;
+                        ctx.request_ctx
+                            .runtime
+                            .cache
+                            .set(key, val.clone(), self.max_age)
+                            .await?;
+                        Ok(val)
+                    }
                 } else {
-                    let val = self.expr.eval(ctx.clone()).await?;
-                    ctx.request_ctx
-                        .runtime
-                        .cache
-                        .set(key, val.clone(), self.max_age)
-                        .await?;
-                    Ok(val)
+                    self.expr.eval(ctx).await
                 }
             } else {
                 Ok(self.expr.eval(ctx).await?)

diff --git a/src/core/lambda/io.rs b/src/core/lambda/io.rs
@@ -52,12 +52,16 @@ impl Eval for IO {
         if ctx.request_ctx.upstream.dedupe {
             Box::pin(async move {
                 let key = self.cache_key(&ctx);
-                ctx.request_ctx
-                    .cache
-                    .get_or_eval(key, move || Box::pin(async { self.eval_inner(ctx).await }))
-                    .await
-                    .as_ref()
-                    .clone()
+                if let Some(key) = key {
+                    ctx.request_ctx
+                        .cache
+                        .get_or_eval(key, move || Box::pin(self.eval_inner(ctx)))
+                        .await
+                        .as_ref()
+                        .clone()
+                } else {
+                    self.eval_inner(ctx).await
+                }
             })
         } else {
             Box::pin(self.eval_inner(ctx))
@@ -139,7 +143,7 @@ impl IO {
 }
 
 impl<'a, Ctx: ResolverContextLike<'a> + Sync + Send> CacheKey<EvaluationContext<'a, Ctx>> for IO {
-    fn cache_key(&self, ctx: &EvaluationContext<'a, Ctx>) -> u64 {
+    fn cache_key(&self, ctx: &EvaluationContext<'a, Ctx>) -> Option<u64> {
         match self {
             IO::Http { req_template, .. } => req_template.cache_key(ctx),
             IO::Grpc { req_template, .. } => req_template.cache_key(ctx),