use std::sync::Arc; use std::time::{Instant, SystemTime, UNIX_EPOCH}; use axiom_rs::Client; use axum::body::Body; use axum::extract::State; use axum::http::{HeaderMap, Request, header}; use axum::middleware::Next; use axum::response::Response; use serde::Serialize; use serde_json::{Map, Value}; use tokio::sync::Mutex; use uuid::Uuid; use crate::config::Config; use crate::state::AppState; /// Shared, request-scoped event handle stored in request extensions. /// /// Usage: /// ```rust /// # use slipnote_backend::logging::SharedEvent; /// # use axum::Extension; /// async fn handler(Extension(event): Extension) { /// let mut event = event.lock().await; /// event.set_nested_field("request", "checked", true); /// } /// ``` pub type SharedEvent = Arc>; /// A wide event that accumulates request-scoped context for logging. /// /// Usage: /// ```rust /// # use slipnote_backend::logging::WideEvent; /// let mut event = WideEvent::default(); /// event.set_nested_field("request", "request_id", "req-123"); /// event.set_nested_field("transcription", "language", "en"); /// ``` #[derive(Debug, Clone, Default, Serialize)] pub struct WideEvent { #[serde(flatten)] fields: Map, } impl WideEvent { /// Builds a new event from request metadata. pub fn from_request(request: &Request) -> Self { let mut event = Self::default(); let request_id = header_value(request.headers(), "x-request-id") .unwrap_or_else(|| Uuid::new_v4().to_string()); let timestamp_ms = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)) .unwrap_or_default(); event.set_nested_field("request", "request_id", request_id); event.set_nested_field("request", "timestamp", timestamp_ms); event.set_nested_field("request", "path", request.uri().path()); let user_agent = request .headers() .get(header::USER_AGENT) .and_then(|value| value.to_str().ok()) .map(str::to_string); event.set_optional_nested_field("request", "user_agent", user_agent); event.set_optional_nested_field("request", "ip", extract_ip(request.headers())); event } /// Sets a nested field if the value is present. /// /// Usage: /// ```rust /// # use slipnote_backend::logging::WideEvent; /// let mut event = WideEvent::default(); /// event.set_optional_nested_field("request", "user_agent", Some("ua")); /// event.set_optional_nested_field::<&str>("request", "ip", None); /// ``` pub fn set_optional_nested_field( &mut self, object_key: impl Into, field_key: impl Into, value: Option, ) { let Some(value) = value else { return; }; self.set_nested_field(object_key, field_key, value); } /// Sets a nested field inside an object, creating it if needed. /// /// Usage: /// ```rust /// # use slipnote_backend::logging::WideEvent; /// let mut event = WideEvent::default(); /// event.set_nested_field("transcription", "language", "en"); /// event.set_nested_field("transcription", "duration", 1.2); /// ``` pub fn set_nested_field( &mut self, object_key: impl Into, field_key: impl Into, value: T, ) { let object_key = object_key.into(); let field_key = field_key.into(); let Ok(value) = serde_json::to_value(value) else { return; }; if let Some(Value::Object(existing)) = self.fields.get_mut(&object_key) { existing.insert(field_key, value); } else { let mut map = Map::new(); map.insert(field_key, value); self.fields.insert(object_key, Value::Object(map)); } } /// Records an error message on the event. /// /// Usage: /// ```rust /// # use slipnote_backend::logging::WideEvent; /// let mut event = WideEvent::default(); /// event.set_error("upstream failed"); /// ``` pub fn set_error(&mut self, err: impl std::fmt::Display) { self.set_nested_field("error", "message", err.to_string()); } } /// A lightweight wrapper around the Axiom client. /// /// Usage: /// ```rust,no_run /// # use slipnote_backend::config::Config; /// # use slipnote_backend::logging::Logger; /// let config = Config::from_env().expect("config"); /// let logger = Logger::from_config(&config).expect("logger"); /// ``` #[derive(Clone, Default)] pub struct Logger { inner: Option>, sample_rate: f64, #[cfg(test)] test_sender: Option>, } #[derive(Debug)] struct LoggerInner { client: Client, dataset: String, } impl Logger { /// Builds a logger from config, enabling ingestion when Axiom is configured. pub fn from_config(config: &Config) -> Result { let sample_rate = clamp_sample_rate(config.log_sample_rate); let Some(token) = config.axiom_token.clone() else { return Ok(Self { inner: None, sample_rate, #[cfg(test)] test_sender: None, }); }; let Some(dataset) = config.axiom_dataset.clone() else { return Ok(Self { inner: None, sample_rate, #[cfg(test)] test_sender: None, }); }; let mut builder = Client::builder().with_token(token); if let Some(url) = config.axiom_url.clone() { builder = builder.with_url(url); } let client = builder.build()?; Ok(Self { inner: Some(Arc::new(LoggerInner { client, dataset })), sample_rate, #[cfg(test)] test_sender: None, }) } /// Returns a disabled logger that drops events. pub fn disabled() -> Self { Self { inner: None, sample_rate: 1.0, #[cfg(test)] test_sender: None, } } #[cfg(test)] /// Builds a logger that captures events for tests. pub fn with_test_sender(sender: tokio::sync::mpsc::Sender) -> Self { Self { inner: None, sample_rate: 1.0, test_sender: Some(sender), } } /// Emits a single event in the background when enabled. pub fn emit(&self, event: WideEvent) { if !should_sample(&event, self.sample_rate) { return; } if cfg!(debug_assertions) && let Ok(line) = serde_json::to_string(&event) { println!("{line}"); } #[cfg(test)] if let Some(sender) = &self.test_sender { let _ = sender.try_send(event.clone()); } let Some(inner) = &self.inner else { return; }; let client = inner.client.clone(); let dataset = inner.dataset.clone(); tokio::spawn(async move { let _ = client.ingest(dataset, vec![event]).await; }); } } fn clamp_sample_rate(rate: f64) -> f64 { if !rate.is_finite() { return 1.0; } rate.clamp(0.0, 1.0) } fn should_sample(event: &WideEvent, rate: f64) -> bool { let request = (event .fields .get("request")) .and_then(|value| value.as_object()); let status_code = request .and_then(|map| map.get("status_code")) .and_then(|value| value.as_u64()); if status_code.unwrap_or(0) >= 500 { return true; } if event.fields.get("error").is_some() { return true; } let duration_ms = request .and_then(|map| map.get("duration_ms")) .and_then(|value| value.as_u64()); if duration_ms.unwrap_or(0) > 2000 { return true; } let rate = clamp_sample_rate(rate); if rate >= 1.0 { return true; } if rate <= 0.0 { return false; } rand::random::() < rate } /// Attaches a request-scoped event, then emits it after the handler completes. pub async fn logging_middleware( State(state): State, mut request: Request, next: Next, ) -> Response { let event = Arc::new(Mutex::new(WideEvent::from_request(&request))); { let mut record = event.lock().await; record.set_nested_field("request", "environment", state.config.slipnote_env.clone()); } request.extensions_mut().insert(event.clone()); let start = Instant::now(); let response = next.run(request).await; let duration_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX); let record = { let mut record = event.lock().await; let status = response.status(); record.set_nested_field("request", "status_code", status.as_u16()); record.set_nested_field("request", "duration_ms", duration_ms); record.set_nested_field( "request", "outcome", if status.is_success() { "success" } else { "error" }, ); record.clone() }; state.logger.emit(record); response } fn extract_ip(headers: &HeaderMap) -> Option { if let Some(value) = header_value(headers, "x-forwarded-for") { return value.split(',').next().map(|value| value.trim().to_string()); } header_value(headers, "x-real-ip") } fn header_value(headers: &HeaderMap, name: &str) -> Option { headers .get(name) .and_then(|value| value.to_str().ok()) .map(str::to_string) } #[cfg(test)] mod tests { use super::*; use axum::routing::get; use axum::{Extension, Router}; use proptest::prelude::*; use tokio::sync::mpsc; use tower::ServiceExt; use crate::test_support::TestStateBuilder; #[tokio::test] async fn middleware_injects_and_emits_event() { let (sender, mut receiver) = mpsc::channel(1); let logger = Logger::with_test_sender(sender); let state = TestStateBuilder::new() .with_logger(logger.clone()) .build(); let app = Router::new() .route( "/health", get(|Extension(event): Extension| async move { let mut event = event.lock().await; event.set_nested_field("request", "checked", true); "ok" }), ) .layer(axum::middleware::from_fn_with_state( state.clone(), logging_middleware, )) .with_state(state); let response = app .oneshot(Request::builder().uri("/health").body(Body::empty()).unwrap()) .await .expect("health response"); assert_eq!(response.status(), axum::http::StatusCode::OK); let emitted = receiver.recv().await.expect("event"); let request = emitted .fields .get("request") .and_then(|value| value.as_object()) .expect("request object"); assert!(request.contains_key("request_id")); } #[test] fn nested_fields_merge_into_object() { let mut event = WideEvent::default(); event.set_nested_field("transcription", "language", "en"); event.set_nested_field("transcription", "duration", 1.2); let transcription = event .fields .get("transcription") .and_then(|value| value.as_object()) .expect("transcription object"); assert_eq!(transcription.get("language"), Some(&Value::String("en".to_string()))); assert_eq!(transcription.get("duration"), Some(&Value::from(1.2))); } #[test] fn optional_nested_fields_skip_none() { let mut event = WideEvent::default(); event.set_optional_nested_field("request", "user_agent", None::); event.set_optional_nested_field("request", "user_agent", Some("ua".to_string())); let request = event .fields .get("request") .and_then(|value| value.as_object()) .expect("request object"); assert_eq!(request.get("user_agent"), Some(&Value::String("ua".to_string()))); } proptest! { #[test] fn nested_field_overwrite_is_consistent( object_key in "obj_[a-z]{1,8}", field_key in "field_[a-z]{1,8}", first in ".*", second in ".*", ) { let mut event = WideEvent::default(); event.set_nested_field(&object_key, &field_key, &first); event.set_nested_field(&object_key, &field_key, &second); let stored = event .fields .get(&object_key) .and_then(|value| value.as_object()) .and_then(|map| map.get(&field_key)) .and_then(|value| value.as_str()) .unwrap_or(""); prop_assert_eq!(stored, second); } } }