fix: cache control headers (tailcallhq#465)

ujjawal-khare-27 · Oct 11, 2023 · 1774b06 · 1774b06
1 parent 61f80ad
commit 1774b06
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 5 deletions.
diff --git a/benches/impl_path_string_for_evaluation_context.rs b/benches/impl_path_string_for_evaluation_context.rs
@@ -130,7 +130,7 @@ fn assert_test(eval_ctx: &EvaluationContext<'_, MockGraphqlContext>) {
 }
 
 fn bench_main(c: &mut Criterion) {
-  let mut req_ctx = RequestContext { req_headers: TEST_HEADERS.clone(), ..Default::default() };
+  let mut req_ctx = RequestContext::default().req_headers(TEST_HEADERS.clone());
 
   req_ctx.server.vars = TEST_VARS.clone();
 

diff --git a/src/http/request_context.rs b/src/http/request_context.rs
@@ -1,3 +1,5 @@
+use std::sync::{Arc, Mutex};
+
 use derive_setters::Setters;
 use hyper::HeaderMap;
 
@@ -9,6 +11,7 @@ pub struct RequestContext {
   pub http_client: DefaultHttpClient,
   pub server: Server,
   pub req_headers: HeaderMap,
+  min_max_age: Arc<Mutex<Option<u64>>>,
 }
 
 impl Default for RequestContext {
@@ -19,12 +22,31 @@ impl Default for RequestContext {
 
 impl RequestContext {
   pub fn new(http_client: DefaultHttpClient, server: Server) -> Self {
-    Self { req_headers: HeaderMap::new(), http_client, server }
+    Self { req_headers: HeaderMap::new(), http_client, server, min_max_age: Arc::new(Mutex::new(None)) }
   }
 
   pub async fn execute(&self, req: reqwest::Request) -> anyhow::Result<Response> {
     Ok(self.http_client.execute(req).await?)
   }
+  fn set_min_max_age_conc(&self, min_max_age: u64) {
+    *self.min_max_age.lock().unwrap() = Some(min_max_age);
+  }
+  pub fn get_min_max_age(&self) -> Option<u64> {
+    *self.min_max_age.lock().unwrap()
+  }
+
+  pub fn set_min_max_age(&self, max_age: u64) {
+    let min_max_age_lock = self.get_min_max_age();
+    match min_max_age_lock {
+      Some(min_max_age) if max_age < min_max_age => {
+        self.set_min_max_age_conc(max_age);
+      }
+      None => {
+        self.set_min_max_age_conc(max_age);
+      }
+      _ => {}
+    }
+  }
 }
 
 impl From<&ServerContext> for RequestContext {
@@ -33,3 +55,32 @@ impl From<&ServerContext> for RequestContext {
     Self::new(http_client, server_ctx.server.clone())
   }
 }
+
+#[cfg(test)]
+mod test {
+
+  use crate::http::RequestContext;
+
+  #[test]
+  fn test_update_max_age_less_than_existing() {
+    let req_ctx = RequestContext::default();
+    req_ctx.set_min_max_age(120);
+    req_ctx.set_min_max_age(60);
+    assert_eq!(req_ctx.get_min_max_age(), Some(60));
+  }
+
+  #[test]
+  fn test_update_max_age_greater_than_existing() {
+    let req_ctx = RequestContext::default();
+    req_ctx.set_min_max_age(60);
+    req_ctx.set_min_max_age(120);
+    assert_eq!(req_ctx.get_min_max_age(), Some(60));
+  }
+
+  #[test]
+  fn test_update_max_age_no_existing_value() {
+    let req_ctx = RequestContext::default();
+    req_ctx.set_min_max_age(120);
+    assert_eq!(req_ctx.get_min_max_age(), Some(120));
+  }
+}
diff --git a/src/http/server.rs b/src/http/server.rs
@@ -27,8 +27,14 @@ async fn graphql_request(req: Request<Body>, server_ctx: &ServerContext) -> Resu
   let bytes = hyper::body::to_bytes(req.into_body()).await?;
   let request: async_graphql_hyper::GraphQLRequest = serde_json::from_slice(&bytes)?;
   let req_ctx = Arc::new(RequestContext::from(server_ctx).req_headers(headers));
-  let response = request.data(req_ctx.clone()).execute(&server_ctx.schema).await;
-  // Todo: Handle Cache-Control headers
+  let mut response = request.data(req_ctx.clone()).execute(&server_ctx.schema).await;
+
+  if server_ctx.server.enable_cache_control() {
+    if let Some(ttl) = req_ctx.get_min_max_age() {
+      response = response.set_cache_control(ttl as i32);
+    }
+  }
+
   response.to_response()
 }
 fn not_found() -> Result<Response<Body>> {

diff --git a/src/lambda/expression.rs b/src/lambda/expression.rs
@@ -11,7 +11,7 @@ use thiserror::Error;
 
 use super::ResolverContextLike;
 use crate::config::group_by::GroupBy;
-use crate::http::{DefaultHttpClient, HttpDataLoader};
+use crate::http::{max_age, DefaultHttpClient, HttpDataLoader};
 #[cfg(feature = "unsafe-js")]
 use crate::javascript;
 use crate::json::JsonLike;
@@ -114,6 +114,11 @@ impl Expression {
                   .await
                   .map_err(|e| EvaluationError::IOException(e.to_string()))?
                   .unwrap_or_default();
+                if ctx.req_ctx.server.enable_cache_control() && resp.status.is_success() {
+                  if let Some(max_age) = max_age(&resp) {
+                    ctx.req_ctx.set_min_max_age(max_age.as_secs());
+                  }
+                }
                 return Ok(resp.body);
               }
 
@@ -130,6 +135,11 @@ impl Expression {
                   .validate(&res.body)
                   .map_err(EvaluationError::from)?;
               }
+              if ctx.req_ctx.server.enable_cache_control() && res.status.is_success() {
+                if let Some(max_age) = max_age(&res) {
+                  ctx.req_ctx.set_min_max_age(max_age.as_secs());
+                }
+              }
               Ok(res.body)
             }
             Operation::JS(input, script) => {