diff --git a/Cargo.lock b/Cargo.lock index 16dfc9d..c886a05 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -519,6 +519,7 @@ dependencies = [ "async-trait", "axum", "axum-test", + "bytes", "chrono", "flagsmith-flag-engine", "http", diff --git a/Cargo.toml b/Cargo.toml index f985108..6ddf89e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,7 @@ thiserror = "2" anyhow = "1" chrono = { version = "0.4", features = ["serde"] } async-trait = "0.1" +bytes = "1" lru = "0.16" http = "1" validator = { version = "0.20", features = ["derive"] } diff --git a/src/cache/environment.rs b/src/cache/environment.rs index a2cbed1..a251160 100644 --- a/src/cache/environment.rs +++ b/src/cache/environment.rs @@ -1,19 +1,26 @@ use async_trait::async_trait; +use bytes::Bytes; use flagsmith_flag_engine::engine_eval::{EngineEvaluationContext, environment_to_context}; use flagsmith_flag_engine::environments::Environment as FlagsmithEnvironment; use serde_json::Value; use std::collections::HashMap; use std::sync::Arc; use tokio::sync::RwLock; +use tracing::error; use crate::models::engine::Environment; #[async_trait] pub trait EnvironmentsCache: Send + Sync { - /// Get the raw environment document (for /environment-document endpoint) - /// Returns Arc to avoid cloning large JSON on every request + /// Get the raw environment document (for flag evaluation paths that + /// need the parsed JSON tree). Returns Arc to avoid cloning. async fn get_environment(&self, environment_key: &str) -> Option>; + /// Get the pre-serialized JSON bytes for the /environment-document + /// endpoint. These are produced once when the document is stored, so + /// the request path is a refcount bump. + async fn get_environment_bytes(&self, environment_key: &str) -> Option; + /// Get the pre-computed evaluation context (for flag evaluation) async fn get_context(&self, environment_key: &str) -> Option; @@ -26,9 +33,11 @@ pub trait EnvironmentsCache: Send + Sync { #[derive(Clone, Default)] pub struct LocalMemEnvironmentsCache { - /// Raw environment documents (for /environment-document endpoint) - /// Stored as Arc to avoid cloning large JSON on every request + /// Raw environment documents (kept parsed for evaluation paths). environments: Arc>>>, + /// Pre-serialized JSON bytes for `/environment-document` responses. + /// Populated on `put_environment`; cheap to clone (refcounted). + environment_bytes: Arc>>, /// Pre-computed evaluation contexts (for flag evaluation) contexts: Arc>>, /// Identity overrides extracted from environments @@ -39,6 +48,7 @@ impl LocalMemEnvironmentsCache { pub fn new() -> Self { Self { environments: Arc::new(RwLock::new(HashMap::new())), + environment_bytes: Arc::new(RwLock::new(HashMap::new())), contexts: Arc::new(RwLock::new(HashMap::new())), identity_overrides: Arc::new(RwLock::new(HashMap::new())), } @@ -52,17 +62,41 @@ impl EnvironmentsCache for LocalMemEnvironmentsCache { environments.get(environment_key).cloned() // Clones Arc (cheap), not Value } + async fn get_environment_bytes(&self, environment_key: &str) -> Option { + let environment_bytes = self.environment_bytes.read().await; + environment_bytes.get(environment_key).cloned() // Bytes clone is a refcount bump + } + async fn get_context(&self, environment_key: &str) -> Option { let contexts = self.contexts.read().await; contexts.get(environment_key).cloned() } async fn put_environment(&self, environment_key: &str, document: Value) -> bool { + // Skip work entirely when the document is unchanged. A short-lived + // read guard on `environments` is enough — the equality check itself + // can be expensive on multi-MB Values, but it doesn't block readers. + { + let environments = self.environments.read().await; + if let Some(existing) = environments.get(environment_key) { + if existing.as_ref() == &document { + return false; + } + } + } + + // Heavy CPU work runs while `document` is still uniquely owned — + // outside any lock — so concurrent flag-evaluation requests aren't + // blocked while we serialize a multi-MB document. + let bytes_result = serde_json::to_vec(&document); + let mut environments = self.environments.write().await; + let mut environment_bytes = self.environment_bytes.write().await; let mut contexts = self.contexts.write().await; let mut identity_overrides = self.identity_overrides.write().await; - // Check if document changed + // Re-check under the write guards in case a concurrent put landed + // an identical document between the read and write guards. let changed = environments .get(environment_key) .map(|existing| existing.as_ref() != &document) @@ -92,6 +126,22 @@ impl EnvironmentsCache for LocalMemEnvironmentsCache { contexts.insert(environment_key.to_string(), context); } + // Pre-serialized so /environment-document requests are a refcount bump. + // Failure leaves the byte cache empty for this key — handler returns 503. + match bytes_result { + Ok(bytes) => { + environment_bytes.insert(environment_key.to_string(), Bytes::from(bytes)); + } + Err(err) => { + error!( + environment_key, + error = %err, + "failed to serialize environment document for byte cache" + ); + environment_bytes.remove(environment_key); + } + } + environments.insert(environment_key.to_string(), Arc::new(document)); } diff --git a/src/routes/environment_document.rs b/src/routes/environment_document.rs index 7e68e3a..ca9b59e 100644 --- a/src/routes/environment_document.rs +++ b/src/routes/environment_document.rs @@ -18,8 +18,8 @@ pub async fn get_environment_document( return Err(EdgeProxyError::FlagsmithUnknownKey(environment_key)); } - // Get pre-serialized bytes (with endpoint caching if enabled) + // Pre-serialized bytes from the environment cache; populated at poll time. let body = service.get_environment_bytes(&environment_key).await?; - Ok(([(header::CONTENT_TYPE, "application/json")], body.to_vec())) + Ok(([(header::CONTENT_TYPE, "application/json")], body)) } diff --git a/src/services/environment.rs b/src/services/environment.rs index 914d7c0..6a71b19 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -3,6 +3,7 @@ use crate::config::settings::{AppSettings, EnvironmentKeyPair}; use crate::error::{EdgeProxyError, Result}; use crate::models::{APIFeatureState, IdentityResponse, IdentityWithTraits}; use crate::services::feature_utils::filter_out_server_key_only_flag_results; +use bytes::Bytes; use chrono::{DateTime, Utc}; use flagsmith_flag_engine::engine::get_evaluation_result; use flagsmith_flag_engine::engine_eval::{FlagResult, add_identity_to_context}; @@ -211,40 +212,27 @@ impl EnvironmentService { .ok_or_else(|| EdgeProxyError::ServiceUnavailable("Environment not loaded".to_string())) } - /// Get pre-serialized environment document bytes with endpoint caching - pub async fn get_environment_bytes(&self, environment_key: &str) -> Result> { - if self.endpoint_cache.is_environment_document_cache_enabled() { - let cache_key = CacheKey::new( + /// Return the pre-serialized JSON bytes for the environment document. + /// + /// Bytes are produced once when polling refreshes the cache, so the + /// request path is a `Bytes` clone (refcount bump). + pub async fn get_environment_bytes(&self, environment_key: &str) -> Result { + if !self.key_mapping.contains_key(environment_key) { + return Err(EdgeProxyError::FlagsmithUnknownKey( environment_key.to_string(), - "environment_document".to_string(), - "".to_string(), - ); - - if let Some(cached_bytes) = self - .endpoint_cache - .get_environment_document(&cache_key) - .await - { - return Ok(cached_bytes); - } + )); } - let document = self.get_environment(environment_key).await?; - - let bytes: Arc<[u8]> = serde_json::to_vec(&*document)?.into(); - - if self.endpoint_cache.is_environment_document_cache_enabled() { - let cache_key = CacheKey::new( - environment_key.to_string(), - "environment_document".to_string(), - "".to_string(), - ); - self.endpoint_cache - .put_environment_document(cache_key, bytes.clone()) - .await; - } + let client_key = self + .server_to_client + .get(environment_key) + .map(|s| s.as_str()) + .unwrap_or(environment_key); - Ok(bytes) + self.cache + .get_environment_bytes(client_key) + .await + .ok_or_else(|| EdgeProxyError::ServiceUnavailable("Environment not loaded".to_string())) } fn extract_server_key_only_ids(document: &serde_json::Value) -> Vec {