feat: Report served-request usage

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

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<String>,
+    // 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<String> {
     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

src/lib.rs

@@ -2,7 +2,9 @@ pub mod cache;
 pub mod config;
 pub mod environments;
 pub mod error;
+pub mod middleware;
 pub mod models;
 pub mod routes;
 pub mod services;
 pub mod state;
+pub mod usage;

src/main.rs

@@ -30,20 +30,25 @@ 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 = state.usage.clone();
+    tokio::spawn(async move {
+        usage.flush_periodically().await;
+    });
+
     let addr = SocketAddr::from((
         settings
             .server

src/middleware/mod.rs

@@ -0,0 +1,2 @@
+pub mod cors;
+pub mod usage;

src/middleware/usage.rs

@@ -0,0 +1,40 @@
+use crate::routes::{ENVIRONMENT_DOCUMENT_PATH, FLAGS_PATH, IDENTITIES_PATH};
+use crate::state::AppState;
+use crate::usage::Resource;
+use axum::extract::{MatchedPath, Request, State};
+use axum::middleware::Next;
+use axum::response::Response;
+
+/// Count a request once the handler has answered: only a 2xx is usage.
+pub async fn track_usage(State(state): State<AppState>, request: Request, next: Next) -> Response {
+    let resource = request
+        .extensions()
+        .get::<MatchedPath>()
+        .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) {
+        if let Some(keys) = state.environments.resolve(&environment_key) {
+            state.usage.track(&keys.client_key, resource);
+        }
+    }
+    response
+}
+
+fn resource_for(route: &str) -> Option<Resource> {
+    match route {
+        FLAGS_PATH => Some(Resource::Flags),
+        IDENTITIES_PATH => Some(Resource::Identities),
+        ENVIRONMENT_DOCUMENT_PATH => Some(Resource::EnvironmentDocument),
+        _ => None,
+    }
+}

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<AppState>,
+    State(service): State<Arc<EnvironmentService>>,
     headers: HeaderMap,
 ) -> Result<impl IntoResponse> {
     let environment_key = extract_environment_key(&headers)?;

src/routes/flags.rs

@@ -1,20 +1,21 @@
 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 {
     pub feature: Option<String>,
 }
 
 pub async fn get_flags(
-    State(service): State<AppState>,
+    State(service): State<Arc<EnvironmentService>>,
     headers: HeaderMap,
     Query(query): Query<FlagsQuery>,
 ) -> Result<Json<serde_json::Value>> {

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<AppState>) -> impl IntoResponse {
+pub async fn health_check(State(service): State<Arc<EnvironmentService>>) -> impl IntoResponse {
     let last_updated = service.last_updated_at.read().await;
 
     match *last_updated {

src/routes/identities.rs

@@ -1,21 +1,22 @@
 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 {
     pub identifier: String,
 }
 
 pub async fn get_identities(
-    State(service): State<AppState>,
+    State(service): State<Arc<EnvironmentService>>,
     headers: HeaderMap,
     Query(query): Query<IdentitiesQuery>,
 ) -> Result<Json<IdentityResponse>> {
@@ -30,7 +31,7 @@ pub async fn get_identities(
 }
 
 pub async fn post_identities(
-    State(service): State<AppState>,
+    State(service): State<Arc<EnvironmentService>>,
     headers: HeaderMap,
     Json(identity): Json<IdentityWithTraits>,
 ) -> Result<Json<IdentityResponse>> {

src/routes/mod.rs

@@ -1,23 +1,31 @@
-pub mod cors;
 pub mod environment_document;
 pub mod extractors;
 pub mod flags;
 pub mod health;
 pub mod identities;
 
 use crate::config::AppSettings;
-use crate::services::EnvironmentService;
+use crate::middleware::{cors, usage};
+use crate::services::{EnvironmentService, UsageProcessor};
+use crate::state::AppState;
 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};
 
-pub fn create_router(settings: AppSettings) -> (Router, Arc<EnvironmentService>) {
+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, 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
@@ -26,25 +34,26 @@ pub fn create_router(settings: AppSettings) -> (Router, Arc<EnvironmentService>)
         .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(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)
 }