perf: pre-serialize environment-document bytes on poll

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"] }

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<Arc<Value>>;
 
+    /// 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<Bytes>;
+
     /// Get the pre-computed evaluation context (for flag evaluation)
     async fn get_context(&self, environment_key: &str) -> Option<EngineEvaluationContext>;
 
@@ -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<Value> to avoid cloning large JSON on every request
+    /// Raw environment documents (kept parsed for evaluation paths).
     environments: Arc<RwLock<HashMap<String, Arc<Value>>>>,
+    /// Pre-serialized JSON bytes for `/environment-document` responses.
+    /// Populated on `put_environment`; cheap to clone (refcounted).
+    environment_bytes: Arc<RwLock<HashMap<String, Bytes>>>,
     /// Pre-computed evaluation contexts (for flag evaluation)
     contexts: Arc<RwLock<HashMap<String, EngineEvaluationContext>>>,
     /// 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<Bytes> {
+        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<EngineEvaluationContext> {
         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));
         }
 

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))
 }

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<Arc<[u8]>> {
-        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<Bytes> {
+        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<u32> {