Iteration

Author: ardatan

lib/executor/src/plugins/response_cache.rs

@@ -1,13 +1,9 @@
 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,
-    },
+    plugins::traits::{ControlFlow, OnExecuteControlFlow, OnExecutePayload, OnSchemaReloadPayload, RouterPlugin},
     utils::consts::TYPENAME_FIELD_NAME,
 };
 
@@ -26,20 +22,15 @@ impl ResponseCachePlugin {
     }
 }
 
-pub struct ResponseCacheContext {
-    key: String,
-}
-
-impl OnExecuteStart for ResponseCachePlugin {
-    fn on_execute_start(&self, payload: OnExecuteStartPayload) -> ControlFlow {
+impl RouterPlugin for ResponseCachePlugin {
+    fn on_execute<'exec>(
+        &'static self,
+        payload: OnExecutePayload<'exec>,
+    ) -> OnExecuteControlFlow<'exec> {
         let key = format!(
             "response_cache:{}:{:?}",
             payload.query_plan, payload.variable_values
         );
-        payload
-            .router_http_request
-            .extensions_mut()
-            .insert(ResponseCacheContext { key: key.clone() });
         if let Ok(mut conn) = self.redis_client.get_connection() {
             let cached_response: Option<Vec<u8>> = conn.get(&key).ok();
             if let Some(cached_response) = cached_response {
@@ -50,61 +41,47 @@ impl OnExecuteStart for ResponseCachePlugin {
                 );
             }
         }
-        ControlFlow::Continue
-    }
-}
+        return ControlFlow::OnEnd(Box::new(move |payload| {
+                // Do not cache if there are errors
+                if !payload.errors.is_empty() {
+                    return ();
+                }
 
-impl OnExecuteEnd for ResponseCachePlugin {
-    fn on_execute_end(&self, payload: OnExecuteEndPayload) -> ControlFlow {
-        // Do not cache if there are errors
-        if !payload.errors.is_empty() {
-            return ControlFlow::Continue;
-        }
-        if let Some(key) = payload
-            .router_http_request
-            .extensions()
-            .get::<ResponseCacheContext>()
-            .map(|ctx| &ctx.key)
-        {
-            if let Ok(mut conn) = self.redis_client.get_connection() {
-                if let Ok(serialized) = sonic_rs::to_vec(&payload.data) {
-                    // Decide on the ttl somehow
-                    // Get the type names
-                    let mut max_ttl = 0;
+                if let Ok(mut conn) = self.redis_client.get_connection() {
+                    if let Ok(serialized) = sonic_rs::to_vec(&payload.data) {
+                        // 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);
+                        // 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;
-                    }
+                        // 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));
+                        // Insert the ttl into extensions for client awareness
+                        payload
+                            .extensions
+                            .insert("response_cache_ttl".to_string(), sonic_rs::json!(max_ttl));
 
-                    // Set the cache with the decided ttl
-                    let _: () = conn.set_ex(key, serialized, max_ttl).unwrap_or(());
+                        // 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

lib/executor/src/plugins/traits.rs

@@ -1,3 +1,4 @@
+use std::hash::Hash;
 use std::{collections::HashMap, sync::Arc};
 
 use hive_router_query_planner::consumer_schema::ConsumerSchema;
@@ -8,8 +9,9 @@ use ntex::web::HttpResponse;
 use crate::response::graphql_error::GraphQLError;
 use crate::response::value::Value;
 
-pub enum ControlFlow {
+pub enum ControlFlow<TEndPayload> {
     Continue,
+    OnEnd(Box<dyn FnOnce(TEndPayload) -> ()>),
     Break(HttpResponse),
 }
 
@@ -19,21 +21,29 @@ pub struct ExecutionResult<'exec> {
     pub extensions: &'exec mut Option<HashMap<String, Value<'exec>>>,
 }
 
-pub struct OnExecuteStartPayload<'exec> {
+pub struct OnExecutePayload<'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 extensions: &'exec mut HashMap<String, 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 type OnExecuteControlFlow<'exec> = ControlFlow<OnExecutePayload<'exec>>;
+
+pub trait RouterPlugin {
+    fn on_execute<'exec>(
+        &'static self, 
+        _payload: OnExecutePayload<'exec>,
+    ) -> ControlFlow<OnExecutePayload<'exec>> {
+        ControlFlow::Continue
+    }
+    fn on_schema_reload(&self, _payload: OnSchemaReloadPayload) {}
 }
 
 pub struct OnExecuteEndPayload<'exec> {
@@ -47,15 +57,7 @@ pub struct OnExecuteEndPayload<'exec> {
     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);
-}