From 84137d221eaf9f0e0abb2c243bd9c3a2a39bcc3f Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Tue, 5 May 2026 16:16:50 +0530 Subject: [PATCH 1/2] perf: pre-serialize environment-document bytes on poll MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The released `/api/v1/environment-document` path runs `serde_json::to_vec(&Arc)` on every request and the handler then `body.to_vec()` clones the resulting `Arc<[u8]>` into a fresh `Vec` for the response body. Two allocations and a full JSON encoding per request — for a document that only changes on the polling refresh interval. This change moves the serialization to the cache layer: - New `EnvironmentsCache::get_environment_bytes` returns `Bytes`. - `LocalMemEnvironmentsCache` produces the bytes once inside `put_environment`, alongside the parsed `Arc` and the pre-computed evaluation context. Failure to serialize leaves the byte cache empty and the handler returns 503. - `EnvironmentService::get_environment_bytes` becomes a thin delegate over the cache; the previous endpoint-cache wrapper (which served the same purpose under `endpoint_caches.environment_document.use_cache`) is removed since the byte cache is always populated now. - The route handler returns the `Bytes` directly as the response body instead of `body.to_vec()`. axum accepts `Bytes` as a body with no copy. Per-request CPU on the env-doc endpoint goes from "serialize a 1-11 MB JSON tree + memcpy" to "refcounted clone". Measured on Fargate (1 vCPU / 2 GB), endpoint cache disabled: small project (1 MB env-doc): peak RPS: 343 -> 566 (1.65x) p50 @ c=25: 82ms -> 36ms medium project (11 MB env-doc): peak RPS: 21 -> 52 (2.46x) p50 @ c=25: 1.03s -> 451ms p99 at high concurrency (c>=500 small / c>=50 medium) gets worse because the bottleneck shifts from CPU to socket-write queueing on the response body — those ranges are past the useful operating point for the endpoint anyway. p50 and p90 improve across the curve. --- Cargo.lock | 1 + Cargo.toml | 1 + src/cache/environment.rs | 40 ++++++++++++++++++++++--- src/routes/environment_document.rs | 4 +-- src/services/environment.rs | 48 +++++++++++------------------- 5 files changed, 58 insertions(+), 36 deletions(-) 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..5e60962 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,6 +62,11 @@ 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() @@ -59,6 +74,7 @@ impl EnvironmentsCache for LocalMemEnvironmentsCache { async fn put_environment(&self, environment_key: &str, document: Value) -> bool { 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; @@ -92,6 +108,22 @@ impl EnvironmentsCache for LocalMemEnvironmentsCache { contexts.insert(environment_key.to_string(), context); } + // Serialize once here so /environment-document requests are an Arc-clone. + // Failure leaves the byte cache empty for this key — handler returns 503. + match serde_json::to_vec(&document) { + 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 { From ab935143de263f2b06cb6936d55241cd967e92fe Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Wed, 6 May 2026 16:03:59 +0530 Subject: [PATCH 2/2] perf: serialize environment document outside cache write locks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Holding all four RwLock write guards across serde_json::to_vec on a multi-MB document blocks every concurrent reader of the cache (flag evaluation, identity evaluation, environment-document reads). Hoist the serialize call out — the document is still uniquely owned at that point, so no lock is needed — and only enter the brief write section to swap in the produced bytes. A short read-lock pre-check returns early on unchanged polls, avoiding wasted serialization. --- src/cache/environment.rs | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/src/cache/environment.rs b/src/cache/environment.rs index 5e60962..a251160 100644 --- a/src/cache/environment.rs +++ b/src/cache/environment.rs @@ -73,12 +73,30 @@ impl EnvironmentsCache for LocalMemEnvironmentsCache { } 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) @@ -108,9 +126,9 @@ impl EnvironmentsCache for LocalMemEnvironmentsCache { contexts.insert(environment_key.to_string(), context); } - // Serialize once here so /environment-document requests are an Arc-clone. + // Pre-serialized so /environment-document requests are a refcount bump. // Failure leaves the byte cache empty for this key — handler returns 503. - match serde_json::to_vec(&document) { + match bytes_result { Ok(bytes) => { environment_bytes.insert(environment_key.to_string(), Bytes::from(bytes)); }