From e2f646ad22b9b880ab0b8ed4ace613af47bd1aac Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 5 Sep 2026 15:50:09 +0530 Subject: [PATCH 01/22] feat: report served-request usage The proxy counts every SDK request it serves, aggregated per environment and resource inside resolve_key so no entry point can forget to count, and flushes to POST {api_url}/proxy/usage/ every usage_flush_interval_seconds (default 60) with the proxy key. - flushes are chunked to the server's 1000-row cap; a rejected (4xx) chunk is dropped, a failed (5xx/network) chunk is kept for next time - document fetches carry X-Proxy-Key so core stops counting the proxy's own polls - statically configured environments keep their old billing: fetches unmarked, served requests unreported - inert without proxy_key; the final partial window is lost on shutdown --- src/config/settings.rs | 19 ++ src/environments.rs | 5 + src/lib.rs | 1 + src/main.rs | 5 + src/services/environment.rs | 116 +++++++++- src/usage.rs | 150 +++++++++++++ tests/test_usage_tracking.rs | 423 +++++++++++++++++++++++++++++++++++ 7 files changed, 713 insertions(+), 6 deletions(-) create mode 100644 src/usage.rs create mode 100644 tests/test_usage_tracking.rs diff --git a/src/config/settings.rs b/src/config/settings.rs index 70f1b30..ba2ef9e 100644 --- a/src/config/settings.rs +++ b/src/config/settings.rs @@ -108,6 +108,10 @@ pub struct AppSettings { pub api_poll_timeout_seconds: u64, #[serde(default = "default_allow_origins")] pub allow_origins: Vec, + // A zero interval would panic the flush task silently. + #[serde(default = "default_usage_flush_interval")] + #[validate(range(min = 1))] + pub usage_flush_interval_seconds: u64, #[serde(default)] pub server: ServerSettings, #[serde(default)] @@ -132,6 +136,10 @@ fn default_allow_origins() -> Vec { vec!["*".to_string()] } +fn default_usage_flush_interval() -> u64 { + 60 +} + impl Default for AppSettings { fn default() -> Self { Self { @@ -141,6 +149,7 @@ impl Default for AppSettings { api_poll_frequency_seconds: default_api_poll_frequency(), api_poll_timeout_seconds: default_api_poll_timeout(), allow_origins: default_allow_origins(), + usage_flush_interval_seconds: default_usage_flush_interval(), server: ServerSettings::default(), logging: LoggingSettings::default(), health_check: HealthCheckSettings::default(), @@ -194,6 +203,16 @@ mod tests { assert!(settings.validate().is_err()); } + #[test] + fn test_config_with_zero_usage_flush_interval_is_invalid() { + // Given an interval that would panic the flush task + let settings: AppSettings = + serde_json::from_str(r#"{"usage_flush_interval_seconds": 0}"#).unwrap(); + + // Then + assert!(settings.validate().is_err()); + } + #[test] fn test_config_with_only_a_proxy_key_is_valid() { // Given a config file that relies entirely on the proxy config diff --git a/src/environments.rs b/src/environments.rs index ac8a7fa..9612d40 100644 --- a/src/environments.rs +++ b/src/environments.rs @@ -80,6 +80,11 @@ impl EnvironmentIndex { index } + /// Whether the key belongs to a statically configured environment. + pub fn is_static(&self, key: &str) -> bool { + self.protected.contains(key) + } + /// Resolve a presented key — client- or server-side — to its /// environment's keys. A server-side key resolves only while it is /// valid, so a deactivation delivered by the proxy config and an diff --git a/src/lib.rs b/src/lib.rs index 2bfc72a..14cb2b3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,3 +6,4 @@ pub mod models; pub mod routes; pub mod services; pub mod state; +pub mod usage; diff --git a/src/main.rs b/src/main.rs index 405c5fd..38b1930 100644 --- a/src/main.rs +++ b/src/main.rs @@ -44,6 +44,11 @@ async fn main() -> anyhow::Result<()> { polling_service.poll_environments().await; }); + let usage_service = environment_service.clone(); + tokio::spawn(async move { + usage_service.flush_usage_periodically().await; + }); + let addr = SocketAddr::from(( settings .server diff --git a/src/services/environment.rs b/src/services/environment.rs index 423a89b..5f786ba 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -6,6 +6,7 @@ use crate::models::{ APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, }; use crate::services::feature_utils::filter_out_server_key_only_flag_results; +use crate::usage::{Resource, UsageCounts, UsageRow}; use chrono::{DateTime, Utc}; use flagsmith_flag_engine::engine::get_evaluation_result; use flagsmith_flag_engine::engine_eval::{FlagResult, add_identity_to_context}; @@ -23,6 +24,7 @@ pub struct EnvironmentService { pub settings: AppSettings, pub last_updated_at: Arc>>>, environments: EnvironmentIndex, + usage: UsageCounts, } impl EnvironmentService { @@ -41,6 +43,7 @@ impl EnvironmentService { settings, last_updated_at: Arc::new(RwLock::new(None)), environments, + usage: UsageCounts::default(), } } @@ -153,10 +156,26 @@ impl EnvironmentService { Ok(response.json().await?) } - fn resolve_key(&self, environment_key: &str) -> Result> { - self.environments + /// Resolve a presented key, counting the request for usage reporting. + /// Every SDK entry point resolves through here, so a served request + /// cannot be missed. + fn resolve_key( + &self, + environment_key: &str, + resource: Resource, + ) -> Result> { + let keys = self + .environments .resolve(environment_key) - .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string())) + .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string()))?; + self.track_usage(&keys.client_key, resource); + Ok(keys) + } + + fn track_usage(&self, client_key: &str, resource: Resource) { + if self.settings.proxy_key.is_some() && !self.environments.is_static(client_key) { + self.usage.increment(client_key, resource); + } } async fn fetch_environment(&self, keys: &EnvironmentKeys) -> Result { @@ -210,6 +229,14 @@ impl EnvironmentService { .client .get(&next_url) .header("X-Environment-Key", server_side_key); + // Core excludes marked fetches from API usage — the proxy + // reports served requests instead. Static environments stay + // unmarked and keep their old billing. + if !self.environments.is_static(server_side_key) { + if let Some(proxy_key) = &self.settings.proxy_key { + request = request.header("X-Proxy-Key", proxy_key); + } + } // If-Modified-Since is meaningful only on the first request; the // upstream pagination cursor (page_id) drives subsequent fetches. if document.is_none() { @@ -263,7 +290,11 @@ impl EnvironmentService { } pub async fn get_environment(&self, environment_key: &str) -> Result> { - let keys = self.resolve_key(environment_key)?; + // Lookup, not an SDK entry point: callers count via resolve_key. + let keys = self + .environments + .resolve(environment_key) + .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string()))?; // Documents are cached under the client key, whichever key was presented self.cache @@ -274,6 +305,7 @@ impl EnvironmentService { /// Get pre-serialized environment document bytes pub async fn get_environment_bytes(&self, environment_key: &str) -> Result> { + self.resolve_key(environment_key, Resource::EnvironmentDocument)?; let document = self.get_environment(environment_key).await?; Ok(serde_json::to_vec(&*document)?.into()) } @@ -298,7 +330,7 @@ impl EnvironmentService { ) -> Result> { // TODO: serve server-side keys here, see // https://github.com/Flagsmith/edge-proxy-rs/issues/24 - self.resolve_key(environment_key)?; + self.resolve_key(environment_key, Resource::Flags)?; let context = self .cache @@ -342,7 +374,7 @@ impl EnvironmentService { identity: &IdentityWithTraits, environment_key: &str, ) -> Result { - self.resolve_key(environment_key)?; + self.resolve_key(environment_key, Resource::Identities)?; // Get pre-computed context from cache let context = self @@ -397,6 +429,78 @@ impl EnvironmentService { self.refresh_environment_caches().await; } } + + /// The usage endpoint's batch cap — MAX_USAGE_ROWS in the edge_proxy + /// app. Flushes are chunked to it so a large environment set can + /// never be rejected outright. + const MAX_ROWS_PER_FLUSH: usize = 1000; + + /// Report the counts accumulated since the last flush to the usage + /// endpoint, in chunks the server accepts. A rejected (4xx) chunk is + /// dropped — retrying cannot heal a rejection, and losing one window + /// beats resending a poisoned batch forever. Any other failure keeps + /// the chunk for the next flush. Returns false when any chunk was + /// not accepted. + pub async fn flush_usage(&self) -> bool { + let Some(proxy_key) = &self.settings.proxy_key else { + return true; + }; + let mut rows = self.usage.drain(); + let url = format!("{}/proxy/usage/", self.settings.api_url); + let mut all_success = true; + + while !rows.is_empty() { + let chunk: Vec = rows + .drain(..rows.len().min(Self::MAX_ROWS_PER_FLUSH)) + .collect(); + let result = self + .client + .post(&url) + .header("X-Proxy-Key", proxy_key) + .json(&chunk) + .send() + .await; + match result { + Ok(response) if response.status().is_success() => {} + Ok(response) if response.status().is_client_error() => { + error!( + "Usage report rejected with {}: dropping {} rows", + response.status(), + chunk.len() + ); + all_success = false; + } + Ok(response) => { + error!("Failed to report usage: {}", response.status()); + self.usage.merge(chunk); + all_success = false; + } + Err(e) => { + error!("Failed to report usage: {}", e); + self.usage.merge(chunk); + all_success = false; + } + } + } + + all_success + } + + pub async fn flush_usage_periodically(self: Arc) { + if self.settings.proxy_key.is_none() { + return; + } + let mut interval = tokio::time::interval(Duration::from_secs( + self.settings.usage_flush_interval_seconds, + )); + // The first tick completes immediately, before anything is counted. + interval.tick().await; + + loop { + interval.tick().await; + self.flush_usage().await; + } + } } /// Format the cached document's `updated_at` as an RFC 2822 `If-Modified-Since` diff --git a/src/usage.rs b/src/usage.rs new file mode 100644 index 0000000..c5c365e --- /dev/null +++ b/src/usage.rs @@ -0,0 +1,150 @@ +use serde::Serialize; +use std::collections::HashMap; +use std::sync::Mutex; + +/// An SDK endpoint as the usage endpoint names it. +#[derive(Serialize, Clone, Copy, PartialEq, Eq, Hash, Debug)] +#[serde(rename_all = "kebab-case")] +pub enum Resource { + Flags, + Identities, + EnvironmentDocument, +} + +/// One row of the `POST /proxy/usage/` body. +#[derive(Serialize, Debug, PartialEq)] +pub struct UsageRow { + pub client_side_key: String, + pub resource: Resource, + pub count: u64, +} + +/// Requests served since the last flush, aggregated per environment and +/// resource. Keyed by resolved client keys only, so the map is bounded by +/// the served environment set. +/// +/// Uses `std::sync::Mutex`: every operation writes, and guards are held +/// only for a map operation, never across an await. +#[derive(Default)] +pub struct UsageCounts { + by_key: Mutex>, +} + +impl UsageCounts { + pub fn increment(&self, client_key: &str, resource: Resource) { + let mut by_key = self.by_key.lock().expect("usage counts lock poisoned"); + *by_key + .entry((client_key.to_string(), resource)) + .or_default() += 1; + } + + /// Take everything counted so far, leaving the map empty. + pub fn drain(&self) -> Vec { + let by_key = std::mem::take(&mut *self.by_key.lock().expect("usage counts lock poisoned")); + by_key + .into_iter() + .map(|((client_side_key, resource), count)| UsageRow { + client_side_key, + resource, + count, + }) + .collect() + } + + /// Add drained counts back, on top of anything counted since — for + /// when a flush fails and the rows must survive until the next one. + pub fn merge(&self, counts: Vec) { + let mut by_key = self.by_key.lock().expect("usage counts lock poisoned"); + for row in counts { + let entry = by_key + .entry((row.client_side_key, row.resource)) + .or_default(); + *entry = entry.saturating_add(row.count); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn increment_aggregates_by_key_and_resource() { + // Given + let counts = UsageCounts::default(); + + // When + counts.increment("client", Resource::Flags); + counts.increment("client", Resource::Flags); + counts.increment("client", Resource::Identities); + + // Then + let mut rows = counts.drain(); + rows.sort_by_key(|row| format!("{:?}", row.resource)); + assert_eq!( + rows, + vec![ + UsageRow { + client_side_key: "client".to_string(), + resource: Resource::Flags, + count: 2, + }, + UsageRow { + client_side_key: "client".to_string(), + resource: Resource::Identities, + count: 1, + }, + ] + ); + } + + #[test] + fn drain_leaves_nothing_behind() { + // Given + let counts = UsageCounts::default(); + counts.increment("client", Resource::Flags); + + // When + counts.drain(); + + // Then + assert!(counts.drain().is_empty()); + } + + #[test] + fn merge_adds_drained_counts_back() { + // Given a drained batch and a request counted since + let counts = UsageCounts::default(); + counts.increment("client", Resource::Flags); + let drained = counts.drain(); + counts.increment("client", Resource::Flags); + + // When + counts.merge(drained); + + // Then + let rows = counts.drain(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].count, 2); + } + + #[test] + fn usage_row_serializes_to_the_contract_shape() { + // Given + let row = UsageRow { + client_side_key: "client".to_string(), + resource: Resource::EnvironmentDocument, + count: 3, + }; + + // Then + assert_eq!( + serde_json::to_value(&row).unwrap(), + serde_json::json!({ + "client_side_key": "client", + "resource": "environment-document", + "count": 3, + }) + ); + } +} diff --git a/tests/test_usage_tracking.rs b/tests/test_usage_tracking.rs new file mode 100644 index 0000000..61a39d1 --- /dev/null +++ b/tests/test_usage_tracking.rs @@ -0,0 +1,423 @@ +use axum_test::TestServer; +use edge_proxy::config::settings::{AppSettings, EnvironmentKeyPair}; +use edge_proxy::routes::create_router; +use edge_proxy::services::EnvironmentService; +use serde_json::{Value, json}; +use wiremock::matchers::{header, method, path}; +use wiremock::{Mock, MockServer, Request, ResponseTemplate}; + +const PROXY_KEY: &str = "pk.test_proxy_key"; +const CLIENT_KEY: &str = "config_client_key"; +const SERVER_KEY: &str = "ser.config_key"; + +fn settings(api_url: &str, proxy_key: Option<&str>, pairs: Vec) -> AppSettings { + AppSettings { + environment_key_pairs: pairs, + proxy_key: proxy_key.map(str::to_string), + api_url: api_url.to_string(), + ..AppSettings::default() + } +} + +fn config_body() -> Value { + json!([{ + "id": 30, + "name": "Test Environment", + "client_side_key": CLIENT_KEY, + "server_side_keys": [ + {"key": SERVER_KEY, "active": true, "expires_at": null} + ], + "updated_at": "2026-08-15T08:57:43.311081Z", + "project_id": 35, + "organisation_id": 82, + }]) +} + +fn document_body() -> Value { + json!({ + "id": 1, + "api_key": CLIENT_KEY, + "name": "Test", + "updated_at": "2026-08-22T00:00:00Z", + "allow_client_traits": true, + "hide_sensitive_data": false, + "hide_disabled_flags": null, + "use_identity_composite_key_for_hashing": true, + "use_identity_overrides_in_local_eval": true, + "project": { + "id": 1, + "name": "project-1", + "hide_disabled_flags": false, + "segments": [], + "server_key_only_feature_ids": [], + "organisation": { + "id": 1, + "name": "org-1", + "feature_analytics": false, + "persist_trait_data": true, + "stop_serving_flags": false, + }, + }, + "feature_states": [ + { + "multivariate_feature_state_values": [], + "feature_state_value": "config_value", + "feature": {"id": 1, "name": "config_flag", "type": "STANDARD"}, + "enabled": true, + "featurestate_uuid": "fs-uuid-1", + } + ], + "identity_overrides": [], + }) +} + +async fn mount_config(mock_server: &MockServer) { + Mock::given(method("GET")) + .and(path("/proxy/config/")) + .and(header("X-Proxy-Key", PROXY_KEY)) + .respond_with(ResponseTemplate::new(200).set_body_json(config_body())) + .mount(mock_server) + .await; +} + +async fn mount_document(mock_server: &MockServer) { + Mock::given(method("GET")) + .and(path("/environment-document/")) + .respond_with(ResponseTemplate::new(200).set_body_json(document_body())) + .mount(mock_server) + .await; +} + +async fn mount_usage(mock_server: &MockServer, status: u16, up_to: Option) { + let mut mock = Mock::given(method("POST")) + .and(path("/proxy/usage/")) + .and(header("X-Proxy-Key", PROXY_KEY)) + .respond_with(ResponseTemplate::new(status)); + if let Some(n) = up_to { + mock = mock.up_to_n_times(n); + } + mock.mount(mock_server).await; +} + +async fn requests_to(mock_server: &MockServer, url_path: &str) -> Vec { + mock_server + .received_requests() + .await + .unwrap() + .into_iter() + .filter(|request| request.url.path() == url_path) + .collect() +} + +/// The usage rows of the request body, sorted by resource for +/// order-independent assertions. +fn usage_rows(request: &Request) -> Vec { + let mut rows: Vec = serde_json::from_slice(&request.body).unwrap(); + rows.sort_by_key(|row| row["resource"].as_str().unwrap().to_string()); + rows +} + +#[tokio::test] +async fn test_served_requests_flush_aggregated_usage() { + // Given a discovered environment served through the full router + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + mount_document(&mock_server).await; + mount_usage(&mock_server, 204, None).await; + let (app, service) = create_router(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + service.refresh_environment_caches().await; + let server = TestServer::new(app).unwrap(); + + // When SDK traffic arrives under both keys, then usage is flushed + server + .get("/api/v1/flags") + .add_header("X-Environment-Key", CLIENT_KEY) + .await + .assert_status_ok(); + server + .get("/api/v1/flags") + .add_header("X-Environment-Key", CLIENT_KEY) + .await + .assert_status_ok(); + server + .get("/api/v1/identities") + .add_query_param("identifier", "user_1") + .add_header("X-Environment-Key", CLIENT_KEY) + .await + .assert_status_ok(); + server + .post("/api/v1/identities") + .json(&json!({"identifier": "user_2"})) + .add_header("X-Environment-Key", CLIENT_KEY) + .await + .assert_status_ok(); + server + .get("/api/v1/environment-document") + .add_header("X-Environment-Key", SERVER_KEY) + .await + .assert_status_ok(); + let flushed = service.flush_usage().await; + + // Then one POST reports everything, keyed by the client key even for + // requests that presented the server key + assert!(flushed); + let posts = requests_to(&mock_server, "/proxy/usage/").await; + assert_eq!(posts.len(), 1); + assert_eq!( + usage_rows(&posts[0]), + vec![ + json!({"client_side_key": CLIENT_KEY, "resource": "environment-document", "count": 1}), + json!({"client_side_key": CLIENT_KEY, "resource": "flags", "count": 2}), + json!({"client_side_key": CLIENT_KEY, "resource": "identities", "count": 2}), + ] + ); +} + +#[tokio::test] +async fn test_unresolved_keys_are_never_counted() { + // Given a proxy serving one environment + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + mount_document(&mock_server).await; + let (app, service) = create_router(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + service.refresh_environment_caches().await; + let server = TestServer::new(app).unwrap(); + + // When requests present an unknown key or none at all + server + .get("/api/v1/flags") + .add_header("X-Environment-Key", "unknown_key") + .await + .assert_status_unauthorized(); + server + .get("/api/v1/flags") + .await + .assert_status_unauthorized(); + + // Then there is nothing to flush and no request is made + assert!(service.flush_usage().await); + assert!(requests_to(&mock_server, "/proxy/usage/").await.is_empty()); +} + +#[tokio::test] +async fn test_failed_flush_merges_counts_into_the_next() { + // Given a served request and a usage endpoint that fails once + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + mount_document(&mock_server).await; + mount_usage(&mock_server, 500, Some(1)).await; + mount_usage(&mock_server, 204, None).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + service.refresh_environment_caches().await; + service + .get_flags_response_data(CLIENT_KEY, None) + .await + .unwrap(); + + // When the first flush fails and another request is served + assert!(!service.flush_usage().await); + service + .get_flags_response_data(CLIENT_KEY, None) + .await + .unwrap(); + + // Then the next flush carries both counts — nothing lost, nothing + // double-counted + assert!(service.flush_usage().await); + let posts = requests_to(&mock_server, "/proxy/usage/").await; + assert_eq!(posts.len(), 2); + assert_eq!( + usage_rows(&posts[1]), + vec![json!({"client_side_key": CLIENT_KEY, "resource": "flags", "count": 2})] + ); +} + +#[tokio::test] +async fn test_flush_without_proxy_key_is_inert() { + // Given a statically configured proxy with no proxy key + let mock_server = MockServer::start().await; + mount_document(&mock_server).await; + let service = EnvironmentService::new(settings( + &mock_server.uri(), + None, + vec![EnvironmentKeyPair { + client_side_key: CLIENT_KEY.to_string(), + server_side_key: SERVER_KEY.to_string(), + }], + )); + service.refresh_environment_caches().await; + service + .get_flags_response_data(CLIENT_KEY, None) + .await + .unwrap(); + + // When / Then: flushing succeeds without reporting anything + assert!(service.flush_usage().await); + assert!(requests_to(&mock_server, "/proxy/usage/").await.is_empty()); +} + +#[tokio::test] +async fn test_rejected_flush_drops_the_batch_instead_of_retrying_it() { + // Given a served request and a usage endpoint that rejects the batch + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + mount_document(&mock_server).await; + mount_usage(&mock_server, 400, None).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + service.refresh_environment_caches().await; + service + .get_flags_response_data(CLIENT_KEY, None) + .await + .unwrap(); + + // When the flush is rejected + assert!(!service.flush_usage().await); + + // Then the rows are dropped, not resent forever: the next flush has + // nothing to send + assert!(service.flush_usage().await); + assert_eq!(requests_to(&mock_server, "/proxy/usage/").await.len(), 1); +} + +#[tokio::test] +async fn test_flush_chunks_batches_to_the_server_cap() { + // Given served requests for more environments than one batch may hold + let environments: Vec = (0..1001) + .map(|n| { + json!({ + "id": n, + "name": format!("env {n}"), + "client_side_key": format!("client_{n}"), + "server_side_keys": [ + {"key": format!("ser.key_{n}"), "active": true, "expires_at": null} + ], + "updated_at": "2026-08-15T08:57:43.311081Z", + "project_id": 35, + "organisation_id": 82, + }) + }) + .collect(); + let mock_server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/proxy/config/")) + .respond_with(ResponseTemplate::new(200).set_body_json(Value::Array(environments))) + .mount(&mock_server) + .await; + mount_document(&mock_server).await; + mount_usage(&mock_server, 204, None).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + service.refresh_environment_caches().await; + for n in 0..1001 { + service + .get_flags_response_data(&format!("client_{n}"), None) + .await + .unwrap(); + } + + // When + assert!(service.flush_usage().await); + + // Then the rows arrive split across two accepted requests + let posts = requests_to(&mock_server, "/proxy/usage/").await; + let row_counts: Vec = posts.iter().map(|post| usage_rows(post).len()).collect(); + assert_eq!(row_counts, vec![1000, 1]); +} + +#[tokio::test] +async fn test_static_environment_usage_is_neither_counted_nor_marked() { + // Given a proxy serving a static environment alongside a discovered one + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + mount_document(&mock_server).await; + mount_usage(&mock_server, 204, None).await; + let service = EnvironmentService::new(settings( + &mock_server.uri(), + Some(PROXY_KEY), + vec![EnvironmentKeyPair { + client_side_key: "static_client".to_string(), + server_side_key: "ser.static_key".to_string(), + }], + )); + service.refresh_environment_caches().await; + + // When both environments serve a request, and usage is flushed + service + .get_flags_response_data("static_client", None) + .await + .unwrap(); + service + .get_flags_response_data(CLIENT_KEY, None) + .await + .unwrap(); + assert!(service.flush_usage().await); + + // Then the static environment keeps its old billing: its document + // fetch is not marked as the proxy's own, and it is not reported + let fetches = requests_to(&mock_server, "/environment-document/").await; + let marked: Vec = fetches + .iter() + .map(|request| request.headers.contains_key("X-Proxy-Key")) + .collect(); + let static_fetches: Vec<&Request> = fetches + .iter() + .filter(|request| request.headers["X-Environment-Key"] == "ser.static_key") + .collect(); + assert!(!static_fetches.is_empty()); + assert!( + static_fetches + .iter() + .all(|request| !request.headers.contains_key("X-Proxy-Key")) + ); + assert!(marked.contains(&true)); + let posts = requests_to(&mock_server, "/proxy/usage/").await; + assert_eq!( + usage_rows(&posts[0]), + vec![json!({"client_side_key": CLIENT_KEY, "resource": "flags", "count": 1})] + ); +} + +#[tokio::test] +async fn test_document_fetch_carries_the_proxy_key() { + // Given a document endpoint that only answers requests marked as the + // proxy's own + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + Mock::given(method("GET")) + .and(path("/environment-document/")) + .and(header("X-Proxy-Key", PROXY_KEY)) + .respond_with(ResponseTemplate::new(200).set_body_json(document_body())) + .mount(&mock_server) + .await; + let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + + // When / Then + assert!(service.refresh_environment_caches().await); + assert!(service.get_environment(CLIENT_KEY).await.is_ok()); +} + +#[tokio::test] +async fn test_static_document_fetch_omits_the_proxy_key() { + // Given a statically configured proxy with no proxy key + let mock_server = MockServer::start().await; + mount_document(&mock_server).await; + let service = EnvironmentService::new(settings( + &mock_server.uri(), + None, + vec![EnvironmentKeyPair { + client_side_key: CLIENT_KEY.to_string(), + server_side_key: SERVER_KEY.to_string(), + }], + )); + + // When + assert!(service.refresh_environment_caches().await); + + // Then the document fetch is byte-identical to today's + let fetches = requests_to(&mock_server, "/environment-document/").await; + assert!(!fetches.is_empty()); + assert!( + fetches + .iter() + .all(|request| !request.headers.contains_key("X-Proxy-Key")) + ); +} From 3a3c672e24aa45c89264ae0a2d55b8fcde4a9a7c Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 5 Sep 2026 15:54:30 +0530 Subject: [PATCH 02/22] refactor: parking_lot mutex for usage counts, name the map by its key --- src/usage.rs | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/src/usage.rs b/src/usage.rs index c5c365e..da6c0b2 100644 --- a/src/usage.rs +++ b/src/usage.rs @@ -1,6 +1,7 @@ -use serde::Serialize; use std::collections::HashMap; -use std::sync::Mutex; + +use parking_lot::Mutex; +use serde::Serialize; /// An SDK endpoint as the usage endpoint names it. #[derive(Serialize, Clone, Copy, PartialEq, Eq, Hash, Debug)] @@ -23,25 +24,26 @@ pub struct UsageRow { /// resource. Keyed by resolved client keys only, so the map is bounded by /// the served environment set. /// -/// Uses `std::sync::Mutex`: every operation writes, and guards are held -/// only for a map operation, never across an await. +/// Uses `parking_lot::Mutex`, not tokio's: every operation writes, and +/// guards are held only for a map operation, never across an await. #[derive(Default)] pub struct UsageCounts { - by_key: Mutex>, + count_by_environment_and_resource: Mutex>, } impl UsageCounts { pub fn increment(&self, client_key: &str, resource: Resource) { - let mut by_key = self.by_key.lock().expect("usage counts lock poisoned"); - *by_key + let mut count_by_environment_and_resource = self.count_by_environment_and_resource.lock(); + *count_by_environment_and_resource .entry((client_key.to_string(), resource)) .or_default() += 1; } /// Take everything counted so far, leaving the map empty. pub fn drain(&self) -> Vec { - let by_key = std::mem::take(&mut *self.by_key.lock().expect("usage counts lock poisoned")); - by_key + let count_by_environment_and_resource = + std::mem::take(&mut *self.count_by_environment_and_resource.lock()); + count_by_environment_and_resource .into_iter() .map(|((client_side_key, resource), count)| UsageRow { client_side_key, @@ -54,9 +56,9 @@ impl UsageCounts { /// Add drained counts back, on top of anything counted since — for /// when a flush fails and the rows must survive until the next one. pub fn merge(&self, counts: Vec) { - let mut by_key = self.by_key.lock().expect("usage counts lock poisoned"); + let mut count_by_environment_and_resource = self.count_by_environment_and_resource.lock(); for row in counts { - let entry = by_key + let entry = count_by_environment_and_resource .entry((row.client_side_key, row.resource)) .or_default(); *entry = entry.saturating_add(row.count); From 174e1c1a79918c458993d3741919cf540a2b4c07 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 5 Sep 2026 16:03:06 +0530 Subject: [PATCH 03/22] fix: count usage from the response, only for 2xx Counting inside resolve_key billed requests that then failed (404 on an unknown feature, 503 before the document loaded). A router middleware now counts after the handler answers, and only when it answered 2xx. --- src/routes/mod.rs | 19 +++++--- src/routes/usage.rs | 43 ++++++++++++++++++ src/services/environment.rs | 41 +++++++---------- tests/test_usage_tracking.rs | 87 ++++++++++++++++++++++++------------ 4 files changed, 133 insertions(+), 57 deletions(-) create mode 100644 src/routes/usage.rs diff --git a/src/routes/mod.rs b/src/routes/mod.rs index fc419b0..cb5a29b 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -4,17 +4,22 @@ pub mod extractors; pub mod flags; pub mod health; pub mod identities; +pub mod usage; use crate::config::AppSettings; use crate::services::EnvironmentService; use axum::{ Router, - middleware::map_response, + middleware::{from_fn_with_state, map_response}, routing::{get, post}, }; use std::sync::Arc; use tower_http::{compression::CompressionLayer, normalize_path::NormalizePath, trace::TraceLayer}; +const FLAGS_PATH: &str = "/api/v1/flags"; +const IDENTITIES_PATH: &str = "/api/v1/identities"; +const ENVIRONMENT_DOCUMENT_PATH: &str = "/api/v1/environment-document"; + pub fn create_router(settings: AppSettings) -> (Router, Arc) { let cors = cors::layer(&settings.allow_origins); let environment_service = Arc::new(EnvironmentService::new(settings)); @@ -26,16 +31,20 @@ pub fn create_router(settings: AppSettings) -> (Router, Arc) .route("/proxy/health/readiness", get(health::health_check)) .route("/proxy/health/liveness", get(health::liveness_check)) // Flags routes (with and without trailing slash) - .route("/api/v1/flags", get(flags::get_flags)) + .route(FLAGS_PATH, get(flags::get_flags)) // Identities routes (with and without trailing slash) - .route("/api/v1/identities", get(identities::get_identities)) - .route("/api/v1/identities", post(identities::post_identities)) + .route(IDENTITIES_PATH, get(identities::get_identities)) + .route(IDENTITIES_PATH, post(identities::post_identities)) // Environment document route .route( - "/api/v1/environment-document", + ENVIRONMENT_DOCUMENT_PATH, get(environment_document::get_environment_document), ) // Middleware layers + .layer(from_fn_with_state( + environment_service.clone(), + usage::track_usage, + )) .layer(CompressionLayer::new()) .layer(cors) .layer(map_response(cors::merge_vary)) diff --git a/src/routes/usage.rs b/src/routes/usage.rs new file mode 100644 index 0000000..b5d4926 --- /dev/null +++ b/src/routes/usage.rs @@ -0,0 +1,43 @@ +use super::{ENVIRONMENT_DOCUMENT_PATH, FLAGS_PATH, IDENTITIES_PATH}; +use crate::services::EnvironmentService; +use crate::usage::Resource; +use axum::extract::{MatchedPath, Request, State}; +use axum::middleware::Next; +use axum::response::Response; +use std::sync::Arc; + +/// Count a request once the handler has answered: only a 2xx is usage. +pub async fn track_usage( + State(service): State>, + request: Request, + next: Next, +) -> Response { + let resource = request + .extensions() + .get::() + .and_then(|path| resource_for(path.as_str())); + let environment_key = request + .headers() + .get("X-Environment-Key") + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + + let response = next.run(request).await; + + if !response.status().is_success() { + return response; + } + if let (Some(resource), Some(environment_key)) = (resource, environment_key) { + service.track_usage(&environment_key, resource); + } + response +} + +fn resource_for(route: &str) -> Option { + match route { + FLAGS_PATH => Some(Resource::Flags), + IDENTITIES_PATH => Some(Resource::Identities), + ENVIRONMENT_DOCUMENT_PATH => Some(Resource::EnvironmentDocument), + _ => None, + } +} diff --git a/src/services/environment.rs b/src/services/environment.rs index 5f786ba..b559ebc 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -156,25 +156,23 @@ impl EnvironmentService { Ok(response.json().await?) } - /// Resolve a presented key, counting the request for usage reporting. - /// Every SDK entry point resolves through here, so a served request - /// cannot be missed. - fn resolve_key( - &self, - environment_key: &str, - resource: Resource, - ) -> Result> { - let keys = self - .environments + fn resolve_key(&self, environment_key: &str) -> Result> { + self.environments .resolve(environment_key) - .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string()))?; - self.track_usage(&keys.client_key, resource); - Ok(keys) + .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string())) } - fn track_usage(&self, client_key: &str, resource: Resource) { - if self.settings.proxy_key.is_some() && !self.environments.is_static(client_key) { - self.usage.increment(client_key, resource); + /// Count a served request against its environment. Static environments + /// stay on core's own billing, so they are not reported. + pub fn track_usage(&self, environment_key: &str, resource: Resource) { + if self.settings.proxy_key.is_none() { + return; + } + let Some(keys) = self.environments.resolve(environment_key) else { + return; + }; + if !self.environments.is_static(&keys.client_key) { + self.usage.increment(&keys.client_key, resource); } } @@ -290,11 +288,7 @@ impl EnvironmentService { } pub async fn get_environment(&self, environment_key: &str) -> Result> { - // Lookup, not an SDK entry point: callers count via resolve_key. - let keys = self - .environments - .resolve(environment_key) - .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string()))?; + let keys = self.resolve_key(environment_key)?; // Documents are cached under the client key, whichever key was presented self.cache @@ -305,7 +299,6 @@ impl EnvironmentService { /// Get pre-serialized environment document bytes pub async fn get_environment_bytes(&self, environment_key: &str) -> Result> { - self.resolve_key(environment_key, Resource::EnvironmentDocument)?; let document = self.get_environment(environment_key).await?; Ok(serde_json::to_vec(&*document)?.into()) } @@ -330,7 +323,7 @@ impl EnvironmentService { ) -> Result> { // TODO: serve server-side keys here, see // https://github.com/Flagsmith/edge-proxy-rs/issues/24 - self.resolve_key(environment_key, Resource::Flags)?; + self.resolve_key(environment_key)?; let context = self .cache @@ -374,7 +367,7 @@ impl EnvironmentService { identity: &IdentityWithTraits, environment_key: &str, ) -> Result { - self.resolve_key(environment_key, Resource::Identities)?; + self.resolve_key(environment_key)?; // Get pre-computed context from cache let context = self diff --git a/tests/test_usage_tracking.rs b/tests/test_usage_tracking.rs index 61a39d1..ef3da6e 100644 --- a/tests/test_usage_tracking.rs +++ b/tests/test_usage_tracking.rs @@ -1,7 +1,9 @@ +use axum::http::StatusCode; use axum_test::TestServer; use edge_proxy::config::settings::{AppSettings, EnvironmentKeyPair}; use edge_proxy::routes::create_router; use edge_proxy::services::EnvironmentService; +use edge_proxy::usage::Resource; use serde_json::{Value, json}; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, Request, ResponseTemplate}; @@ -199,6 +201,56 @@ async fn test_unresolved_keys_are_never_counted() { assert!(requests_to(&mock_server, "/proxy/usage/").await.is_empty()); } +#[tokio::test] +async fn test_failed_requests_are_not_counted() { + // Given one environment that serves and one whose document never loaded + let mock_server = MockServer::start().await; + let mut environments = config_body(); + environments.as_array_mut().unwrap().push(json!({ + "id": 31, + "name": "Broken Environment", + "client_side_key": "broken_client", + "server_side_keys": [ + {"key": "ser.broken_key", "active": true, "expires_at": null} + ], + "updated_at": "2026-08-15T08:57:43.311081Z", + "project_id": 35, + "organisation_id": 82, + })); + Mock::given(method("GET")) + .and(path("/proxy/config/")) + .respond_with(ResponseTemplate::new(200).set_body_json(environments)) + .mount(&mock_server) + .await; + Mock::given(method("GET")) + .and(path("/environment-document/")) + .and(header("X-Environment-Key", "ser.broken_key")) + .respond_with(ResponseTemplate::new(500)) + .mount(&mock_server) + .await; + mount_document(&mock_server).await; + let (app, service) = create_router(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + service.refresh_environment_caches().await; + let server = TestServer::new(app).unwrap(); + + // When requests fail after their key resolved + server + .get("/api/v1/flags") + .add_query_param("feature", "missing") + .add_header("X-Environment-Key", CLIENT_KEY) + .await + .assert_status(StatusCode::NOT_FOUND); + server + .get("/api/v1/flags") + .add_header("X-Environment-Key", "broken_client") + .await + .assert_status(StatusCode::SERVICE_UNAVAILABLE); + + // Then nothing is reported + assert!(service.flush_usage().await); + assert!(requests_to(&mock_server, "/proxy/usage/").await.is_empty()); +} + #[tokio::test] async fn test_failed_flush_merges_counts_into_the_next() { // Given a served request and a usage endpoint that fails once @@ -209,17 +261,11 @@ async fn test_failed_flush_merges_counts_into_the_next() { mount_usage(&mock_server, 204, None).await; let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); service.refresh_environment_caches().await; - service - .get_flags_response_data(CLIENT_KEY, None) - .await - .unwrap(); + service.track_usage(CLIENT_KEY, Resource::Flags); // When the first flush fails and another request is served assert!(!service.flush_usage().await); - service - .get_flags_response_data(CLIENT_KEY, None) - .await - .unwrap(); + service.track_usage(CLIENT_KEY, Resource::Flags); // Then the next flush carries both counts — nothing lost, nothing // double-counted @@ -246,10 +292,7 @@ async fn test_flush_without_proxy_key_is_inert() { }], )); service.refresh_environment_caches().await; - service - .get_flags_response_data(CLIENT_KEY, None) - .await - .unwrap(); + service.track_usage(CLIENT_KEY, Resource::Flags); // When / Then: flushing succeeds without reporting anything assert!(service.flush_usage().await); @@ -265,10 +308,7 @@ async fn test_rejected_flush_drops_the_batch_instead_of_retrying_it() { mount_usage(&mock_server, 400, None).await; let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); service.refresh_environment_caches().await; - service - .get_flags_response_data(CLIENT_KEY, None) - .await - .unwrap(); + service.track_usage(CLIENT_KEY, Resource::Flags); // When the flush is rejected assert!(!service.flush_usage().await); @@ -308,10 +348,7 @@ async fn test_flush_chunks_batches_to_the_server_cap() { let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); service.refresh_environment_caches().await; for n in 0..1001 { - service - .get_flags_response_data(&format!("client_{n}"), None) - .await - .unwrap(); + service.track_usage(&format!("client_{n}"), Resource::Flags); } // When @@ -341,14 +378,8 @@ async fn test_static_environment_usage_is_neither_counted_nor_marked() { service.refresh_environment_caches().await; // When both environments serve a request, and usage is flushed - service - .get_flags_response_data("static_client", None) - .await - .unwrap(); - service - .get_flags_response_data(CLIENT_KEY, None) - .await - .unwrap(); + service.track_usage("static_client", Resource::Flags); + service.track_usage(CLIENT_KEY, Resource::Flags); assert!(service.flush_usage().await); // Then the static environment keeps its old billing: its document From 7c2c853bfc71efa646eb7f156ad4cc8bce020c6e Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 5 Sep 2026 16:19:26 +0530 Subject: [PATCH 04/22] fix: refuse cross-origin pagination links The next page request carries the environment and proxy keys, so a Link header must not be able to send them off api_url's origin. --- src/services/environment.rs | 63 +++++++++++++++++++++++++++---------- 1 file changed, 47 insertions(+), 16 deletions(-) diff --git a/src/services/environment.rs b/src/services/environment.rs index b559ebc..853d597 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -251,7 +251,7 @@ impl EnvironmentService { response.error_for_status_ref()?; - let next_link = parse_next_link(response.headers(), &self.settings.api_url); + let next_link = parse_next_link(response.headers(), &self.settings.api_url)?; let body: serde_json::Value = response.json().await?; match document.as_mut() { @@ -537,11 +537,18 @@ fn merge_paginated_overrides(base: &mut serde_json::Value, page: serde_json::Val /// Parse the `Link` response header for the next-page URL (RFC 5988). /// /// Returns an absolute URL, resolving relative targets against `api_url`. -fn parse_next_link(headers: &HeaderMap, api_url: &str) -> Option { - let base = Url::parse(api_url).ok()?; +/// The `rel=next` pagination target, if any. A target off `api_url`'s +/// origin is an error: the next request carries the environment and proxy +/// keys, and a Link header must not be able to send them elsewhere. +fn parse_next_link(headers: &HeaderMap, api_url: &str) -> Result> { + let Ok(base) = Url::parse(api_url) else { + return Ok(None); + }; for header_value in headers.get_all(reqwest::header::LINK).iter() { - let raw = header_value.to_str().ok()?; + let Ok(raw) = header_value.to_str() else { + return Ok(None); + }; for segment in raw.split(',') { let segment = segment.trim(); let target = match (segment.find('<'), segment.find('>')) { @@ -552,12 +559,18 @@ fn parse_next_link(headers: &HeaderMap, api_url: &str) -> Option { if !is_next_rel(params) { continue; } - if let Ok(absolute) = base.join(target) { - return Some(absolute.into()); + let Ok(absolute) = base.join(target) else { + continue; + }; + if absolute.origin() != base.origin() { + return Err(EdgeProxyError::ServiceUnavailable(format!( + "refusing cross-origin pagination link {absolute}" + ))); } + return Ok(Some(absolute.into())); } } - None + Ok(None) } fn is_next_rel(params: &str) -> bool { @@ -595,40 +608,58 @@ mod tests { ); let next = parse_next_link(&headers, "https://edge.api.flagsmith.com/api/v1").unwrap(); assert_eq!( - next, + next.unwrap(), "https://edge.api.flagsmith.com/api/v1/environment-document/?page_id=identity_override%3A1%3Aabc" ); } #[test] - fn parse_next_link_absolute() { - let headers = - link("; rel=\"next\""); + fn parse_next_link_absolute_same_origin() { + let headers = link( + "; rel=\"next\"", + ); let next = parse_next_link(&headers, "https://edge.api.flagsmith.com/api/v1").unwrap(); assert_eq!( - next, - "https://example.test/api/v1/environment-document/?page_id=x" + next.unwrap(), + "https://edge.api.flagsmith.com/api/v1/environment-document/?page_id=x" ); } + #[test] + fn parse_next_link_rejects_cross_origin() { + let headers = + link("; rel=\"next\""); + assert!(parse_next_link(&headers, "https://edge.api.flagsmith.com/api/v1").is_err()); + } + #[test] fn parse_next_link_picks_next_among_multiple_rels() { let headers = link("; rel=\"prev\", ; rel=\"next\""); let next = parse_next_link(&headers, "https://edge.api.flagsmith.com/api/v1").unwrap(); - assert_eq!(next, "https://edge.api.flagsmith.com/api/v1/page/next"); + assert_eq!( + next.unwrap(), + "https://edge.api.flagsmith.com/api/v1/page/next" + ); } #[test] fn parse_next_link_returns_none_when_only_other_rels() { let headers = link("; rel=\"prev\", ; rel=\"self\""); - assert!(parse_next_link(&headers, "https://edge.api.flagsmith.com/api/v1").is_none()); + assert!( + parse_next_link(&headers, "https://edge.api.flagsmith.com/api/v1") + .unwrap() + .is_none() + ); } #[test] fn parse_next_link_handles_unquoted_rel() { let headers = link("; rel=next"); let next = parse_next_link(&headers, "https://edge.api.flagsmith.com/api/v1").unwrap(); - assert_eq!(next, "https://edge.api.flagsmith.com/api/v1/page/next"); + assert_eq!( + next.unwrap(), + "https://edge.api.flagsmith.com/api/v1/page/next" + ); } #[test] From 027803f683d57d2225aa55ac1a4c65d4152c4046 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 5 Sep 2026 16:48:50 +0530 Subject: [PATCH 05/22] fix: resend a failed usage batch unchanged under an idempotency key A batch that failed with 5xx or a network error may still have been processed. It is now kept intact and resent with the same Idempotency-Key so the server can recognise it, instead of being merged into the next flush and counted twice. Nothing new is drained while a batch is pending; counts keep aggregating in the bounded map. --- Cargo.lock | 2 + Cargo.toml | 1 + src/services/environment.rs | 110 +++++++++++++++++++++-------------- src/usage.rs | 39 +++++-------- tests/test_usage_tracking.rs | 37 ++++++++---- 5 files changed, 108 insertions(+), 81 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bb2d529..89aef42 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -529,6 +529,7 @@ dependencies = [ "tracing", "tracing-subscriber", "tracing-test", + "uuid 1.19.0", "validator", "wiremock", ] @@ -2561,6 +2562,7 @@ version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e2e054861b4bd027cd373e18e8d8d8e6548085000e41290d95ce0c373a654b4a" dependencies = [ + "getrandom 0.3.4", "js-sys", "wasm-bindgen", ] diff --git a/Cargo.toml b/Cargo.toml index a6dde39..7cc2a07 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,6 +10,7 @@ axum = { version = "0.8", features = ["macros"] } tokio = { version = "1", features = ["full"] } tower = "0.5" parking_lot = "0.12" +uuid = { version = "1", features = ["v4"] } tower-http = { version = "0.6", features = ["compression-gzip", "cors", "trace", "normalize-path"] } reqwest = { version = "0.13", features = ["json", "gzip"] } serde = { version = "1", features = ["derive"] } diff --git a/src/services/environment.rs b/src/services/environment.rs index 853d597..53a5dd3 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -6,11 +6,12 @@ use crate::models::{ APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, }; use crate::services::feature_utils::filter_out_server_key_only_flag_results; -use crate::usage::{Resource, UsageCounts, UsageRow}; +use crate::usage::{Resource, UsageBatch, UsageCounts}; use chrono::{DateTime, Utc}; use flagsmith_flag_engine::engine::get_evaluation_result; use flagsmith_flag_engine::engine_eval::{FlagResult, add_identity_to_context}; use flagsmith_flag_engine::identities::Trait as FlagsmithTrait; +use parking_lot::Mutex; use reqwest::header::HeaderMap; use reqwest::{Client, Url}; use std::sync::Arc; @@ -25,6 +26,13 @@ pub struct EnvironmentService { pub last_updated_at: Arc>>>, environments: EnvironmentIndex, usage: UsageCounts, + pending_usage: Mutex>, +} + +enum UsageOutcome { + Accepted, + Rejected, + Failed, } impl EnvironmentService { @@ -44,6 +52,7 @@ impl EnvironmentService { last_updated_at: Arc::new(RwLock::new(None)), environments, usage: UsageCounts::default(), + pending_usage: Mutex::default(), } } @@ -423,62 +432,77 @@ impl EnvironmentService { } } - /// The usage endpoint's batch cap — MAX_USAGE_ROWS in the edge_proxy - /// app. Flushes are chunked to it so a large environment set can - /// never be rejected outright. + /// The usage endpoint's batch cap — MAX_USAGE_ROWS in the + /// edge_control_plane app. Flushes are chunked to it so a large + /// environment set can never be rejected outright. const MAX_ROWS_PER_FLUSH: usize = 1000; - /// Report the counts accumulated since the last flush to the usage - /// endpoint, in chunks the server accepts. A rejected (4xx) chunk is - /// dropped — retrying cannot heal a rejection, and losing one window - /// beats resending a poisoned batch forever. Any other failure keeps - /// the chunk for the next flush. Returns false when any chunk was - /// not accepted. + /// Returns false when any batch was not accepted. pub async fn flush_usage(&self) -> bool { let Some(proxy_key) = &self.settings.proxy_key else { return true; }; - let mut rows = self.usage.drain(); - let url = format!("{}/proxy/usage/", self.settings.api_url); - let mut all_success = true; - while !rows.is_empty() { - let chunk: Vec = rows - .drain(..rows.len().min(Self::MAX_ROWS_PER_FLUSH)) - .collect(); - let result = self - .client - .post(&url) - .header("X-Proxy-Key", proxy_key) - .json(&chunk) - .send() - .await; - match result { - Ok(response) if response.status().is_success() => {} - Ok(response) if response.status().is_client_error() => { - error!( - "Usage report rejected with {}: dropping {} rows", - response.status(), - chunk.len() - ); - all_success = false; - } - Ok(response) => { - error!("Failed to report usage: {}", response.status()); - self.usage.merge(chunk); - all_success = false; - } - Err(e) => { - error!("Failed to report usage: {}", e); - self.usage.merge(chunk); + // A batch that failed may still have been processed, so it is + // resent unchanged under the same key and nothing new is drained + // until it is through. Meanwhile counts keep aggregating in the + // map, which is bounded by the served environment set. + let mut batches = std::mem::take(&mut *self.pending_usage.lock()); + if batches.is_empty() { + let mut rows = self.usage.drain(); + while !rows.is_empty() { + let chunk = rows.drain(..rows.len().min(Self::MAX_ROWS_PER_FLUSH)); + batches.push(UsageBatch::new(chunk.collect())); + } + } + + let mut all_success = true; + for batch in batches { + match self.post_usage(proxy_key, &batch).await { + UsageOutcome::Accepted => {} + UsageOutcome::Rejected => all_success = false, + UsageOutcome::Failed => { + self.pending_usage.lock().push(batch); all_success = false; } } } - all_success } + async fn post_usage(&self, proxy_key: &str, batch: &UsageBatch) -> UsageOutcome { + let url = format!("{}/proxy/usage/", self.settings.api_url); + let result = self + .client + .post(&url) + .header("X-Proxy-Key", proxy_key) + .header("Idempotency-Key", &batch.id) + .json(&batch.rows) + .send() + .await; + match result { + Ok(response) if response.status().is_success() => UsageOutcome::Accepted, + // Retrying cannot heal a rejection: drop the batch rather than + // resend it forever. + Ok(response) if response.status().is_client_error() => { + error!( + "Usage report rejected with {}: dropping {} rows", + response.status(), + batch.rows.len() + ); + UsageOutcome::Rejected + } + Ok(response) => { + error!("Failed to report usage: {}", response.status()); + UsageOutcome::Failed + } + Err(e) => { + error!("Failed to report usage: {}", e); + UsageOutcome::Failed + } + } + } + pub async fn flush_usage_periodically(self: Arc) { if self.settings.proxy_key.is_none() { return; diff --git a/src/usage.rs b/src/usage.rs index da6c0b2..cfca73c 100644 --- a/src/usage.rs +++ b/src/usage.rs @@ -52,16 +52,20 @@ impl UsageCounts { }) .collect() } +} - /// Add drained counts back, on top of anything counted since — for - /// when a flush fails and the rows must survive until the next one. - pub fn merge(&self, counts: Vec) { - let mut count_by_environment_and_resource = self.count_by_environment_and_resource.lock(); - for row in counts { - let entry = count_by_environment_and_resource - .entry((row.client_side_key, row.resource)) - .or_default(); - *entry = entry.saturating_add(row.count); +/// One `POST /proxy/usage/` body with the idempotency key it is sent +/// under, so a retry after a lost response is recognised, not recounted. +pub struct UsageBatch { + pub id: String, + pub rows: Vec, +} + +impl UsageBatch { + pub fn new(rows: Vec) -> Self { + Self { + id: uuid::Uuid::new_v4().to_string(), + rows, } } } @@ -113,23 +117,6 @@ mod tests { assert!(counts.drain().is_empty()); } - #[test] - fn merge_adds_drained_counts_back() { - // Given a drained batch and a request counted since - let counts = UsageCounts::default(); - counts.increment("client", Resource::Flags); - let drained = counts.drain(); - counts.increment("client", Resource::Flags); - - // When - counts.merge(drained); - - // Then - let rows = counts.drain(); - assert_eq!(rows.len(), 1); - assert_eq!(rows[0].count, 2); - } - #[test] fn usage_row_serializes_to_the_contract_shape() { // Given diff --git a/tests/test_usage_tracking.rs b/tests/test_usage_tracking.rs index ef3da6e..982e4f2 100644 --- a/tests/test_usage_tracking.rs +++ b/tests/test_usage_tracking.rs @@ -161,10 +161,11 @@ async fn test_served_requests_flush_aggregated_usage() { let flushed = service.flush_usage().await; // Then one POST reports everything, keyed by the client key even for - // requests that presented the server key + // requests that presented the server key, under an idempotency key assert!(flushed); let posts = requests_to(&mock_server, "/proxy/usage/").await; assert_eq!(posts.len(), 1); + assert!(!posts[0].headers["Idempotency-Key"].is_empty()); assert_eq!( usage_rows(&posts[0]), vec![ @@ -252,7 +253,7 @@ async fn test_failed_requests_are_not_counted() { } #[tokio::test] -async fn test_failed_flush_merges_counts_into_the_next() { +async fn test_failed_flush_retries_the_same_batch_before_new_counts() { // Given a served request and a usage endpoint that fails once let mock_server = MockServer::start().await; mount_config(&mock_server).await; @@ -263,19 +264,26 @@ async fn test_failed_flush_merges_counts_into_the_next() { service.refresh_environment_caches().await; service.track_usage(CLIENT_KEY, Resource::Flags); - // When the first flush fails and another request is served + // When the first flush fails, another request is served, and two + // more flushes run assert!(!service.flush_usage().await); service.track_usage(CLIENT_KEY, Resource::Flags); - - // Then the next flush carries both counts — nothing lost, nothing - // double-counted assert!(service.flush_usage().await); + assert!(service.flush_usage().await); + + // Then the failed batch is resent unchanged under its original key, + // and the count served meanwhile follows as its own batch let posts = requests_to(&mock_server, "/proxy/usage/").await; - assert_eq!(posts.len(), 2); - assert_eq!( - usage_rows(&posts[1]), - vec![json!({"client_side_key": CLIENT_KEY, "resource": "flags", "count": 2})] - ); + assert_eq!(posts.len(), 3); + let keys: Vec<&str> = posts + .iter() + .map(|post| post.headers["Idempotency-Key"].to_str().unwrap()) + .collect(); + assert_eq!(keys[0], keys[1]); + assert_ne!(keys[1], keys[2]); + let one_flag = vec![json!({"client_side_key": CLIENT_KEY, "resource": "flags", "count": 1})]; + assert_eq!(usage_rows(&posts[1]), one_flag); + assert_eq!(usage_rows(&posts[2]), one_flag); } #[tokio::test] @@ -354,10 +362,15 @@ async fn test_flush_chunks_batches_to_the_server_cap() { // When assert!(service.flush_usage().await); - // Then the rows arrive split across two accepted requests + // Then the rows arrive split across two accepted requests, each its + // own batch let posts = requests_to(&mock_server, "/proxy/usage/").await; let row_counts: Vec = posts.iter().map(|post| usage_rows(post).len()).collect(); assert_eq!(row_counts, vec![1000, 1]); + assert_ne!( + posts[0].headers["Idempotency-Key"], + posts[1].headers["Idempotency-Key"] + ); } #[tokio::test] From 78d2f69a50a488bd6d45180781b71e213a81fb87 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 13:30:39 +0530 Subject: [PATCH 06/22] fix: wait a full interval after a slow usage flush tokio intervals burst missed ticks by default, so a flush that overran the interval would be followed by back-to-back flushes. --- src/services/environment.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/services/environment.rs b/src/services/environment.rs index 53a5dd3..5d451d8 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -510,6 +510,9 @@ impl EnvironmentService { let mut interval = tokio::time::interval(Duration::from_secs( self.settings.usage_flush_interval_seconds, )); + // A flush that overruns the interval (core slow or down) must not be + // followed by a burst of catch-up flushes hammering it further. + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); // The first tick completes immediately, before anything is counted. interval.tick().await; From 0444a39f41817d09b031223f98b04e83e02efbb8 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 14:39:03 +0530 Subject: [PATCH 07/22] docs: drop the UsageCounts doc comment --- src/usage.rs | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/usage.rs b/src/usage.rs index cfca73c..a8f88c7 100644 --- a/src/usage.rs +++ b/src/usage.rs @@ -20,12 +20,6 @@ pub struct UsageRow { pub count: u64, } -/// Requests served since the last flush, aggregated per environment and -/// resource. Keyed by resolved client keys only, so the map is bounded by -/// the served environment set. -/// -/// Uses `parking_lot::Mutex`, not tokio's: every operation writes, and -/// guards are held only for a map operation, never across an await. #[derive(Default)] pub struct UsageCounts { count_by_environment_and_resource: Mutex>, From efa494ee2bfee3c871d5be474b3e5865500f4304 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 14:54:37 +0530 Subject: [PATCH 08/22] docs: drop the Resource doc comment --- src/usage.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/usage.rs b/src/usage.rs index a8f88c7..18c9cd3 100644 --- a/src/usage.rs +++ b/src/usage.rs @@ -3,7 +3,6 @@ use std::collections::HashMap; use parking_lot::Mutex; use serde::Serialize; -/// An SDK endpoint as the usage endpoint names it. #[derive(Serialize, Clone, Copy, PartialEq, Eq, Hash, Debug)] #[serde(rename_all = "kebab-case")] pub enum Resource { From c6bb20d317b35f999822c109e9605f5b86209868 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 14:58:01 +0530 Subject: [PATCH 09/22] test: aggregate across two client keys --- src/usage.rs | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/usage.rs b/src/usage.rs index 18c9cd3..df28da1 100644 --- a/src/usage.rs +++ b/src/usage.rs @@ -73,26 +73,32 @@ mod tests { let counts = UsageCounts::default(); // When - counts.increment("client", Resource::Flags); - counts.increment("client", Resource::Flags); - counts.increment("client", Resource::Identities); + counts.increment("client_a", Resource::Flags); + counts.increment("client_a", Resource::Flags); + counts.increment("client_a", Resource::Identities); + counts.increment("client_b", Resource::Flags); // Then let mut rows = counts.drain(); - rows.sort_by_key(|row| format!("{:?}", row.resource)); + rows.sort_by_key(|row| (row.client_side_key.clone(), format!("{:?}", row.resource))); assert_eq!( rows, vec![ UsageRow { - client_side_key: "client".to_string(), + client_side_key: "client_a".to_string(), resource: Resource::Flags, count: 2, }, UsageRow { - client_side_key: "client".to_string(), + client_side_key: "client_a".to_string(), resource: Resource::Identities, count: 1, }, + UsageRow { + client_side_key: "client_b".to_string(), + resource: Resource::Flags, + count: 1, + }, ] ); } From 85eba1b827c4f6ee8586421a11760403fc61bf29 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 14:58:14 +0530 Subject: [PATCH 10/22] docs: drop the UsageBatch doc comment --- src/usage.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/usage.rs b/src/usage.rs index df28da1..3ca7642 100644 --- a/src/usage.rs +++ b/src/usage.rs @@ -47,8 +47,6 @@ impl UsageCounts { } } -/// One `POST /proxy/usage/` body with the idempotency key it is sent -/// under, so a retry after a lost response is recognised, not recounted. pub struct UsageBatch { pub id: String, pub rows: Vec, From 583a101d60bd6b7bbb8ded88952c532f00c065a9 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 15:02:37 +0530 Subject: [PATCH 11/22] docs: drop the parse_next_link doc comment --- src/services/environment.rs | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/services/environment.rs b/src/services/environment.rs index 5d451d8..00c9cca 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -563,10 +563,6 @@ fn merge_paginated_overrides(base: &mut serde_json::Value, page: serde_json::Val /// Parse the `Link` response header for the next-page URL (RFC 5988). /// -/// Returns an absolute URL, resolving relative targets against `api_url`. -/// The `rel=next` pagination target, if any. A target off `api_url`'s -/// origin is an error: the next request carries the environment and proxy -/// keys, and a Link header must not be able to send them elsewhere. fn parse_next_link(headers: &HeaderMap, api_url: &str) -> Result> { let Ok(base) = Url::parse(api_url) else { return Ok(None); From 3507a91bfb6d89dc750fd2b3c2973ffacdfef629 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 15:12:10 +0530 Subject: [PATCH 12/22] refactor: move the layers into a middleware module --- src/lib.rs | 1 + src/{routes => middleware}/cors.rs | 0 src/middleware/mod.rs | 2 ++ src/{routes/usage.rs => middleware/usage_tracking.rs} | 2 +- src/routes/mod.rs | 11 +++++------ 5 files changed, 9 insertions(+), 7 deletions(-) rename src/{routes => middleware}/cors.rs (100%) create mode 100644 src/middleware/mod.rs rename src/{routes/usage.rs => middleware/usage_tracking.rs} (94%) diff --git a/src/lib.rs b/src/lib.rs index 14cb2b3..ce2f691 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ pub mod cache; pub mod config; pub mod environments; pub mod error; +pub mod middleware; pub mod models; pub mod routes; pub mod services; diff --git a/src/routes/cors.rs b/src/middleware/cors.rs similarity index 100% rename from src/routes/cors.rs rename to src/middleware/cors.rs diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs new file mode 100644 index 0000000..8c71dee --- /dev/null +++ b/src/middleware/mod.rs @@ -0,0 +1,2 @@ +pub mod cors; +pub mod usage_tracking; diff --git a/src/routes/usage.rs b/src/middleware/usage_tracking.rs similarity index 94% rename from src/routes/usage.rs rename to src/middleware/usage_tracking.rs index b5d4926..3134b71 100644 --- a/src/routes/usage.rs +++ b/src/middleware/usage_tracking.rs @@ -1,4 +1,4 @@ -use super::{ENVIRONMENT_DOCUMENT_PATH, FLAGS_PATH, IDENTITIES_PATH}; +use crate::routes::{ENVIRONMENT_DOCUMENT_PATH, FLAGS_PATH, IDENTITIES_PATH}; use crate::services::EnvironmentService; use crate::usage::Resource; use axum::extract::{MatchedPath, Request, State}; diff --git a/src/routes/mod.rs b/src/routes/mod.rs index cb5a29b..d7e9af5 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -1,12 +1,11 @@ -pub mod cors; pub mod environment_document; pub mod extractors; pub mod flags; pub mod health; pub mod identities; -pub mod usage; use crate::config::AppSettings; +use crate::middleware::{cors, usage_tracking}; use crate::services::EnvironmentService; use axum::{ Router, @@ -16,9 +15,9 @@ use axum::{ use std::sync::Arc; use tower_http::{compression::CompressionLayer, normalize_path::NormalizePath, trace::TraceLayer}; -const FLAGS_PATH: &str = "/api/v1/flags"; -const IDENTITIES_PATH: &str = "/api/v1/identities"; -const ENVIRONMENT_DOCUMENT_PATH: &str = "/api/v1/environment-document"; +pub(crate) const FLAGS_PATH: &str = "/api/v1/flags"; +pub(crate) const IDENTITIES_PATH: &str = "/api/v1/identities"; +pub(crate) const ENVIRONMENT_DOCUMENT_PATH: &str = "/api/v1/environment-document"; pub fn create_router(settings: AppSettings) -> (Router, Arc) { let cors = cors::layer(&settings.allow_origins); @@ -43,7 +42,7 @@ pub fn create_router(settings: AppSettings) -> (Router, Arc) // Middleware layers .layer(from_fn_with_state( environment_service.clone(), - usage::track_usage, + usage_tracking::track_usage, )) .layer(CompressionLayer::new()) .layer(cors) From 3595450745cb588c66d0a5ba84ed51e2d15036be Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 15:15:20 +0530 Subject: [PATCH 13/22] refactor: name the usage middleware module usage --- src/middleware/mod.rs | 2 +- src/middleware/{usage_tracking.rs => usage.rs} | 0 src/routes/mod.rs | 4 ++-- 3 files changed, 3 insertions(+), 3 deletions(-) rename src/middleware/{usage_tracking.rs => usage.rs} (100%) diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs index 8c71dee..d8d781e 100644 --- a/src/middleware/mod.rs +++ b/src/middleware/mod.rs @@ -1,2 +1,2 @@ pub mod cors; -pub mod usage_tracking; +pub mod usage; diff --git a/src/middleware/usage_tracking.rs b/src/middleware/usage.rs similarity index 100% rename from src/middleware/usage_tracking.rs rename to src/middleware/usage.rs diff --git a/src/routes/mod.rs b/src/routes/mod.rs index d7e9af5..5a7fe4a 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -5,7 +5,7 @@ pub mod health; pub mod identities; use crate::config::AppSettings; -use crate::middleware::{cors, usage_tracking}; +use crate::middleware::{cors, usage}; use crate::services::EnvironmentService; use axum::{ Router, @@ -42,7 +42,7 @@ pub fn create_router(settings: AppSettings) -> (Router, Arc) // Middleware layers .layer(from_fn_with_state( environment_service.clone(), - usage_tracking::track_usage, + usage::track_usage, )) .layer(CompressionLayer::new()) .layer(cors) From 3e1d7f2e665647cfeda3d21ce74378f462add509 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 15:28:18 +0530 Subject: [PATCH 14/22] docs: one-line MAX_ROWS_PER_FLUSH comment --- Cargo.lock | 233 ++++++++++++++++++++++++++++++++++++ Cargo.toml | 5 + src/services/environment.rs | 4 +- 3 files changed, 239 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 89aef42..82b9f44 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -26,6 +26,12 @@ dependencies = [ "libc", ] +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + [[package]] name = "anstyle" version = "1.0.13" @@ -260,6 +266,12 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6bd91ee7b2422bcb158d90ef4d14f75ef67f340943fc4149891dcce8f8b972a3" +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.2.51" @@ -304,6 +316,58 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c133bc6a41be0d194c306b5506d15e6feeea7b1d6604bd3f8310dfb2ca96486" + [[package]] name = "cmake" version = "0.1.57" @@ -394,6 +458,73 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "criterion" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f" +dependencies = [ + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "is-terminal", + "itertools", + "num-traits", + "once_cell", + "oorandom", + "plotters", + "rayon", + "regex", + "serde", + "serde_derive", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1" +dependencies = [ + "cast", + "itertools", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "622f3fc73690be383c7214310406f28a90e6edeadc3cea882f9d71e495b9711a" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc74980687109a3b14c72fd458107bf0baa1da1a1a805e178d15501ba9b86d9d" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.1.6" @@ -514,6 +645,7 @@ dependencies = [ "axum", "axum-test", "chrono", + "criterion", "flagsmith-flag-engine", "http", "mockall", @@ -534,6 +666,12 @@ dependencies = [ "wiremock", ] +[[package]] +name = "either" +version = "1.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" + [[package]] name = "email_address" version = "0.2.9" @@ -805,6 +943,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.16.0" @@ -1100,6 +1249,26 @@ dependencies = [ "serde", ] +[[package]] +name = "is-terminal" +version = "0.4.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" +dependencies = [ + "hermit-abi", + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.15" @@ -1391,6 +1560,12 @@ version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "openssl-probe" version = "0.2.0" @@ -1438,6 +1613,34 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "potential_utf" version = "0.1.3" @@ -1629,6 +1832,26 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2236,6 +2459,16 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tinyvec" version = "1.10.0" diff --git a/Cargo.toml b/Cargo.toml index 7cc2a07..a38574e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,7 @@ validator = { version = "0.20", features = ["derive"] } flagsmith-flag-engine = "0.6" [dev-dependencies] +criterion = "0.5" axum-test = "18" mockall = "0.14" tokio-test = "0.4" @@ -37,3 +38,7 @@ opt-level = 3 lto = true codegen-units = 1 strip = true + +[[bench]] +name = "usage" +harness = false diff --git a/src/services/environment.rs b/src/services/environment.rs index 00c9cca..b6db4f5 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -432,9 +432,7 @@ impl EnvironmentService { } } - /// The usage endpoint's batch cap — MAX_USAGE_ROWS in the - /// edge_control_plane app. Flushes are chunked to it so a large - /// environment set can never be rejected outright. + /// Mirrors MAX_USAGE_ROWS on the usage endpoint. const MAX_ROWS_PER_FLUSH: usize = 1000; /// Returns false when any batch was not accepted. From b62d9fd8a1584d39d0bfb48c480681155ba9fcab Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 15:30:49 +0530 Subject: [PATCH 15/22] docs: drop the pending-batch comment --- src/services/environment.rs | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/services/environment.rs b/src/services/environment.rs index b6db4f5..f1755fe 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -441,10 +441,6 @@ impl EnvironmentService { return true; }; - // A batch that failed may still have been processed, so it is - // resent unchanged under the same key and nothing new is drained - // until it is through. Meanwhile counts keep aggregating in the - // map, which is bounded by the served environment set. let mut batches = std::mem::take(&mut *self.pending_usage.lock()); if batches.is_empty() { let mut rows = self.usage.drain(); From a91c7ed3f022847e464d74f4e720367e5e6ed458 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 15:36:15 +0530 Subject: [PATCH 16/22] docs: drop the track_usage doc comment --- src/services/environment.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/services/environment.rs b/src/services/environment.rs index f1755fe..bf07290 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -171,8 +171,6 @@ impl EnvironmentService { .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string())) } - /// Count a served request against its environment. Static environments - /// stay on core's own billing, so they are not reported. pub fn track_usage(&self, environment_key: &str, resource: Resource) { if self.settings.proxy_key.is_none() { return; From 1168a6e0e3551260c1ee9c70457a5413813bb941 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 15:46:33 +0530 Subject: [PATCH 17/22] fix: bill every served environment the same once a proxy key is set Static and discovered environments were billed differently: static ones kept unmarked fetches and unreported requests. With a proxy key, all fetches are now marked and all served requests reported; the server's grant check decides what counts. --- src/environments.rs | 5 ----- src/services/environment.rs | 15 +++++---------- tests/test_usage_tracking.rs | 34 +++++++++++++++------------------- 3 files changed, 20 insertions(+), 34 deletions(-) diff --git a/src/environments.rs b/src/environments.rs index 9612d40..ac8a7fa 100644 --- a/src/environments.rs +++ b/src/environments.rs @@ -80,11 +80,6 @@ impl EnvironmentIndex { index } - /// Whether the key belongs to a statically configured environment. - pub fn is_static(&self, key: &str) -> bool { - self.protected.contains(key) - } - /// Resolve a presented key — client- or server-side — to its /// environment's keys. A server-side key resolves only while it is /// valid, so a deactivation delivered by the proxy config and an diff --git a/src/services/environment.rs b/src/services/environment.rs index bf07290..c587041 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -178,9 +178,7 @@ impl EnvironmentService { let Some(keys) = self.environments.resolve(environment_key) else { return; }; - if !self.environments.is_static(&keys.client_key) { - self.usage.increment(&keys.client_key, resource); - } + self.usage.increment(&keys.client_key, resource); } async fn fetch_environment(&self, keys: &EnvironmentKeys) -> Result { @@ -234,13 +232,10 @@ impl EnvironmentService { .client .get(&next_url) .header("X-Environment-Key", server_side_key); - // Core excludes marked fetches from API usage — the proxy - // reports served requests instead. Static environments stay - // unmarked and keep their old billing. - if !self.environments.is_static(server_side_key) { - if let Some(proxy_key) = &self.settings.proxy_key { - request = request.header("X-Proxy-Key", proxy_key); - } + // Core excludes marked fetches from API usage; the proxy reports + // served requests instead. + if let Some(proxy_key) = &self.settings.proxy_key { + request = request.header("X-Proxy-Key", proxy_key); } // If-Modified-Since is meaningful only on the first request; the // upstream pagination cursor (page_id) drives subsequent fetches. diff --git a/tests/test_usage_tracking.rs b/tests/test_usage_tracking.rs index 982e4f2..99a263a 100644 --- a/tests/test_usage_tracking.rs +++ b/tests/test_usage_tracking.rs @@ -374,8 +374,8 @@ async fn test_flush_chunks_batches_to_the_server_cap() { } #[tokio::test] -async fn test_static_environment_usage_is_neither_counted_nor_marked() { - // Given a proxy serving a static environment alongside a discovered one +async fn test_static_environment_with_proxy_key_is_billed_like_a_discovered_one() { + // Given a proxy key, and a static environment alongside a discovered one let mock_server = MockServer::start().await; mount_config(&mock_server).await; mount_document(&mock_server).await; @@ -395,28 +395,24 @@ async fn test_static_environment_usage_is_neither_counted_nor_marked() { service.track_usage(CLIENT_KEY, Resource::Flags); assert!(service.flush_usage().await); - // Then the static environment keeps its old billing: its document - // fetch is not marked as the proxy's own, and it is not reported + // Then every document fetch is marked and both environments are reported let fetches = requests_to(&mock_server, "/environment-document/").await; - let marked: Vec = fetches - .iter() - .map(|request| request.headers.contains_key("X-Proxy-Key")) - .collect(); - let static_fetches: Vec<&Request> = fetches - .iter() - .filter(|request| request.headers["X-Environment-Key"] == "ser.static_key") - .collect(); - assert!(!static_fetches.is_empty()); + assert!(!fetches.is_empty()); assert!( - static_fetches + fetches .iter() - .all(|request| !request.headers.contains_key("X-Proxy-Key")) + .all(|request| request.headers["X-Proxy-Key"] == PROXY_KEY) ); - assert!(marked.contains(&true)); let posts = requests_to(&mock_server, "/proxy/usage/").await; + assert_eq!(posts.len(), 1); + let mut rows = usage_rows(&posts[0]); + rows.sort_by_key(|row| row["client_side_key"].as_str().unwrap().to_string()); assert_eq!( - usage_rows(&posts[0]), - vec![json!({"client_side_key": CLIENT_KEY, "resource": "flags", "count": 1})] + rows, + vec![ + json!({"client_side_key": CLIENT_KEY, "resource": "flags", "count": 1}), + json!({"client_side_key": "static_client", "resource": "flags", "count": 1}), + ] ); } @@ -440,7 +436,7 @@ async fn test_document_fetch_carries_the_proxy_key() { } #[tokio::test] -async fn test_static_document_fetch_omits_the_proxy_key() { +async fn test_document_fetch_without_proxy_key_is_unmarked() { // Given a statically configured proxy with no proxy key let mock_server = MockServer::start().await; mount_document(&mock_server).await; From feb2b49c2364b57828fb6d8a2f404a5f46830e5e Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 15:58:37 +0530 Subject: [PATCH 18/22] chore: drop the bench dependency committed by mistake 3e1d7f2 swept criterion and a [[bench]] target into an unrelated commit; the bench file itself was never added. --- Cargo.lock | 233 ----------------------------------------------------- Cargo.toml | 5 -- 2 files changed, 238 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 82b9f44..89aef42 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -26,12 +26,6 @@ dependencies = [ "libc", ] -[[package]] -name = "anes" -version = "0.1.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" - [[package]] name = "anstyle" version = "1.0.13" @@ -266,12 +260,6 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6bd91ee7b2422bcb158d90ef4d14f75ef67f340943fc4149891dcce8f8b972a3" -[[package]] -name = "cast" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" - [[package]] name = "cc" version = "1.2.51" @@ -316,58 +304,6 @@ dependencies = [ "windows-link 0.2.1", ] -[[package]] -name = "ciborium" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" -dependencies = [ - "ciborium-io", - "ciborium-ll", - "serde", -] - -[[package]] -name = "ciborium-io" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" - -[[package]] -name = "ciborium-ll" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" -dependencies = [ - "ciborium-io", - "half", -] - -[[package]] -name = "clap" -version = "4.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" -dependencies = [ - "clap_builder", -] - -[[package]] -name = "clap_builder" -version = "4.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" -dependencies = [ - "anstyle", - "clap_lex", -] - -[[package]] -name = "clap_lex" -version = "1.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1c133bc6a41be0d194c306b5506d15e6feeea7b1d6604bd3f8310dfb2ca96486" - [[package]] name = "cmake" version = "0.1.57" @@ -458,73 +394,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "criterion" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f" -dependencies = [ - "anes", - "cast", - "ciborium", - "clap", - "criterion-plot", - "is-terminal", - "itertools", - "num-traits", - "once_cell", - "oorandom", - "plotters", - "rayon", - "regex", - "serde", - "serde_derive", - "serde_json", - "tinytemplate", - "walkdir", -] - -[[package]] -name = "criterion-plot" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1" -dependencies = [ - "cast", - "itertools", -] - -[[package]] -name = "crossbeam-deque" -version = "0.8.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "622f3fc73690be383c7214310406f28a90e6edeadc3cea882f9d71e495b9711a" -dependencies = [ - "crossbeam-epoch", - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-epoch" -version = "0.9.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc74980687109a3b14c72fd458107bf0baa1da1a1a805e178d15501ba9b86d9d" -dependencies = [ - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-utils" -version = "0.8.23" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" - -[[package]] -name = "crunchy" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" - [[package]] name = "crypto-common" version = "0.1.6" @@ -645,7 +514,6 @@ dependencies = [ "axum", "axum-test", "chrono", - "criterion", "flagsmith-flag-engine", "http", "mockall", @@ -666,12 +534,6 @@ dependencies = [ "wiremock", ] -[[package]] -name = "either" -version = "1.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" - [[package]] name = "email_address" version = "0.2.9" @@ -943,17 +805,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "half" -version = "2.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" -dependencies = [ - "cfg-if", - "crunchy", - "zerocopy", -] - [[package]] name = "hashbrown" version = "0.16.0" @@ -1249,26 +1100,6 @@ dependencies = [ "serde", ] -[[package]] -name = "is-terminal" -version = "0.4.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" -dependencies = [ - "hermit-abi", - "libc", - "windows-sys 0.61.2", -] - -[[package]] -name = "itertools" -version = "0.10.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" -dependencies = [ - "either", -] - [[package]] name = "itoa" version = "1.0.15" @@ -1560,12 +1391,6 @@ version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" -[[package]] -name = "oorandom" -version = "11.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" - [[package]] name = "openssl-probe" version = "0.2.0" @@ -1613,34 +1438,6 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" -[[package]] -name = "plotters" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" -dependencies = [ - "num-traits", - "plotters-backend", - "plotters-svg", - "wasm-bindgen", - "web-sys", -] - -[[package]] -name = "plotters-backend" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" - -[[package]] -name = "plotters-svg" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" -dependencies = [ - "plotters-backend", -] - [[package]] name = "potential_utf" version = "0.1.3" @@ -1832,26 +1629,6 @@ dependencies = [ "getrandom 0.3.4", ] -[[package]] -name = "rayon" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" -dependencies = [ - "either", - "rayon-core", -] - -[[package]] -name = "rayon-core" -version = "1.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" -dependencies = [ - "crossbeam-deque", - "crossbeam-utils", -] - [[package]] name = "redox_syscall" version = "0.5.18" @@ -2459,16 +2236,6 @@ dependencies = [ "zerovec", ] -[[package]] -name = "tinytemplate" -version = "1.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" -dependencies = [ - "serde", - "serde_json", -] - [[package]] name = "tinyvec" version = "1.10.0" diff --git a/Cargo.toml b/Cargo.toml index a38574e..7cc2a07 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,7 +26,6 @@ validator = { version = "0.20", features = ["derive"] } flagsmith-flag-engine = "0.6" [dev-dependencies] -criterion = "0.5" axum-test = "18" mockall = "0.14" tokio-test = "0.4" @@ -38,7 +37,3 @@ opt-level = 3 lto = true codegen-units = 1 strip = true - -[[bench]] -name = "usage" -harness = false From e2d65c5e8a318d26cf3cccf4834b4df49a432461 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 15:58:51 +0530 Subject: [PATCH 19/22] refactor: move usage counting and reporting into UsageProcessor EnvironmentService was also counting served requests and posting usage batches. That now lives in services/usage.rs; the router state is a struct holding both services, and the middleware resolves the key through one and records on the other. --- src/main.rs | 10 +-- src/middleware/usage.rs | 13 ++- src/routes/environment_document.rs | 5 +- src/routes/flags.rs | 5 +- src/routes/health.rs | 5 +- src/routes/identities.rs | 7 +- src/routes/mod.rs | 19 +++-- src/services/environment.rs | 110 +------------------------ src/services/mod.rs | 2 + src/services/usage.rs | 128 +++++++++++++++++++++++++++++ src/state.rs | 9 +- tests/test_cors.rs | 5 +- tests/test_health.rs | 12 +-- tests/test_server.rs | 10 ++- tests/test_usage_tracking.rs | 74 +++++++++-------- 15 files changed, 231 insertions(+), 183 deletions(-) create mode 100644 src/services/usage.rs diff --git a/src/main.rs b/src/main.rs index 38b1930..894d18c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -30,23 +30,23 @@ async fn main() -> anyhow::Result<()> { ); } - let (app, environment_service) = create_router(settings.clone()); + let (app, state) = create_router(settings.clone()); // Refreshes must never overlap: a delayed older poll finishing after a // newer one could restore removed environments or rotated keys. The // poll loop is serial, so it just has to start after the initial // refresh completes. info!("Loading initial environment data..."); - environment_service.refresh_environment_caches().await; + state.environments.refresh_environment_caches().await; - let polling_service = environment_service.clone(); + let polling_service = state.environments.clone(); tokio::spawn(async move { polling_service.poll_environments().await; }); - let usage_service = environment_service.clone(); + let usage = state.usage.clone(); tokio::spawn(async move { - usage_service.flush_usage_periodically().await; + usage.flush_periodically().await; }); let addr = SocketAddr::from(( diff --git a/src/middleware/usage.rs b/src/middleware/usage.rs index 3134b71..fe4e28c 100644 --- a/src/middleware/usage.rs +++ b/src/middleware/usage.rs @@ -1,17 +1,12 @@ use crate::routes::{ENVIRONMENT_DOCUMENT_PATH, FLAGS_PATH, IDENTITIES_PATH}; -use crate::services::EnvironmentService; +use crate::state::AppState; use crate::usage::Resource; use axum::extract::{MatchedPath, Request, State}; use axum::middleware::Next; use axum::response::Response; -use std::sync::Arc; /// Count a request once the handler has answered: only a 2xx is usage. -pub async fn track_usage( - State(service): State>, - request: Request, - next: Next, -) -> Response { +pub async fn track_usage(State(state): State, request: Request, next: Next) -> Response { let resource = request .extensions() .get::() @@ -28,7 +23,9 @@ pub async fn track_usage( return response; } if let (Some(resource), Some(environment_key)) = (resource, environment_key) { - service.track_usage(&environment_key, resource); + if let Some(keys) = state.environments.resolve(&environment_key) { + state.usage.record(&keys.client_key, resource); + } } response } diff --git a/src/routes/environment_document.rs b/src/routes/environment_document.rs index 7e68e3a..da361f0 100644 --- a/src/routes/environment_document.rs +++ b/src/routes/environment_document.rs @@ -1,14 +1,15 @@ use crate::error::{EdgeProxyError, Result}; use crate::routes::extractors::extract_environment_key; -use crate::state::AppState; +use crate::services::EnvironmentService; use axum::{ extract::State, http::{HeaderMap, header}, response::IntoResponse, }; +use std::sync::Arc; pub async fn get_environment_document( - State(service): State, + State(service): State>, headers: HeaderMap, ) -> Result { let environment_key = extract_environment_key(&headers)?; diff --git a/src/routes/flags.rs b/src/routes/flags.rs index 136f868..b51ed0c 100644 --- a/src/routes/flags.rs +++ b/src/routes/flags.rs @@ -1,12 +1,13 @@ use crate::error::Result; use crate::routes::extractors::extract_environment_key; -use crate::state::AppState; +use crate::services::EnvironmentService; use axum::{ Json, extract::{Query, State}, http::HeaderMap, }; use serde::Deserialize; +use std::sync::Arc; #[derive(Deserialize)] pub struct FlagsQuery { @@ -14,7 +15,7 @@ pub struct FlagsQuery { } pub async fn get_flags( - State(service): State, + State(service): State>, headers: HeaderMap, Query(query): Query, ) -> Result> { diff --git a/src/routes/health.rs b/src/routes/health.rs index 14738b8..2720652 100644 --- a/src/routes/health.rs +++ b/src/routes/health.rs @@ -1,7 +1,8 @@ -use crate::state::AppState; +use crate::services::EnvironmentService; use axum::{Json, extract::State, http::StatusCode, response::IntoResponse}; use chrono::Utc; use serde::{Deserialize, Serialize}; +use std::sync::Arc; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct HealthCheckResponse { @@ -28,7 +29,7 @@ impl HealthCheckResponse { } } -pub async fn health_check(State(service): State) -> impl IntoResponse { +pub async fn health_check(State(service): State>) -> impl IntoResponse { let last_updated = service.last_updated_at.read().await; match *last_updated { diff --git a/src/routes/identities.rs b/src/routes/identities.rs index 2fabaaf..83904f5 100644 --- a/src/routes/identities.rs +++ b/src/routes/identities.rs @@ -1,13 +1,14 @@ use crate::error::Result; use crate::models::{IdentityResponse, IdentityWithTraits}; use crate::routes::extractors::extract_environment_key; -use crate::state::AppState; +use crate::services::EnvironmentService; use axum::{ Json, extract::{Query, State}, http::HeaderMap, }; use serde::Deserialize; +use std::sync::Arc; #[derive(Deserialize)] pub struct IdentitiesQuery { @@ -15,7 +16,7 @@ pub struct IdentitiesQuery { } pub async fn get_identities( - State(service): State, + State(service): State>, headers: HeaderMap, Query(query): Query, ) -> Result> { @@ -30,7 +31,7 @@ pub async fn get_identities( } pub async fn post_identities( - State(service): State, + State(service): State>, headers: HeaderMap, Json(identity): Json, ) -> Result> { diff --git a/src/routes/mod.rs b/src/routes/mod.rs index 5a7fe4a..8bcda0d 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -6,7 +6,8 @@ pub mod identities; use crate::config::AppSettings; use crate::middleware::{cors, usage}; -use crate::services::EnvironmentService; +use crate::services::{EnvironmentService, UsageProcessor}; +use crate::state::AppState; use axum::{ Router, middleware::{from_fn_with_state, map_response}, @@ -19,9 +20,12 @@ pub(crate) const FLAGS_PATH: &str = "/api/v1/flags"; pub(crate) const IDENTITIES_PATH: &str = "/api/v1/identities"; pub(crate) const ENVIRONMENT_DOCUMENT_PATH: &str = "/api/v1/environment-document"; -pub fn create_router(settings: AppSettings) -> (Router, Arc) { +pub fn create_router(settings: AppSettings) -> (Router, AppState) { let cors = cors::layer(&settings.allow_origins); - let environment_service = Arc::new(EnvironmentService::new(settings)); + let state = AppState { + usage: Arc::new(UsageProcessor::new(&settings)), + environments: Arc::new(EnvironmentService::new(settings)), + }; let router = Router::new() // Health check routes @@ -40,19 +44,16 @@ pub fn create_router(settings: AppSettings) -> (Router, Arc) get(environment_document::get_environment_document), ) // Middleware layers - .layer(from_fn_with_state( - environment_service.clone(), - usage::track_usage, - )) + .layer(from_fn_with_state(state.clone(), usage::track_usage)) .layer(CompressionLayer::new()) .layer(cors) .layer(map_response(cors::merge_vary)) .layer(TraceLayer::new_for_http()) - .with_state(environment_service.clone()); + .with_state(state.clone()); // Trailing-slash normalization must wrap the router itself: axum matches // routes before `Router::layer` middleware runs let app = Router::new().fallback_service(NormalizePath::trim_trailing_slash(router)); - (app, environment_service) + (app, state) } diff --git a/src/services/environment.rs b/src/services/environment.rs index c587041..e97337d 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -6,12 +6,10 @@ use crate::models::{ APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, }; use crate::services::feature_utils::filter_out_server_key_only_flag_results; -use crate::usage::{Resource, UsageBatch, UsageCounts}; use chrono::{DateTime, Utc}; use flagsmith_flag_engine::engine::get_evaluation_result; use flagsmith_flag_engine::engine_eval::{FlagResult, add_identity_to_context}; use flagsmith_flag_engine::identities::Trait as FlagsmithTrait; -use parking_lot::Mutex; use reqwest::header::HeaderMap; use reqwest::{Client, Url}; use std::sync::Arc; @@ -25,14 +23,6 @@ pub struct EnvironmentService { pub settings: AppSettings, pub last_updated_at: Arc>>>, environments: EnvironmentIndex, - usage: UsageCounts, - pending_usage: Mutex>, -} - -enum UsageOutcome { - Accepted, - Rejected, - Failed, } impl EnvironmentService { @@ -51,8 +41,6 @@ impl EnvironmentService { settings, last_updated_at: Arc::new(RwLock::new(None)), environments, - usage: UsageCounts::default(), - pending_usage: Mutex::default(), } } @@ -165,22 +153,16 @@ impl EnvironmentService { Ok(response.json().await?) } + pub fn resolve(&self, environment_key: &str) -> Option> { + self.environments.resolve(environment_key) + } + fn resolve_key(&self, environment_key: &str) -> Result> { self.environments .resolve(environment_key) .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string())) } - pub fn track_usage(&self, environment_key: &str, resource: Resource) { - if self.settings.proxy_key.is_none() { - return; - } - let Some(keys) = self.environments.resolve(environment_key) else { - return; - }; - self.usage.increment(&keys.client_key, resource); - } - async fn fetch_environment(&self, keys: &EnvironmentKeys) -> Result { let server_key = keys.valid_server_key().ok_or_else(|| { EdgeProxyError::ServiceUnavailable(format!( @@ -424,90 +406,6 @@ impl EnvironmentService { self.refresh_environment_caches().await; } } - - /// Mirrors MAX_USAGE_ROWS on the usage endpoint. - const MAX_ROWS_PER_FLUSH: usize = 1000; - - /// Returns false when any batch was not accepted. - pub async fn flush_usage(&self) -> bool { - let Some(proxy_key) = &self.settings.proxy_key else { - return true; - }; - - let mut batches = std::mem::take(&mut *self.pending_usage.lock()); - if batches.is_empty() { - let mut rows = self.usage.drain(); - while !rows.is_empty() { - let chunk = rows.drain(..rows.len().min(Self::MAX_ROWS_PER_FLUSH)); - batches.push(UsageBatch::new(chunk.collect())); - } - } - - let mut all_success = true; - for batch in batches { - match self.post_usage(proxy_key, &batch).await { - UsageOutcome::Accepted => {} - UsageOutcome::Rejected => all_success = false, - UsageOutcome::Failed => { - self.pending_usage.lock().push(batch); - all_success = false; - } - } - } - all_success - } - - async fn post_usage(&self, proxy_key: &str, batch: &UsageBatch) -> UsageOutcome { - let url = format!("{}/proxy/usage/", self.settings.api_url); - let result = self - .client - .post(&url) - .header("X-Proxy-Key", proxy_key) - .header("Idempotency-Key", &batch.id) - .json(&batch.rows) - .send() - .await; - match result { - Ok(response) if response.status().is_success() => UsageOutcome::Accepted, - // Retrying cannot heal a rejection: drop the batch rather than - // resend it forever. - Ok(response) if response.status().is_client_error() => { - error!( - "Usage report rejected with {}: dropping {} rows", - response.status(), - batch.rows.len() - ); - UsageOutcome::Rejected - } - Ok(response) => { - error!("Failed to report usage: {}", response.status()); - UsageOutcome::Failed - } - Err(e) => { - error!("Failed to report usage: {}", e); - UsageOutcome::Failed - } - } - } - - pub async fn flush_usage_periodically(self: Arc) { - if self.settings.proxy_key.is_none() { - return; - } - let mut interval = tokio::time::interval(Duration::from_secs( - self.settings.usage_flush_interval_seconds, - )); - // A flush that overruns the interval (core slow or down) must not be - // followed by a burst of catch-up flushes hammering it further. - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); - // The first tick completes immediately, before anything is counted. - interval.tick().await; - - loop { - interval.tick().await; - self.flush_usage().await; - } - } } /// Format the cached document's `updated_at` as an RFC 2822 `If-Modified-Since` diff --git a/src/services/mod.rs b/src/services/mod.rs index efaf4e3..f9517e1 100644 --- a/src/services/mod.rs +++ b/src/services/mod.rs @@ -1,4 +1,6 @@ pub mod environment; pub mod feature_utils; +pub mod usage; pub use environment::EnvironmentService; +pub use usage::UsageProcessor; diff --git a/src/services/usage.rs b/src/services/usage.rs new file mode 100644 index 0000000..02a0b5d --- /dev/null +++ b/src/services/usage.rs @@ -0,0 +1,128 @@ +use crate::config::settings::AppSettings; +use crate::usage::{Resource, UsageBatch, UsageCounts}; +use parking_lot::Mutex; +use reqwest::Client; +use std::sync::Arc; +use std::time::Duration; +use tracing::error; + +enum Outcome { + Accepted, + Rejected, + Failed, +} + +pub struct UsageProcessor { + counts: UsageCounts, + pending: Mutex>, + client: Client, + api_url: String, + proxy_key: Option, + flush_interval: Duration, +} + +impl UsageProcessor { + /// Mirrors MAX_USAGE_ROWS on the usage endpoint. + const MAX_ROWS_PER_FLUSH: usize = 1000; + + pub fn new(settings: &AppSettings) -> Self { + let client = Client::builder() + .timeout(Duration::from_secs(settings.api_poll_timeout_seconds)) + .build() + .expect("Failed to create HTTP client"); + Self { + counts: UsageCounts::default(), + pending: Mutex::default(), + client, + api_url: settings.api_url.clone(), + proxy_key: settings.proxy_key.clone(), + flush_interval: Duration::from_secs(settings.usage_flush_interval_seconds), + } + } + + pub fn record(&self, client_key: &str, resource: Resource) { + if self.proxy_key.is_none() { + return; + } + self.counts.increment(client_key, resource); + } + + /// Returns false when any batch was not accepted. + pub async fn flush(&self) -> bool { + let Some(proxy_key) = &self.proxy_key else { + return true; + }; + + let mut batches = std::mem::take(&mut *self.pending.lock()); + if batches.is_empty() { + let mut rows = self.counts.drain(); + while !rows.is_empty() { + let chunk = rows.drain(..rows.len().min(Self::MAX_ROWS_PER_FLUSH)); + batches.push(UsageBatch::new(chunk.collect())); + } + } + + let mut all_success = true; + for batch in batches { + match self.post(proxy_key, &batch).await { + Outcome::Accepted => {} + Outcome::Rejected => all_success = false, + Outcome::Failed => { + self.pending.lock().push(batch); + all_success = false; + } + } + } + all_success + } + + async fn post(&self, proxy_key: &str, batch: &UsageBatch) -> Outcome { + let url = format!("{}/proxy/usage/", self.api_url); + let result = self + .client + .post(&url) + .header("X-Proxy-Key", proxy_key) + .header("Idempotency-Key", &batch.id) + .json(&batch.rows) + .send() + .await; + match result { + Ok(response) if response.status().is_success() => Outcome::Accepted, + // Retrying cannot heal a rejection: drop the batch rather than + // resend it forever. + Ok(response) if response.status().is_client_error() => { + error!( + "Usage report rejected with {}: dropping {} rows", + response.status(), + batch.rows.len() + ); + Outcome::Rejected + } + Ok(response) => { + error!("Failed to report usage: {}", response.status()); + Outcome::Failed + } + Err(e) => { + error!("Failed to report usage: {}", e); + Outcome::Failed + } + } + } + + pub async fn flush_periodically(self: Arc) { + if self.proxy_key.is_none() { + return; + } + let mut interval = tokio::time::interval(self.flush_interval); + // A flush that overruns the interval (core slow or down) must not be + // followed by a burst of catch-up flushes hammering it further. + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + // The first tick completes immediately, before anything is counted. + interval.tick().await; + + loop { + interval.tick().await; + self.flush().await; + } + } +} diff --git a/src/state.rs b/src/state.rs index 43ae606..c946f9c 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,4 +1,9 @@ -use crate::services::EnvironmentService; +use crate::services::{EnvironmentService, UsageProcessor}; +use axum::extract::FromRef; use std::sync::Arc; -pub type AppState = Arc; +#[derive(Clone, FromRef)] +pub struct AppState { + pub environments: Arc, + pub usage: Arc, +} diff --git a/tests/test_cors.rs b/tests/test_cors.rs index c34414c..0474c4f 100644 --- a/tests/test_cors.rs +++ b/tests/test_cors.rs @@ -17,8 +17,9 @@ async fn server_allowing(allow_origins: &[&str]) -> TestServer { allow_origins: allow_origins.iter().map(|o| o.to_string()).collect(), ..AppSettings::new() }; - let (app, service) = create_router(settings); - service + let (app, state) = create_router(settings); + state + .environments .cache .put_environment(&environment_1_api_key(), environment_1()) .await; diff --git a/tests/test_health.rs b/tests/test_health.rs index 8016c49..ef99212 100644 --- a/tests/test_health.rs +++ b/tests/test_health.rs @@ -21,9 +21,9 @@ async fn test_liveness_check() { async fn test_health_check_returns_200_if_cache_was_updated_recently() { // Given let settings = AppSettings::new(); - let (app, service) = create_router(settings); + let (app, state) = create_router(settings); { - let mut last_updated = service.last_updated_at.write().await; + let mut last_updated = state.environments.last_updated_at.write().await; *last_updated = Some(Utc::now()); } let server = TestServer::new(app).unwrap(); @@ -62,9 +62,9 @@ async fn test_health_check_returns_503_if_cache_was_not_updated() { async fn test_health_check_returns_503_if_cache_is_stale() { // Given let settings = AppSettings::new(); - let (app, service) = create_router(settings); + let (app, state) = create_router(settings); { - let mut last_updated = service.last_updated_at.write().await; + let mut last_updated = state.environments.last_updated_at.write().await; *last_updated = Some(Utc::now() - chrono::Duration::days(10)); } let server = TestServer::new(app).unwrap(); @@ -88,10 +88,10 @@ async fn test_health_check_returns_200_if_cache_is_never_stale() { settings.health_check = HealthCheckSettings { environment_update_grace_period_seconds: None, }; - let (app, service) = create_router(settings); + let (app, state) = create_router(settings); let last_update_time = Utc::now() - chrono::Duration::days(10); { - let mut last_updated = service.last_updated_at.write().await; + let mut last_updated = state.environments.last_updated_at.write().await; *last_updated = Some(last_update_time); } let server = TestServer::new(app).unwrap(); diff --git a/tests/test_server.rs b/tests/test_server.rs index 2a599af..2ddf01a 100644 --- a/tests/test_server.rs +++ b/tests/test_server.rs @@ -16,9 +16,10 @@ async fn setup_test_server() -> TestServer { ..AppSettings::new() }; - let (app, service) = create_router(settings); + let (app, state) = create_router(settings); - service + state + .environments .cache .put_environment(&environment_1_api_key(), environment_1()) .await; @@ -208,8 +209,9 @@ async fn test_get_environment_document() { }], ..AppSettings::new() }; - let (app, service) = create_router(settings); - service + let (app, state) = create_router(settings); + state + .environments .cache .put_environment("client_key", environment_1()) .await; diff --git a/tests/test_usage_tracking.rs b/tests/test_usage_tracking.rs index 99a263a..f17ba70 100644 --- a/tests/test_usage_tracking.rs +++ b/tests/test_usage_tracking.rs @@ -2,7 +2,7 @@ use axum::http::StatusCode; use axum_test::TestServer; use edge_proxy::config::settings::{AppSettings, EnvironmentKeyPair}; use edge_proxy::routes::create_router; -use edge_proxy::services::EnvironmentService; +use edge_proxy::services::{EnvironmentService, UsageProcessor}; use edge_proxy::usage::Resource; use serde_json::{Value, json}; use wiremock::matchers::{header, method, path}; @@ -126,8 +126,8 @@ async fn test_served_requests_flush_aggregated_usage() { mount_config(&mock_server).await; mount_document(&mock_server).await; mount_usage(&mock_server, 204, None).await; - let (app, service) = create_router(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); - service.refresh_environment_caches().await; + let (app, state) = create_router(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + state.environments.refresh_environment_caches().await; let server = TestServer::new(app).unwrap(); // When SDK traffic arrives under both keys, then usage is flushed @@ -158,7 +158,7 @@ async fn test_served_requests_flush_aggregated_usage() { .add_header("X-Environment-Key", SERVER_KEY) .await .assert_status_ok(); - let flushed = service.flush_usage().await; + let flushed = state.usage.flush().await; // Then one POST reports everything, keyed by the client key even for // requests that presented the server key, under an idempotency key @@ -182,8 +182,8 @@ async fn test_unresolved_keys_are_never_counted() { let mock_server = MockServer::start().await; mount_config(&mock_server).await; mount_document(&mock_server).await; - let (app, service) = create_router(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); - service.refresh_environment_caches().await; + let (app, state) = create_router(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + state.environments.refresh_environment_caches().await; let server = TestServer::new(app).unwrap(); // When requests present an unknown key or none at all @@ -198,7 +198,7 @@ async fn test_unresolved_keys_are_never_counted() { .assert_status_unauthorized(); // Then there is nothing to flush and no request is made - assert!(service.flush_usage().await); + assert!(state.usage.flush().await); assert!(requests_to(&mock_server, "/proxy/usage/").await.is_empty()); } @@ -230,8 +230,8 @@ async fn test_failed_requests_are_not_counted() { .mount(&mock_server) .await; mount_document(&mock_server).await; - let (app, service) = create_router(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); - service.refresh_environment_caches().await; + let (app, state) = create_router(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + state.environments.refresh_environment_caches().await; let server = TestServer::new(app).unwrap(); // When requests fail after their key resolved @@ -248,7 +248,7 @@ async fn test_failed_requests_are_not_counted() { .assert_status(StatusCode::SERVICE_UNAVAILABLE); // Then nothing is reported - assert!(service.flush_usage().await); + assert!(state.usage.flush().await); assert!(requests_to(&mock_server, "/proxy/usage/").await.is_empty()); } @@ -260,16 +260,18 @@ async fn test_failed_flush_retries_the_same_batch_before_new_counts() { mount_document(&mock_server).await; mount_usage(&mock_server, 500, Some(1)).await; mount_usage(&mock_server, 204, None).await; - let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + let config = settings(&mock_server.uri(), Some(PROXY_KEY), vec![]); + let service = EnvironmentService::new(config.clone()); + let usage = UsageProcessor::new(&config); service.refresh_environment_caches().await; - service.track_usage(CLIENT_KEY, Resource::Flags); + usage.record(CLIENT_KEY, Resource::Flags); // When the first flush fails, another request is served, and two // more flushes run - assert!(!service.flush_usage().await); - service.track_usage(CLIENT_KEY, Resource::Flags); - assert!(service.flush_usage().await); - assert!(service.flush_usage().await); + assert!(!usage.flush().await); + usage.record(CLIENT_KEY, Resource::Flags); + assert!(usage.flush().await); + assert!(usage.flush().await); // Then the failed batch is resent unchanged under its original key, // and the count served meanwhile follows as its own batch @@ -291,19 +293,21 @@ async fn test_flush_without_proxy_key_is_inert() { // Given a statically configured proxy with no proxy key let mock_server = MockServer::start().await; mount_document(&mock_server).await; - let service = EnvironmentService::new(settings( + let config = settings( &mock_server.uri(), None, vec![EnvironmentKeyPair { client_side_key: CLIENT_KEY.to_string(), server_side_key: SERVER_KEY.to_string(), }], - )); + ); + let service = EnvironmentService::new(config.clone()); + let usage = UsageProcessor::new(&config); service.refresh_environment_caches().await; - service.track_usage(CLIENT_KEY, Resource::Flags); + usage.record(CLIENT_KEY, Resource::Flags); // When / Then: flushing succeeds without reporting anything - assert!(service.flush_usage().await); + assert!(usage.flush().await); assert!(requests_to(&mock_server, "/proxy/usage/").await.is_empty()); } @@ -314,16 +318,18 @@ async fn test_rejected_flush_drops_the_batch_instead_of_retrying_it() { mount_config(&mock_server).await; mount_document(&mock_server).await; mount_usage(&mock_server, 400, None).await; - let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + let config = settings(&mock_server.uri(), Some(PROXY_KEY), vec![]); + let service = EnvironmentService::new(config.clone()); + let usage = UsageProcessor::new(&config); service.refresh_environment_caches().await; - service.track_usage(CLIENT_KEY, Resource::Flags); + usage.record(CLIENT_KEY, Resource::Flags); // When the flush is rejected - assert!(!service.flush_usage().await); + assert!(!usage.flush().await); // Then the rows are dropped, not resent forever: the next flush has // nothing to send - assert!(service.flush_usage().await); + assert!(usage.flush().await); assert_eq!(requests_to(&mock_server, "/proxy/usage/").await.len(), 1); } @@ -353,14 +359,16 @@ async fn test_flush_chunks_batches_to_the_server_cap() { .await; mount_document(&mock_server).await; mount_usage(&mock_server, 204, None).await; - let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + let config = settings(&mock_server.uri(), Some(PROXY_KEY), vec![]); + let service = EnvironmentService::new(config.clone()); + let usage = UsageProcessor::new(&config); service.refresh_environment_caches().await; for n in 0..1001 { - service.track_usage(&format!("client_{n}"), Resource::Flags); + usage.record(&format!("client_{n}"), Resource::Flags); } // When - assert!(service.flush_usage().await); + assert!(usage.flush().await); // Then the rows arrive split across two accepted requests, each its // own batch @@ -380,20 +388,22 @@ async fn test_static_environment_with_proxy_key_is_billed_like_a_discovered_one( mount_config(&mock_server).await; mount_document(&mock_server).await; mount_usage(&mock_server, 204, None).await; - let service = EnvironmentService::new(settings( + let config = settings( &mock_server.uri(), Some(PROXY_KEY), vec![EnvironmentKeyPair { client_side_key: "static_client".to_string(), server_side_key: "ser.static_key".to_string(), }], - )); + ); + let service = EnvironmentService::new(config.clone()); + let usage = UsageProcessor::new(&config); service.refresh_environment_caches().await; // When both environments serve a request, and usage is flushed - service.track_usage("static_client", Resource::Flags); - service.track_usage(CLIENT_KEY, Resource::Flags); - assert!(service.flush_usage().await); + usage.record("static_client", Resource::Flags); + usage.record(CLIENT_KEY, Resource::Flags); + assert!(usage.flush().await); // Then every document fetch is marked and both environments are reported let fetches = requests_to(&mock_server, "/environment-document/").await; From 1bfece516b1e1d03b56b48f235d8294247d3e7b1 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 16:05:28 +0530 Subject: [PATCH 20/22] refactor: UsageProcessor::record is track --- src/middleware/usage.rs | 2 +- src/services/usage.rs | 2 +- tests/test_usage_tracking.rs | 14 +++++++------- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/middleware/usage.rs b/src/middleware/usage.rs index fe4e28c..21417f5 100644 --- a/src/middleware/usage.rs +++ b/src/middleware/usage.rs @@ -24,7 +24,7 @@ pub async fn track_usage(State(state): State, request: Request, next: } if let (Some(resource), Some(environment_key)) = (resource, environment_key) { if let Some(keys) = state.environments.resolve(&environment_key) { - state.usage.record(&keys.client_key, resource); + state.usage.track(&keys.client_key, resource); } } response diff --git a/src/services/usage.rs b/src/services/usage.rs index 02a0b5d..081cad8 100644 --- a/src/services/usage.rs +++ b/src/services/usage.rs @@ -40,7 +40,7 @@ impl UsageProcessor { } } - pub fn record(&self, client_key: &str, resource: Resource) { + pub fn track(&self, client_key: &str, resource: Resource) { if self.proxy_key.is_none() { return; } diff --git a/tests/test_usage_tracking.rs b/tests/test_usage_tracking.rs index f17ba70..d6a971b 100644 --- a/tests/test_usage_tracking.rs +++ b/tests/test_usage_tracking.rs @@ -264,12 +264,12 @@ async fn test_failed_flush_retries_the_same_batch_before_new_counts() { let service = EnvironmentService::new(config.clone()); let usage = UsageProcessor::new(&config); service.refresh_environment_caches().await; - usage.record(CLIENT_KEY, Resource::Flags); + usage.track(CLIENT_KEY, Resource::Flags); // When the first flush fails, another request is served, and two // more flushes run assert!(!usage.flush().await); - usage.record(CLIENT_KEY, Resource::Flags); + usage.track(CLIENT_KEY, Resource::Flags); assert!(usage.flush().await); assert!(usage.flush().await); @@ -304,7 +304,7 @@ async fn test_flush_without_proxy_key_is_inert() { let service = EnvironmentService::new(config.clone()); let usage = UsageProcessor::new(&config); service.refresh_environment_caches().await; - usage.record(CLIENT_KEY, Resource::Flags); + usage.track(CLIENT_KEY, Resource::Flags); // When / Then: flushing succeeds without reporting anything assert!(usage.flush().await); @@ -322,7 +322,7 @@ async fn test_rejected_flush_drops_the_batch_instead_of_retrying_it() { let service = EnvironmentService::new(config.clone()); let usage = UsageProcessor::new(&config); service.refresh_environment_caches().await; - usage.record(CLIENT_KEY, Resource::Flags); + usage.track(CLIENT_KEY, Resource::Flags); // When the flush is rejected assert!(!usage.flush().await); @@ -364,7 +364,7 @@ async fn test_flush_chunks_batches_to_the_server_cap() { let usage = UsageProcessor::new(&config); service.refresh_environment_caches().await; for n in 0..1001 { - usage.record(&format!("client_{n}"), Resource::Flags); + usage.track(&format!("client_{n}"), Resource::Flags); } // When @@ -401,8 +401,8 @@ async fn test_static_environment_with_proxy_key_is_billed_like_a_discovered_one( service.refresh_environment_caches().await; // When both environments serve a request, and usage is flushed - usage.record("static_client", Resource::Flags); - usage.record(CLIENT_KEY, Resource::Flags); + usage.track("static_client", Resource::Flags); + usage.track(CLIENT_KEY, Resource::Flags); assert!(usage.flush().await); // Then every document fetch is marked and both environments are reported From 7b97c2144c8e89a860304eeaba40544a251dbfd1 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 16:19:35 +0530 Subject: [PATCH 21/22] refactor: rename pending to pending_batches --- src/services/usage.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/services/usage.rs b/src/services/usage.rs index 081cad8..b4d4e04 100644 --- a/src/services/usage.rs +++ b/src/services/usage.rs @@ -14,7 +14,7 @@ enum Outcome { pub struct UsageProcessor { counts: UsageCounts, - pending: Mutex>, + pending_batches: Mutex>, client: Client, api_url: String, proxy_key: Option, @@ -32,7 +32,7 @@ impl UsageProcessor { .expect("Failed to create HTTP client"); Self { counts: UsageCounts::default(), - pending: Mutex::default(), + pending_batches: Mutex::default(), client, api_url: settings.api_url.clone(), proxy_key: settings.proxy_key.clone(), @@ -53,7 +53,7 @@ impl UsageProcessor { return true; }; - let mut batches = std::mem::take(&mut *self.pending.lock()); + let mut batches = std::mem::take(&mut *self.pending_batches.lock()); if batches.is_empty() { let mut rows = self.counts.drain(); while !rows.is_empty() { @@ -68,7 +68,7 @@ impl UsageProcessor { Outcome::Accepted => {} Outcome::Rejected => all_success = false, Outcome::Failed => { - self.pending.lock().push(batch); + self.pending_batches.lock().push(batch); all_success = false; } } From 27d28dba829cd8658897927a895a97e4a4758f60 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 19 Sep 2026 16:57:49 +0530 Subject: [PATCH 22/22] fix: never follow redirects on requests carrying keys reqwest strips only its own list of sensitive headers on a cross-host redirect; X-Proxy-Key and X-Environment-Key would be forwarded. --- src/services/environment.rs | 3 ++ src/services/usage.rs | 3 ++ tests/test_usage_tracking.rs | 56 ++++++++++++++++++++++++++++++++++++ 3 files changed, 62 insertions(+) diff --git a/src/services/environment.rs b/src/services/environment.rs index e97337d..a742fb4 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -30,6 +30,9 @@ impl EnvironmentService { let client = Client::builder() .timeout(Duration::from_secs(settings.api_poll_timeout_seconds)) .gzip(true) + // reqwest keeps custom headers across redirects; never let one + // carry the environment or proxy key to another host. + .redirect(reqwest::redirect::Policy::none()) .build() .expect("Failed to create HTTP client"); diff --git a/src/services/usage.rs b/src/services/usage.rs index b4d4e04..cab67e5 100644 --- a/src/services/usage.rs +++ b/src/services/usage.rs @@ -28,6 +28,9 @@ impl UsageProcessor { pub fn new(settings: &AppSettings) -> Self { let client = Client::builder() .timeout(Duration::from_secs(settings.api_poll_timeout_seconds)) + // reqwest keeps custom headers across redirects; never let one + // carry the proxy key to another host. + .redirect(reqwest::redirect::Policy::none()) .build() .expect("Failed to create HTTP client"); Self { diff --git a/tests/test_usage_tracking.rs b/tests/test_usage_tracking.rs index d6a971b..c43ec85 100644 --- a/tests/test_usage_tracking.rs +++ b/tests/test_usage_tracking.rs @@ -471,3 +471,59 @@ async fn test_document_fetch_without_proxy_key_is_unmarked() { .all(|request| !request.headers.contains_key("X-Proxy-Key")) ); } + +#[tokio::test] +async fn test_document_fetch_does_not_follow_redirects() { + // Given a document endpoint that redirects to another host + let mock_server = MockServer::start().await; + let other_host = MockServer::start().await; + mount_config(&mock_server).await; + Mock::given(method("GET")) + .and(path("/environment-document/")) + .respond_with(ResponseTemplate::new(302).insert_header( + "Location", + format!("{}/environment-document/", other_host.uri()), + )) + .mount(&mock_server) + .await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(200).set_body_json(document_body())) + .mount(&other_host) + .await; + let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + + // When + let refreshed = service.refresh_environment_caches().await; + + // Then the fetch fails and the keys never reach the other host + assert!(!refreshed); + assert!(other_host.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn test_usage_flush_does_not_follow_redirects() { + // Given a usage endpoint that redirects to another host + let mock_server = MockServer::start().await; + let other_host = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/proxy/usage/")) + .respond_with( + ResponseTemplate::new(307) + .insert_header("Location", format!("{}/proxy/usage/", other_host.uri())), + ) + .mount(&mock_server) + .await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(204)) + .mount(&other_host) + .await; + let usage = UsageProcessor::new(&settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + usage.track(CLIENT_KEY, Resource::Flags); + + // When + let flushed = usage.flush().await; + + // Then the batch is not accepted and the key never reaches the other host + assert!(!flushed); + assert!(other_host.received_requests().await.unwrap().is_empty()); +}