diff --git a/src/cache/environment.rs b/src/cache/environment.rs index a2cbed1..5f025ad 100644 --- a/src/cache/environment.rs +++ b/src/cache/environment.rs @@ -14,8 +14,10 @@ pub trait EnvironmentsCache: Send + Sync { /// Returns Arc to avoid cloning large JSON on every request async fn get_environment(&self, environment_key: &str) -> Option>; - /// Get the pre-computed evaluation context (for flag evaluation) - async fn get_context(&self, environment_key: &str) -> Option; + /// Get the pre-computed evaluation context (for flag evaluation). + /// Returns an Arc so the caller doesn't pay a deep clone per request — + /// the context is read-heavy and only replaced on polling-interval refresh. + async fn get_context(&self, environment_key: &str) -> Option>; /// Store environment document and compute context. Returns true if changed. async fn put_environment(&self, environment_key: &str, document: Value) -> bool; @@ -29,8 +31,10 @@ pub struct LocalMemEnvironmentsCache { /// Raw environment documents (for /environment-document endpoint) /// Stored as Arc to avoid cloning large JSON on every request environments: Arc>>>, - /// Pre-computed evaluation contexts (for flag evaluation) - contexts: Arc>>, + /// Pre-computed evaluation contexts (for flag evaluation). + /// Stored as Arc so readers share the same allocation — avoids a deep + /// clone of the full context (features + segments) per request. + contexts: Arc>>>, /// Identity overrides extracted from environments identity_overrides: Arc>>>, } @@ -52,9 +56,9 @@ impl EnvironmentsCache for LocalMemEnvironmentsCache { environments.get(environment_key).cloned() // Clones Arc (cheap), not Value } - async fn get_context(&self, environment_key: &str) -> Option { + async fn get_context(&self, environment_key: &str) -> Option> { let contexts = self.contexts.read().await; - contexts.get(environment_key).cloned() + contexts.get(environment_key).cloned() // Arc clone, not deep clone } async fn put_environment(&self, environment_key: &str, document: Value) -> bool { @@ -89,7 +93,7 @@ impl EnvironmentsCache for LocalMemEnvironmentsCache { if let Ok(environment) = serde_json::from_value::(document.clone()) { let flagsmith_env: FlagsmithEnvironment = environment.to_flagsmith_environment(); let context = environment_to_context(flagsmith_env); - contexts.insert(environment_key.to_string(), context); + contexts.insert(environment_key.to_string(), Arc::new(context)); } environments.insert(environment_key.to_string(), Arc::new(document)); diff --git a/src/routes/flags.rs b/src/routes/flags.rs index 136f868..0cb9681 100644 --- a/src/routes/flags.rs +++ b/src/routes/flags.rs @@ -1,10 +1,12 @@ use crate::error::Result; +use crate::models::APIFeatureState; use crate::routes::extractors::extract_environment_key; use crate::state::AppState; use axum::{ Json, extract::{Query, State}, http::HeaderMap, + response::IntoResponse, }; use serde::Deserialize; @@ -13,22 +15,34 @@ pub struct FlagsQuery { pub feature: Option, } +pub enum FlagsResponse { + Single(APIFeatureState), + Multiple(Vec), +} + +impl IntoResponse for FlagsResponse { + fn into_response(self) -> axum::response::Response { + match self { + FlagsResponse::Single(f) => Json(f).into_response(), + FlagsResponse::Multiple(f) => Json(f).into_response(), + } + } +} + pub async fn get_flags( State(service): State, headers: HeaderMap, Query(query): Query, -) -> Result> { +) -> Result { let environment_key = extract_environment_key(&headers)?; - let flags = service + let mut flags = service .get_flags_response_data(&environment_key, query.feature.as_deref()) .await?; if query.feature.is_some() && flags.len() == 1 { - // Return single feature - Ok(Json(serde_json::to_value(&flags[0])?)) + Ok(FlagsResponse::Single(flags.swap_remove(0))) } else { - // Return all features - Ok(Json(serde_json::to_value(flags)?)) + Ok(FlagsResponse::Multiple(flags)) } }