Response Cache Plugin

Author: ardatan

Cargo.lock

@@ -33,6 +33,7 @@ ahash = "0.8.12"
 regex-automata = "0.4.10"
 vrl = { version = "0.27.0", features = ["compiler", "parser", "value", "diagnostic", "stdlib", "core"] }
 
+ntex = { version = "2", features = ["tokio"] }
 ntex-http = "0.1.15"
 hyper-tls = { version = "0.6.0", features = ["vendored"] }
 hyper-util = { version = "0.1.16", features = [
@@ -47,6 +48,7 @@ itoa = "1.0.15"
 ryu = "1.0.20"
 indexmap = "2.10.0"
 bumpalo = "3.19.0"
+redis = "0.32.7"
 
 [dev-dependencies]
 subgraphs = { path = "../../bench/subgraphs" }

lib/executor/Cargo.toml

@@ -8,6 +8,7 @@ pub mod projection;
 pub mod response;
 pub mod utils;
 pub mod variables;
+pub mod plugins;
 
 pub use execution::plan::execute_query_plan;
 pub use executors::map::SubgraphExecutorMap;

lib/executor/src/lib.rs

@@ -0,0 +1,2 @@
+mod traits;
+mod response_cache;

lib/executor/src/plugins/mod.rs

@@ -0,0 +1,135 @@
+use bytes::Bytes;
+use dashmap::DashMap;
+use ntex::web::HttpResponse;
+use redis::Commands;
+use sonic_rs::json;
+
+use crate::{
+    plugins::traits::{
+        ControlFlow, OnExecuteEnd, OnExecuteEndPayload, OnExecuteStart, OnExecuteStartPayload,
+        OnSchemaReload, OnSchemaReloadPayload,
+    }, response::value::Value, utils::consts::TYPENAME_FIELD_NAME
+};
+
+pub struct ResponseCachePlugin {
+    redis_client: redis::Client,
+    ttl_per_type: DashMap<String, u64>,
+}
+
+impl ResponseCachePlugin {
+    pub fn try_new(redis_url: &str) -> Result<Self, redis::RedisError> {
+        let redis_client = redis::Client::open(redis_url)?;
+        Ok(Self {
+            redis_client,
+            ttl_per_type: DashMap::new(),
+        })
+    }
+}
+
+struct ResponseCacheContext {
+    key: String,
+}
+
+impl OnExecuteStart for ResponseCachePlugin {
+    fn on_execute_start(&self, mut payload: OnExecuteStartPayload) -> ControlFlow {
+        let key = format!(
+            "response_cache:{}:{:?}",
+            payload.query_plan, payload.variable_values
+        );
+        payload
+            .router_http_request
+            .extensions_mut()
+            .insert(ResponseCacheContext { key: key.clone() });
+        if let Some(mut conn) = self.redis_client.get_connection().ok() {
+            let cached_response: Option<Vec<u8>> = conn.get(&key).ok();
+            if let Some(cached_response) = cached_response {
+                return ControlFlow::Break(
+                    HttpResponse::Ok()
+                        .header("Content-Type", "application/json")
+                        .body(cached_response),
+                );
+            }
+        }
+        ControlFlow::Continue
+    }
+}
+
+impl OnExecuteEnd for ResponseCachePlugin {
+    fn on_execute_end(&self, mut payload: OnExecuteEndPayload) -> ControlFlow {
+        // Do not cache if there are errors
+        if payload.errors.len() > 0 {
+            return ControlFlow::Continue;
+        }
+        if let Some(key) = payload
+            .router_http_request
+            .extensions()
+            .get::<ResponseCacheContext>()
+            .map(|ctx| &ctx.key)
+        {
+            if let Some(mut conn) = self.redis_client.get_connection().ok() {
+                if let Some(serialized) = sonic_rs::to_vec(&payload.data).ok() {
+                    // Decide on the ttl somehow
+                    // Get the type names
+                    let mut max_ttl = 0;
+
+                    // Imagine this code is traversing the response data to find type names
+                    if let Some(obj) = payload.data.as_object() {
+                        if let Some(typename) = obj
+                            .iter()
+                            .position(|(k, _)| k == &TYPENAME_FIELD_NAME)
+                            .and_then(|idx| obj[idx].1.as_str())
+                        {
+                            if let Some(ttl) = self.ttl_per_type.get(typename).map(|v| *v) {
+                                max_ttl = max_ttl.max(ttl);
+                            }
+                        }
+                    }
+
+                    // If no ttl found, default to 60 seconds
+                    if max_ttl == 0 {
+                        max_ttl = 60;
+                    }
+
+                    // Insert the ttl into extensions for client awareness
+                    payload.extensions
+                        .insert("response_cache_ttl".to_string(), json!(max_ttl));
+
+                    // Set the cache with the decided ttl
+                    let _: () = conn.set_ex(key, serialized, max_ttl).unwrap_or(());
+                }
+            }
+        }
+        ControlFlow::Continue
+    }
+}
+
+impl OnSchemaReload for ResponseCachePlugin {
+    fn on_schema_reload(&self, payload: OnSchemaReloadPayload) {
+        // Visit the schema and update ttl_per_type based on some directive
+        payload
+            .new_schema
+            .document
+            .definitions
+            .iter()
+            .for_each(|def| {
+                if let graphql_parser::schema::Definition::TypeDefinition(type_def) = def {
+                    if let graphql_parser::schema::TypeDefinition::Object(obj_type) = type_def {
+                        for directive in &obj_type.directives {
+                            if directive.name == "cacheControl" {
+                                for arg in &directive.arguments {
+                                    if arg.0 == "maxAge" {
+                                        if let graphql_parser::query::Value::Int(max_age) = &arg.1 {
+                                            if let Some(max_age) = max_age.as_i64() {
+                                                self.ttl_per_type
+                                                    .insert(obj_type.name.clone(), max_age as u64);
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            });
+    }
+}

lib/executor/src/plugins/response_cache.rs

@@ -0,0 +1,63 @@
+use std::{collections::HashMap, sync::Arc};
+
+use async_trait::async_trait;
+use hive_router_query_planner::consumer_schema::ConsumerSchema;
+use hive_router_query_planner::planner::plan_nodes::QueryPlan;
+use ntex::web::HttpRequest;
+use ntex::web::HttpResponse;
+
+use crate::response::graphql_error::GraphQLError;
+use crate::response::value::Value;
+
+pub enum ControlFlow {
+    Continue,
+    Break(HttpResponse)
+}
+
+pub struct ExecutionResult<'exec> {
+    pub data: &'exec mut Value<'exec>,
+    pub errors: &'exec mut Vec<GraphQLError>,
+    pub extensions: &'exec mut Option<HashMap<String, Value<'exec>>>,
+}
+
+pub struct OnExecuteStartPayload<'exec> {
+    pub router_http_request: &'exec HttpRequest,
+    pub query_plan: Arc<QueryPlan>,
+
+    pub data: &'exec mut Value<'exec>,
+    pub errors: &'exec mut  Vec<GraphQLError>,
+    pub extensions: Option<&'exec mut  sonic_rs::Value>,
+
+    pub skip_execution: bool,
+
+    pub variable_values: &'exec Option<HashMap<String, sonic_rs::Value>>,
+}
+
+pub trait OnExecuteStart {
+    fn on_execute_start(&self, payload: OnExecuteStartPayload) -> ControlFlow;
+}
+
+pub struct OnExecuteEndPayload<'exec> {
+    pub router_http_request: &'exec HttpRequest,
+    pub query_plan: Arc<QueryPlan>,
+    
+    pub data: &'exec Value<'exec>,
+    pub errors: &'exec Vec<GraphQLError>,
+    pub extensions: &'exec mut HashMap<String, sonic_rs::Value>,
+
+    pub variable_values: &'exec Option<HashMap<String, sonic_rs::Value>>,
+}
+
+pub trait OnExecuteEnd {
+    fn on_execute_end(&self, payload: OnExecuteEndPayload) -> ControlFlow;
+}
+
+pub struct OnSchemaReloadPayload {
+    pub old_schema: &'static ConsumerSchema,
+    pub new_schema: &'static mut ConsumerSchema,
+}
+
+pub trait OnSchemaReload {
+    fn on_schema_reload(&self, payload: OnSchemaReloadPayload);
+}
+