Compare commits

...
Author SHA1 Message Date
Brandon Thomas a19e989962 updated graceful behavior 2026-05-23 10:57:57 -04:00
Brandon Thomas 463197edea datadog client, metrics, worker 2026-05-23 09:41:59 -04:00
25 changed files with 1598 additions and 2 deletions
Generated
+28
View File
@@ -3038,6 +3038,22 @@ version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c297a1c74b71ae29df00c3e22dd9534821d60eb9af5a0192823fa2acea70c2a"
[[package]]
name = "datadog_client"
version = "0.0.1"
dependencies = [
"env",
"env_logger 0.11.9",
"log",
"reqwest 0.12.23",
"serde",
"serde_derive",
"serde_json",
"thiserror 2.0.18",
"tokio",
"workspace-build-acceleration",
]
[[package]]
name = "datetimes"
version = "0.0.1"
@@ -6950,6 +6966,16 @@ dependencies = [
"workspace-build-acceleration",
]
[[package]]
name = "metrics"
version = "0.0.1"
dependencies = [
"datadog_client",
"log",
"tokio",
"workspace-build-acceleration",
]
[[package]]
name = "midjourney_client"
version = "0.0.1"
@@ -11455,6 +11481,7 @@ dependencies = [
"composite_identifiers",
"config",
"crockford",
"datadog_client",
"datetimes",
"derive_more 0.99.17",
"diesel",
@@ -11487,6 +11514,7 @@ dependencies = [
"md-5 0.10.6",
"media",
"memory_caching",
"metrics",
"migration",
"mimetypes",
"mysql_queries",
+4
View File
@@ -10,6 +10,7 @@ members = [
"crates/api_clients/artcraft/artcraft_client",
"crates/api_clients/artcraft/artcraft_router",
"crates/api_clients/cloud_storage",
"crates/api_clients/datadog_client",
"crates/api_clients/fal_client",
"crates/api_clients/google_drive_downloader",
"crates/api_clients/google_sign_in",
@@ -79,6 +80,7 @@ members = [
"crates/lib/jwt_signer",
"crates/lib/logging",
"crates/lib/markdown",
"crates/lib/metrics",
"crates/lib/pager",
"crates/lib/password",
"crates/lib/server_environment",
@@ -240,6 +242,7 @@ component_traits = { path = "crates/service/plugins/component_traits" }
concurrency = { path = "crates/lib/concurrency" }
cookie_store = { path = "crates/lib/cookie_store" }
crockford = { path = "crates/lib/data/crockford" }
datadog_client = { path = "crates/api_clients/datadog_client" }
easyenv = { path = "crates/lib/easyenv" }
enums = { path = "crates/schema/public/enums" }
error_utils = { path = "crates/lib/error_utils" }
@@ -265,6 +268,7 @@ memory_store = { path = "crates/lib/caching/memory_store" }
midjourney_client = { path = "crates/api_clients/midjourney_client" }
mimetypes = { path = "crates/lib/files/mimetypes" }
muapi_client = { path = "crates/api_clients/muapi_client" }
metrics = { path = "crates/lib/metrics" }
mysql_queries = { path = "crates/schema/database/mysql_queries" }
pager = { path = "crates/lib/pager" }
openai_sora_client = { path = "crates/api_clients/openai_sora_client" }
@@ -0,0 +1,29 @@
[package]
name = "datadog_client"
edition = "2024"
version = "0.0.1"
publish = false
[package.metadata.cargo-udeps.ignore]
normal = ["workspace-build-acceleration"]
[lib]
name = "datadog_client"
path = "src/lib.rs"
[dependencies]
# Workspace hack (faster compile times)
workspace-build-acceleration.workspace = true
# External
log.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_derive.workspace = true
serde_json.workspace = true
thiserror.workspace = true
[dev-dependencies]
env.workspace = true
env_logger.workspace = true
tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread", "time"] }
@@ -0,0 +1,251 @@
use log::debug;
use serde_derive::Serialize;
use crate::creds::DatadogApiKey;
use crate::error::DatadogError;
pub const DEFAULT_BASE_URL: &str = "https://api.datadoghq.com";
/// Cheap-to-clone HTTP client for the Datadog metrics API.
#[derive(Clone)]
pub struct DatadogClient {
base_url: String,
api_key: DatadogApiKey,
http: reqwest::Client,
}
impl DatadogClient {
pub fn new(api_key: DatadogApiKey) -> Self {
Self::with_base_url(api_key, DEFAULT_BASE_URL.to_string())
}
pub fn with_base_url(api_key: DatadogApiKey, base_url: String) -> Self {
Self {
base_url,
api_key,
http: reqwest::Client::new(),
}
}
/// POST `/api/v1/distribution_points`. Each `DistributionSeries.points`
/// entry contains a list of *raw* observation values; Datadog computes
/// percentiles server-side, so no t-digest plumbing here.
pub async fn submit_distribution_points(
&self,
series: Vec<DistributionSeries>,
) -> Result<(), DatadogError> {
if series.is_empty() {
return Ok(());
}
let body = DistributionsBody { series };
let url = format!("{}/api/v1/distribution_points", self.base_url);
debug!("submitting {} distribution series to {}", body.series.len(), url);
let response = self.http
.post(&url)
.header("DD-API-KEY", self.api_key.as_str())
.json(&body)
.send()
.await?;
expect_2xx(response).await
}
/// POST `/api/v1/series`. For pre-aggregated counts/gauges/rates.
pub async fn submit_series(
&self,
series: Vec<MetricSeries>,
) -> Result<(), DatadogError> {
if series.is_empty() {
return Ok(());
}
let body = SeriesBody { series };
let url = format!("{}/api/v1/series", self.base_url);
debug!("submitting {} metric series to {}", body.series.len(), url);
let response = self.http
.post(&url)
.header("DD-API-KEY", self.api_key.as_str())
.json(&body)
.send()
.await?;
expect_2xx(response).await
}
}
async fn expect_2xx(response: reqwest::Response) -> Result<(), DatadogError> {
let status = response.status();
if status.is_success() {
return Ok(());
}
let body = response.text().await.unwrap_or_default();
Err(DatadogError::Api { status: status.as_u16(), body })
}
// ==================== Wire-format types ====================
#[derive(Serialize)]
struct DistributionsBody {
series: Vec<DistributionSeries>,
}
#[derive(Clone, Debug, Serialize)]
pub struct DistributionSeries {
pub metric: String,
/// Each point is `[unix_timestamp_secs, [v1, v2, ...]]`.
pub points: Vec<DistributionPoint>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub tags: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub host: Option<String>,
}
/// Tuple-style serialization: serializes as `[timestamp, [values]]`.
#[derive(Clone, Debug, Serialize)]
pub struct DistributionPoint(pub i64, pub Vec<f64>);
#[derive(Serialize)]
struct SeriesBody {
series: Vec<MetricSeries>,
}
#[derive(Clone, Debug, Serialize)]
pub struct MetricSeries {
pub metric: String,
/// Each point is `[unix_timestamp_secs, value]`.
pub points: Vec<MetricPoint>,
/// `"count"`, `"gauge"`, or `"rate"`.
#[serde(rename = "type")]
pub metric_type: String,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub tags: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub host: Option<String>,
/// Seconds. Required by Datadog for `count`/`rate`.
#[serde(skip_serializing_if = "Option::is_none")]
pub interval: Option<i64>,
}
/// Tuple-style serialization: serializes as `[timestamp, value]`.
#[derive(Clone, Debug, Serialize)]
pub struct MetricPoint(pub i64, pub f64);
// ==================== Tests ====================
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn distribution_point_serializes_as_tuple() {
let point = DistributionPoint(1_700_000_000, vec![1.0, 2.5, 3.0]);
let json = serde_json::to_string(&point).unwrap();
assert_eq!(json, "[1700000000,[1.0,2.5,3.0]]");
}
#[test]
fn metric_point_serializes_as_tuple() {
let point = MetricPoint(1_700_000_000, 42.0);
let json = serde_json::to_string(&point).unwrap();
assert_eq!(json, "[1700000000,42.0]");
}
#[test]
fn metric_series_uses_type_key() {
let s = MetricSeries {
metric: "x".to_string(),
points: vec![MetricPoint(1, 1.0)],
metric_type: "count".to_string(),
tags: vec!["a:b".to_string()],
host: Some("h".to_string()),
interval: Some(10),
};
let json = serde_json::to_string(&s).unwrap();
assert!(json.contains(r#""type":"count""#), "got: {}", json);
assert!(json.contains(r#""interval":10"#), "got: {}", json);
}
#[test]
fn empty_tags_omitted() {
let s = MetricSeries {
metric: "x".to_string(),
points: vec![MetricPoint(1, 1.0)],
metric_type: "gauge".to_string(),
tags: vec![],
host: None,
interval: None,
};
let json = serde_json::to_string(&s).unwrap();
assert!(!json.contains("tags"));
assert!(!json.contains("host"));
assert!(!json.contains("interval"));
}
#[test]
fn api_key_debug_is_redacted() {
let key = crate::creds::DatadogApiKey::new("supersecret123");
let debug = format!("{:?}", key);
assert!(!debug.contains("supersecret123"));
assert!(debug.contains("***"));
}
// ==================== Real-API smoke tests ====================
//
// These are marked `#[ignore]` because they hit api.datadoghq.com with a
// real API key from `datadog_key.txt` at the workspace root and submit a
// tiny named metric (`datadog_client.smoke_test`).
//
// Run explicitly with:
// cargo test -p datadog_client -- --ignored
//
// After running, look for the metric in Datadog within ~1 minute.
mod real_api_smoke {
use super::*;
use crate::creds::DatadogApiKey;
use crate::test_utils::get_test_api_key::get_test_api_key;
use std::time::{SystemTime, UNIX_EPOCH};
fn now_secs() -> i64 {
SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs() as i64).unwrap_or(0)
}
fn smoke_tags() -> Vec<String> {
vec![
"source:datadog_client_smoke_test".to_string(),
"env:test".to_string(),
]
}
#[ignore] // hits the real Datadog API
#[tokio::test]
async fn submit_distribution_points_real_api() {
let _ = env_logger::builder().is_test(true).try_init();
let key = get_test_api_key().expect("could not load datadog_key.txt");
let client = DatadogClient::new(DatadogApiKey::new(key));
let series = vec![DistributionSeries {
metric: "datadog_client.smoke_test.duration".to_string(),
points: vec![DistributionPoint(now_secs(), vec![1.0, 2.5, 7.0])],
tags: smoke_tags(),
host: Some("smoke-test-host".to_string()),
}];
client.submit_distribution_points(series).await
.expect("distribution submission should succeed");
}
#[ignore] // hits the real Datadog API
#[tokio::test]
async fn submit_series_real_api() {
let _ = env_logger::builder().is_test(true).try_init();
let key = get_test_api_key().expect("could not load datadog_key.txt");
let client = DatadogClient::new(DatadogApiKey::new(key));
let series = vec![MetricSeries {
metric: "datadog_client.smoke_test.count".to_string(),
points: vec![MetricPoint(now_secs(), 1.0)],
metric_type: "count".to_string(),
tags: smoke_tags(),
host: Some("smoke-test-host".to_string()),
interval: Some(10),
}];
client.submit_series(series).await
.expect("series submission should succeed");
}
}
}
@@ -0,0 +1,20 @@
/// Opaque wrapper around a Datadog API key. `Debug` is implemented manually
/// so the secret never leaks to logs.
#[derive(Clone)]
pub struct DatadogApiKey(String);
impl DatadogApiKey {
pub fn new(key: impl Into<String>) -> Self {
Self(key.into())
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Debug for DatadogApiKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("DatadogApiKey(***)")
}
}
@@ -0,0 +1,13 @@
use thiserror::Error;
#[derive(Debug, Error)]
pub enum DatadogError {
#[error("HTTP transport error: {0}")]
Transport(#[from] reqwest::Error),
#[error("Datadog API returned non-2xx: status={status} body={body}")]
Api { status: u16, body: String },
#[error("Serialization error: {0}")]
Serialization(#[from] serde_json::Error),
}
@@ -0,0 +1,19 @@
//! Minimal Datadog HTTP API client.
//!
//! This is a *transport* crate — it knows how to talk to the Datadog metrics
//! API and nothing else. Buffering, aggregation, and the actix middleware
//! that feeds it live in the `metrics` library. Mirrors the
//! `rootly_client` + `pager` split.
//!
//! Supports two endpoints:
//!
//! - `POST /api/v1/distribution_points` — submit raw observation lists.
//! Datadog computes p50/p75/p90/p95/p99/avg/min/max/count server-side.
//! - `POST /api/v1/series` — submit pre-aggregated counts/gauges/rates.
pub mod client;
pub mod creds;
pub mod error;
#[cfg(test)]
mod test_utils;
@@ -0,0 +1,40 @@
#![cfg(test)]
use std::fs::read_to_string;
use std::path::{Path, PathBuf};
/// Reads the Datadog API key from `datadog_key.txt` at the workspace root.
/// Used only by `#[ignore]` smoke tests that hit the real API.
pub(crate) fn get_test_api_key() -> Result<String, String> {
let candidates = candidate_paths();
for path in &candidates {
if let Ok(contents) = read_to_string(path) {
let trimmed = contents.trim().to_string();
if trimmed.is_empty() {
return Err(format!("datadog_key.txt at {} is empty", path.display()));
}
return Ok(trimmed);
}
}
Err(format!(
"Could not locate datadog_key.txt. Tried: {}",
candidates.iter().map(|p| p.display().to_string()).collect::<Vec<_>>().join(", "),
))
}
/// Walk upward from the test's CWD looking for `datadog_key.txt`. Tests run
/// from the crate dir (`crates/api_clients/datadog_client`), so we need to
/// climb a few levels to reach the workspace root.
fn candidate_paths() -> Vec<PathBuf> {
let mut out = Vec::new();
let mut current: PathBuf = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
for _ in 0..6 {
out.push(current.join("datadog_key.txt"));
if !current.pop() {
break;
}
}
// Last-resort hardcoded path matching this repo's known layout.
out.push(Path::new("/Users/bt/dev/storyteller/artcraft/datadog_key.txt").to_path_buf());
out
}
@@ -0,0 +1,3 @@
#![cfg(test)]
pub(crate) mod get_test_api_key;
+23
View File
@@ -0,0 +1,23 @@
[package]
name = "metrics"
edition = "2024"
version = "0.0.1"
publish = false
[package.metadata.cargo-udeps.ignore]
normal = ["workspace-build-acceleration"]
[lib]
name = "metrics"
path = "src/lib.rs"
[dependencies]
# Workspace hack (faster compile times)
workspace-build-acceleration.workspace = true
# Internal
datadog_client.workspace = true
# External
log.workspace = true
tokio = { workspace = true, features = ["macros", "rt", "time", "sync"] }
+113
View File
@@ -0,0 +1,113 @@
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::time::Duration;
use datadog_client::client::DatadogClient;
use crate::collector::MetricsCollector;
use crate::queue::SampleQueue;
use crate::worker::{MetricsWorker, MetricsWorkerConfig};
pub const DEFAULT_FLUSH_INTERVAL_SECS: u64 = 10;
pub const DEFAULT_DURATION_METRIC: &str = "http.server.request.duration_ms";
pub const DEFAULT_COUNT_METRIC: &str = "http.server.request.count";
pub struct MetricsBuilder {
client: Option<DatadogClient>,
flush_interval: Duration,
duration_metric_name: String,
count_metric_name: String,
host: Option<String>,
env: Option<String>,
service_name: Option<String>,
max_buffered_samples: Option<usize>,
}
impl MetricsBuilder {
pub fn new() -> Self {
Self {
client: None,
flush_interval: Duration::from_secs(DEFAULT_FLUSH_INTERVAL_SECS),
duration_metric_name: DEFAULT_DURATION_METRIC.to_string(),
count_metric_name: DEFAULT_COUNT_METRIC.to_string(),
host: None,
env: None,
service_name: None,
max_buffered_samples: None,
}
}
pub fn datadog_client(mut self, client: DatadogClient) -> Self {
self.client = Some(client); self
}
pub fn flush_interval(mut self, interval: Duration) -> Self {
self.flush_interval = interval; self
}
pub fn duration_metric_name(mut self, name: impl Into<String>) -> Self {
self.duration_metric_name = name.into(); self
}
pub fn count_metric_name(mut self, name: impl Into<String>) -> Self {
self.count_metric_name = name.into(); self
}
pub fn host(mut self, host: impl Into<String>) -> Self {
self.host = Some(host.into()); self
}
pub fn env(mut self, env: impl Into<String>) -> Self {
self.env = Some(env.into()); self
}
pub fn service_name(mut self, name: impl Into<String>) -> Self {
self.service_name = Some(name.into()); self
}
pub fn max_buffered_samples(mut self, n: usize) -> Self {
self.max_buffered_samples = Some(n); self
}
/// Build the (collector, worker) pair plus a shutdown flag. Spawn
/// `worker.run()` on a long-lived tokio task; clone `collector` into
/// middleware/handlers; call `shutdown.store(true, Ordering::Relaxed)`
/// to drain the worker on graceful shutdown.
///
/// Returns `Err` only if no Datadog client was provided. Callers that
/// want to disable metrics should use [`MetricsCollector::noop`] instead.
pub fn build(self) -> Result<BuildOutput, String> {
let client = self.client
.ok_or_else(|| "MetricsBuilder: datadog_client is required".to_string())?;
let queue = match self.max_buffered_samples {
Some(n) => SampleQueue::with_capacity(n),
None => SampleQueue::new(),
};
let shutdown = Arc::new(AtomicBool::new(false));
let config = MetricsWorkerConfig {
flush_interval: self.flush_interval,
duration_metric_name: self.duration_metric_name,
count_metric_name: self.count_metric_name,
host: self.host,
env: self.env.unwrap_or_else(|| "unknown".to_string()),
service_name: self.service_name.unwrap_or_else(|| "unknown".to_string()),
};
let collector = MetricsCollector::new(queue.clone());
let worker = MetricsWorker::new(queue, client, config, shutdown.clone());
Ok(BuildOutput { collector, worker, shutdown })
}
}
impl Default for MetricsBuilder {
fn default() -> Self { Self::new() }
}
pub struct BuildOutput {
pub collector: MetricsCollector,
pub worker: MetricsWorker,
pub shutdown: Arc<AtomicBool>,
}
+196
View File
@@ -0,0 +1,196 @@
use std::time::SystemTime;
use crate::queue::SampleQueue;
use crate::sample::{CounterSample, ObservationSample, RequestSample, Sample};
/// Cheap-to-clone handle that records metric observations. All `record_*`
/// methods are non-blocking: they grab a short-lived mutex on the in-memory
/// queue and never do I/O.
///
/// Built via [`crate::builder::MetricsBuilder`]. Use [`Self::noop()`] when
/// metrics should be disabled (no Datadog API key configured, dev env, …).
#[derive(Clone)]
pub struct MetricsCollector {
queue: SampleQueue,
enabled: bool,
}
impl MetricsCollector {
pub(crate) fn new(queue: SampleQueue) -> Self {
Self { queue, enabled: true }
}
/// Collector that swallows every recording call. Use when no API key is
/// configured, so callers can wire the collector unconditionally.
pub fn noop() -> Self {
Self { queue: SampleQueue::new(), enabled: false }
}
pub fn is_enabled(&self) -> bool { self.enabled }
// ==================== Request observations ====================
/// Record one HTTP request observation. Called by the actix middleware.
pub fn record_request(
&self,
route: impl Into<String>,
method: impl Into<String>,
status_code: u16,
duration_ms: f64,
) {
if !self.enabled {
return;
}
self.queue.push(Sample::Request(RequestSample {
route: route.into(),
method: method.into(),
status_code,
duration_ms,
timestamp_secs: now_secs(),
}));
}
// ==================== Custom counters ====================
/// Increment a custom counter metric by 1. Emitted to Datadog as `count`
/// (with the worker's `flush_interval` as the rollup interval). Tags are
/// shipped verbatim — keep their cardinality bounded.
pub fn increment(&self, metric: impl Into<String>, tags: Vec<String>) {
self.add(metric, 1, tags);
}
/// Add `value` to a custom counter metric. Use [`Self::increment`] for the
/// common +1 case.
pub fn add(&self, metric: impl Into<String>, value: u64, tags: Vec<String>) {
if !self.enabled || value == 0 {
return;
}
let metric = metric.into();
let timestamp_secs = now_secs();
// We model `add(n)` as n increments of the same `(metric, tags)`
// bucket. The worker counts entries per group; emitting n separate
// samples keeps the aggregation logic uniform.
for _ in 0..value {
self.queue.push(Sample::Counter(CounterSample {
metric: metric.clone(),
tags: tags.clone(),
timestamp_secs,
}));
}
}
// ==================== Custom observations ====================
/// Record a single observation (`f64`) for a distribution metric. Datadog
/// will compute p50/p75/p90/p95/p99/avg/min/max/count server-side.
pub fn observe(&self, metric: impl Into<String>, value: f64, tags: Vec<String>) {
if !self.enabled {
return;
}
self.queue.push(Sample::Observation(ObservationSample {
metric: metric.into(),
value,
tags,
timestamp_secs: now_secs(),
}));
}
}
fn now_secs() -> i64 {
SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0)
}
#[cfg(test)]
mod tests {
use super::*;
mod request_path {
use super::*;
#[test]
fn noop_records_nothing() {
let c = MetricsCollector::noop();
c.record_request("/x", "GET", 200, 1.0);
assert!(!c.is_enabled());
}
#[test]
fn enabled_pushes_to_shared_queue() {
let q = SampleQueue::new();
let c = MetricsCollector::new(q.clone());
c.record_request("/x", "POST", 201, 2.5);
assert_eq!(q.len(), 1);
}
}
mod counter_path {
use super::*;
#[test]
fn increment_pushes_one_counter_sample() {
let q = SampleQueue::new();
let c = MetricsCollector::new(q.clone());
c.increment("my.metric", vec!["a:b".to_string()]);
assert_eq!(q.len(), 1);
match &q.drain()[0] {
Sample::Counter(c) => {
assert_eq!(c.metric, "my.metric");
assert_eq!(c.tags, vec!["a:b".to_string()]);
}
other => panic!("expected Counter, got {:?}", other),
}
}
#[test]
fn add_emits_n_samples() {
let q = SampleQueue::new();
let c = MetricsCollector::new(q.clone());
c.add("my.metric", 5, vec![]);
assert_eq!(q.len(), 5);
}
#[test]
fn add_zero_is_noop() {
let q = SampleQueue::new();
let c = MetricsCollector::new(q.clone());
c.add("my.metric", 0, vec![]);
assert_eq!(q.len(), 0);
}
#[test]
fn noop_collector_ignores_counters() {
let c = MetricsCollector::noop();
c.increment("x", vec![]);
c.add("x", 5, vec![]);
}
}
mod observation_path {
use super::*;
#[test]
fn observe_pushes_observation_sample() {
let q = SampleQueue::new();
let c = MetricsCollector::new(q.clone());
c.observe("my.dist", 3.14, vec!["k:v".to_string()]);
assert_eq!(q.len(), 1);
match &q.drain()[0] {
Sample::Observation(o) => {
assert_eq!(o.metric, "my.dist");
assert_eq!(o.value, 3.14);
assert_eq!(o.tags, vec!["k:v".to_string()]);
}
other => panic!("expected Observation, got {:?}", other),
}
}
#[test]
fn noop_collector_ignores_observations() {
let c = MetricsCollector::noop();
c.observe("x", 1.0, vec![]);
}
}
}
+22
View File
@@ -0,0 +1,22 @@
//! Per-request metrics for HTTP services, backed by Datadog.
//!
//! Three pieces:
//!
//! - [`collector::MetricsCollector`] — cheap-to-clone handle that
//! middleware/handlers call to record a request observation. Pushes the
//! sample onto a shared in-memory buffer; never blocks on I/O.
//! - [`queue::SampleQueue`] — the buffer itself (bounded, oldest dropped).
//! - [`worker::MetricsWorker`] — a long-lived tokio task that drains the
//! queue every `flush_interval`, groups samples by (route, method,
//! status_code), and submits one distribution-points POST plus one
//! series POST to Datadog. Submission failures are logged, not
//! propagated.
//!
//! Mirrors the `pager` library's collector/worker split. Build with
//! [`builder::MetricsBuilder`].
pub mod builder;
pub mod collector;
pub mod queue;
pub mod sample;
pub mod worker;
+107
View File
@@ -0,0 +1,107 @@
use std::sync::{Arc, Mutex};
use log::warn;
use crate::sample::Sample;
/// Soft cap on buffered samples. The worker drains every few seconds, so
/// hitting this would mean the Datadog API is unreachable AND we're seeing
/// >10k req/s sustained. Hitting it drops oldest-first.
pub const DEFAULT_MAX_SAMPLES: usize = 100_000;
/// Thread-safe bounded buffer of samples. Cheap to clone (Arc inside).
#[derive(Clone)]
pub struct SampleQueue {
inner: Arc<Mutex<Vec<Sample>>>,
max_samples: usize,
}
impl SampleQueue {
pub fn new() -> Self {
Self::with_capacity(DEFAULT_MAX_SAMPLES)
}
pub fn with_capacity(max_samples: usize) -> Self {
Self {
inner: Arc::new(Mutex::new(Vec::with_capacity(1024))),
max_samples,
}
}
/// Non-blocking push. Drops the *oldest* sample if the buffer is full so
/// the worker still gets a freshness-biased view if Datadog is down.
pub fn push(&self, sample: Sample) {
let mut guard = match self.inner.lock() {
Ok(g) => g,
Err(poison) => poison.into_inner(),
};
if guard.len() >= self.max_samples {
// Drop oldest. `swap_remove(0)` is O(1) but reorders; we don't care
// about insertion order since aggregation groups by tuple key.
let _ = guard.swap_remove(0);
}
guard.push(sample);
}
/// Atomically take all currently-buffered samples.
pub fn drain(&self) -> Vec<Sample> {
let mut guard = match self.inner.lock() {
Ok(g) => g,
Err(poison) => {
warn!("metrics SampleQueue lock poisoned on drain; recovering");
poison.into_inner()
}
};
std::mem::take(&mut *guard)
}
pub fn len(&self) -> usize {
self.inner.lock().map(|g| g.len()).unwrap_or(0)
}
}
impl Default for SampleQueue {
fn default() -> Self { Self::new() }
}
#[cfg(test)]
mod tests {
use super::*;
use crate::sample::RequestSample;
fn request_sample(route: &str) -> Sample {
Sample::Request(RequestSample {
route: route.to_string(),
method: "GET".to_string(),
status_code: 200,
duration_ms: 1.0,
timestamp_secs: 1,
})
}
#[test]
fn drain_returns_pushed_samples_and_empties() {
let q = SampleQueue::new();
q.push(request_sample("/a"));
q.push(request_sample("/b"));
let drained = q.drain();
assert_eq!(drained.len(), 2);
assert_eq!(q.drain().len(), 0);
}
#[test]
fn full_queue_drops_oldest() {
let q = SampleQueue::with_capacity(3);
q.push(request_sample("/old"));
q.push(request_sample("/older"));
q.push(request_sample("/oldest"));
q.push(request_sample("/new"));
let drained = q.drain();
assert_eq!(drained.len(), 3);
let routes: Vec<_> = drained.iter().filter_map(|s| match s {
Sample::Request(r) => Some(r.route.clone()),
_ => None,
}).collect();
assert!(routes.contains(&"/new".to_string()));
}
}
+42
View File
@@ -0,0 +1,42 @@
/// A request observation captured by the HTTP middleware. Aggregated into
/// two metrics per `(route, method, status_code)` group: a distribution
/// over `duration_ms` values, and a count.
#[derive(Clone, Debug)]
pub struct RequestSample {
/// Low-cardinality route pattern from actix (`req.match_pattern()`), e.g.
/// `/v1/users/{user_token}/profile`. Falls back to `"unmatched"` when
/// the request hit no route.
pub route: String,
pub method: String,
pub status_code: u16,
pub duration_ms: f64,
pub timestamp_secs: i64,
}
/// A single hit on a domain counter. Aggregated into a `count` metric per
/// `(metric, tags_set)` group.
#[derive(Clone, Debug)]
pub struct CounterSample {
pub metric: String,
pub tags: Vec<String>,
pub timestamp_secs: i64,
}
/// A single observation of a domain value. Aggregated into a `distribution`
/// metric per `(metric, tags_set)` group so Datadog can compute percentiles
/// server-side.
#[derive(Clone, Debug)]
pub struct ObservationSample {
pub metric: String,
pub value: f64,
pub tags: Vec<String>,
pub timestamp_secs: i64,
}
/// Wire-level sample union pushed onto the queue.
#[derive(Clone, Debug)]
pub enum Sample {
Request(RequestSample),
Counter(CounterSample),
Observation(ObservationSample),
}
+462
View File
@@ -0,0 +1,462 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use log::{debug, info, warn};
use datadog_client::client::{
DatadogClient, DistributionPoint, DistributionSeries, MetricPoint, MetricSeries,
};
use crate::queue::SampleQueue;
use crate::sample::{CounterSample, ObservationSample, RequestSample, Sample};
#[derive(Clone)]
pub struct MetricsWorkerConfig {
pub flush_interval: Duration,
pub duration_metric_name: String,
pub count_metric_name: String,
pub host: Option<String>,
pub env: String,
pub service_name: String,
}
/// Long-lived async loop: every `flush_interval`, drain the queue, aggregate
/// each sample kind, and POST to Datadog. Stops when the shared `shutdown`
/// flag is set and flushes one final batch on the way out.
pub struct MetricsWorker {
queue: SampleQueue,
client: DatadogClient,
config: MetricsWorkerConfig,
shutdown: Arc<AtomicBool>,
}
impl MetricsWorker {
pub(crate) fn new(
queue: SampleQueue,
client: DatadogClient,
config: MetricsWorkerConfig,
shutdown: Arc<AtomicBool>,
) -> Self {
Self { queue, client, config, shutdown }
}
pub fn shutdown_handle(&self) -> Arc<AtomicBool> {
self.shutdown.clone()
}
pub async fn run(&self) {
info!(
"Metrics worker started (flush every {:?}, service={}, env={})",
self.config.flush_interval, self.config.service_name, self.config.env,
);
let mut ticker = tokio::time::interval(self.config.flush_interval);
// `tokio::time::interval` fires immediately on the first tick; skip it
// so we don't flush a half-empty bucket the moment we boot.
ticker.tick().await;
while !self.shutdown.load(Ordering::Relaxed) {
ticker.tick().await;
if self.shutdown.load(Ordering::Relaxed) { break; }
self.flush_once().await;
}
info!("Metrics worker draining final batch on shutdown");
self.flush_once().await;
info!("Metrics worker stopped");
}
async fn flush_once(&self) {
let samples = self.queue.drain();
if samples.is_empty() {
return;
}
let sample_count = samples.len();
let MetricsBuild { distributions, counts } = self.build_metrics(samples);
debug!(
"metrics flush: {} samples → {} distribution series, {} count series",
sample_count, distributions.len(), counts.len(),
);
if let Err(e) = self.client.submit_distribution_points(distributions).await {
warn!("metrics: distribution submission failed: {}", e);
}
if let Err(e) = self.client.submit_series(counts).await {
warn!("metrics: series submission failed: {}", e);
}
}
/// Partition samples by kind and aggregate each:
/// - Request samples → per `(route, method, status_code)`: distribution of
/// duration_ms + count.
/// - Counter samples → per `(metric, sorted_tags)`: count (entry count
/// in that group).
/// - Observation samples → per `(metric, sorted_tags)`: distribution of
/// values.
fn build_metrics(&self, samples: Vec<Sample>) -> MetricsBuild {
let mut requests = Vec::new();
let mut counters = Vec::new();
let mut observations = Vec::new();
for s in samples {
match s {
Sample::Request(r) => requests.push(r),
Sample::Counter(c) => counters.push(c),
Sample::Observation(o) => observations.push(o),
}
}
let interval = self.config.flush_interval.as_secs() as i64;
let mut distributions = Vec::new();
let mut counts = Vec::new();
self.build_request_metrics(requests, interval, &mut distributions, &mut counts);
self.build_counter_metrics(counters, interval, &mut counts);
self.build_observation_metrics(observations, &mut distributions);
MetricsBuild { distributions, counts }
}
fn build_request_metrics(
&self,
samples: Vec<RequestSample>,
interval: i64,
distributions: &mut Vec<DistributionSeries>,
counts: &mut Vec<MetricSeries>,
) {
let mut by_key: HashMap<RequestKey, RequestBucket> = HashMap::new();
for s in samples {
let key = RequestKey { route: s.route, method: s.method, status_code: s.status_code };
let bucket = by_key.entry(key).or_insert_with(|| RequestBucket {
durations_ms: Vec::new(),
last_timestamp_secs: s.timestamp_secs,
});
bucket.durations_ms.push(s.duration_ms);
if s.timestamp_secs > bucket.last_timestamp_secs {
bucket.last_timestamp_secs = s.timestamp_secs;
}
}
for (key, bucket) in by_key {
let tags = self.build_request_tags(&key);
let count = bucket.durations_ms.len() as f64;
distributions.push(DistributionSeries {
metric: self.config.duration_metric_name.clone(),
points: vec![DistributionPoint(bucket.last_timestamp_secs, bucket.durations_ms)],
tags: tags.clone(),
host: self.config.host.clone(),
});
counts.push(MetricSeries {
metric: self.config.count_metric_name.clone(),
points: vec![MetricPoint(bucket.last_timestamp_secs, count)],
metric_type: "count".to_string(),
tags,
host: self.config.host.clone(),
interval: Some(interval),
});
}
}
fn build_counter_metrics(
&self,
samples: Vec<CounterSample>,
interval: i64,
counts: &mut Vec<MetricSeries>,
) {
let mut by_key: HashMap<(String, Vec<String>), CounterBucket> = HashMap::new();
for s in samples {
let tags = sort_tags(s.tags);
let entry = by_key.entry((s.metric, tags)).or_insert_with(|| CounterBucket {
value: 0.0,
last_timestamp_secs: s.timestamp_secs,
});
entry.value += 1.0;
if s.timestamp_secs > entry.last_timestamp_secs {
entry.last_timestamp_secs = s.timestamp_secs;
}
}
for ((metric, tags), bucket) in by_key {
counts.push(MetricSeries {
metric,
points: vec![MetricPoint(bucket.last_timestamp_secs, bucket.value)],
metric_type: "count".to_string(),
tags: with_service_env_tags(tags, &self.config),
host: self.config.host.clone(),
interval: Some(interval),
});
}
}
fn build_observation_metrics(
&self,
samples: Vec<ObservationSample>,
distributions: &mut Vec<DistributionSeries>,
) {
let mut by_key: HashMap<(String, Vec<String>), ObservationBucket> = HashMap::new();
for s in samples {
let tags = sort_tags(s.tags);
let entry = by_key.entry((s.metric, tags)).or_insert_with(|| ObservationBucket {
values: Vec::new(),
last_timestamp_secs: s.timestamp_secs,
});
entry.values.push(s.value);
if s.timestamp_secs > entry.last_timestamp_secs {
entry.last_timestamp_secs = s.timestamp_secs;
}
}
for ((metric, tags), bucket) in by_key {
distributions.push(DistributionSeries {
metric,
points: vec![DistributionPoint(bucket.last_timestamp_secs, bucket.values)],
tags: with_service_env_tags(tags, &self.config),
host: self.config.host.clone(),
});
}
}
fn build_request_tags(&self, key: &RequestKey) -> Vec<String> {
let class = status_class(key.status_code);
vec![
format!("route:{}", key.route),
format!("method:{}", key.method),
format!("status_code:{}", key.status_code),
format!("status_class:{}", class),
format!("env:{}", self.config.env),
format!("service:{}", self.config.service_name),
]
}
}
#[derive(Eq, Hash, PartialEq)]
struct RequestKey {
route: String,
method: String,
status_code: u16,
}
struct RequestBucket {
durations_ms: Vec<f64>,
last_timestamp_secs: i64,
}
struct CounterBucket {
value: f64,
last_timestamp_secs: i64,
}
struct ObservationBucket {
values: Vec<f64>,
last_timestamp_secs: i64,
}
struct MetricsBuild {
distributions: Vec<DistributionSeries>,
counts: Vec<MetricSeries>,
}
fn status_class(status: u16) -> &'static str {
match status {
100..=199 => "1xx",
200..=299 => "2xx",
300..=399 => "3xx",
400..=499 => "4xx",
_ => "5xx",
}
}
/// Sort so `(metric, tags)` keys hash the same regardless of caller ordering.
fn sort_tags(mut tags: Vec<String>) -> Vec<String> {
tags.sort();
tags.dedup();
tags
}
/// Append the service/env tags to a custom-metric tag list (after dedup),
/// matching what request metrics get for free.
fn with_service_env_tags(mut tags: Vec<String>, config: &MetricsWorkerConfig) -> Vec<String> {
tags.push(format!("env:{}", config.env));
tags.push(format!("service:{}", config.service_name));
tags
}
#[cfg(test)]
mod tests {
use super::*;
fn make_worker() -> MetricsWorker {
let queue = SampleQueue::new();
let client = DatadogClient::new(datadog_client::creds::DatadogApiKey::new("test"));
let config = MetricsWorkerConfig {
flush_interval: Duration::from_secs(10),
duration_metric_name: "test.duration_ms".to_string(),
count_metric_name: "test.count".to_string(),
host: Some("test-host".to_string()),
env: "test".to_string(),
service_name: "test-service".to_string(),
};
MetricsWorker::new(queue, client, config, Arc::new(AtomicBool::new(false)))
}
mod status_class_tests {
use super::*;
#[test]
fn classifies_each_range() {
assert_eq!(status_class(100), "1xx");
assert_eq!(status_class(200), "2xx");
assert_eq!(status_class(299), "2xx");
assert_eq!(status_class(301), "3xx");
assert_eq!(status_class(404), "4xx");
assert_eq!(status_class(500), "5xx");
assert_eq!(status_class(0), "5xx");
}
}
mod request_aggregation {
use super::*;
fn sample(route: &str, method: &str, status: u16, dur: f64, ts: i64) -> Sample {
Sample::Request(RequestSample {
route: route.to_string(),
method: method.to_string(),
status_code: status,
duration_ms: dur,
timestamp_secs: ts,
})
}
#[test]
fn groups_by_route_method_status() {
let w = make_worker();
let samples = vec![
sample("/a", "GET", 200, 1.0, 100),
sample("/a", "GET", 200, 2.0, 110),
sample("/a", "GET", 500, 9.0, 120),
sample("/b", "POST", 201, 5.0, 130),
];
let build = w.build_metrics(samples);
assert_eq!(build.distributions.len(), 3);
assert_eq!(build.counts.len(), 3);
let a_200 = build.distributions.iter()
.find(|d| d.tags.iter().any(|t| t == "route:/a")
&& d.tags.iter().any(|t| t == "status_code:200"))
.expect("missing /a 200 series");
assert_eq!(a_200.points[0].1, vec![1.0, 2.0]);
}
#[test]
fn tags_include_class_env_service() {
let w = make_worker();
let samples = vec![sample("/x", "GET", 500, 1.0, 1)];
let build = w.build_metrics(samples);
let tags = &build.distributions[0].tags;
assert!(tags.contains(&"status_class:5xx".to_string()));
assert!(tags.contains(&"env:test".to_string()));
assert!(tags.contains(&"service:test-service".to_string()));
}
#[test]
fn count_series_uses_count_type_and_interval() {
let w = make_worker();
let samples = vec![sample("/x", "GET", 200, 1.0, 1)];
let build = w.build_metrics(samples);
assert_eq!(build.counts[0].metric_type, "count");
assert_eq!(build.counts[0].interval, Some(10));
assert_eq!(build.counts[0].points[0].1, 1.0);
}
}
mod counter_aggregation {
use super::*;
fn counter(metric: &str, tags: Vec<&str>, ts: i64) -> Sample {
Sample::Counter(CounterSample {
metric: metric.to_string(),
tags: tags.into_iter().map(|s| s.to_string()).collect(),
timestamp_secs: ts,
})
}
#[test]
fn groups_by_metric_and_tag_set() {
let w = make_worker();
let samples = vec![
counter("orders.placed", vec!["region:us"], 1),
counter("orders.placed", vec!["region:us"], 2),
counter("orders.placed", vec!["region:eu"], 3),
counter("orders.failed", vec!["region:us"], 4),
];
let build = w.build_metrics(samples);
assert_eq!(build.distributions.len(), 0);
assert_eq!(build.counts.len(), 3);
let us_placed = build.counts.iter().find(|m| {
m.metric == "orders.placed"
&& m.tags.iter().any(|t| t == "region:us")
}).expect("missing orders.placed/us");
assert_eq!(us_placed.points[0].1, 2.0);
}
#[test]
fn tag_order_does_not_split_groups() {
let w = make_worker();
let samples = vec![
counter("x", vec!["a:1", "b:2"], 1),
counter("x", vec!["b:2", "a:1"], 2),
];
let build = w.build_metrics(samples);
assert_eq!(build.counts.len(), 1);
assert_eq!(build.counts[0].points[0].1, 2.0);
}
#[test]
fn service_and_env_tags_appended() {
let w = make_worker();
let samples = vec![counter("foo", vec![], 1)];
let build = w.build_metrics(samples);
let tags = &build.counts[0].tags;
assert!(tags.contains(&"env:test".to_string()));
assert!(tags.contains(&"service:test-service".to_string()));
}
}
mod observation_aggregation {
use super::*;
fn obs(metric: &str, value: f64, tags: Vec<&str>, ts: i64) -> Sample {
Sample::Observation(ObservationSample {
metric: metric.to_string(),
value,
tags: tags.into_iter().map(|s| s.to_string()).collect(),
timestamp_secs: ts,
})
}
#[test]
fn groups_into_distributions() {
let w = make_worker();
let samples = vec![
obs("queue.depth", 10.0, vec!["queue:a"], 1),
obs("queue.depth", 20.0, vec!["queue:a"], 2),
obs("queue.depth", 30.0, vec!["queue:b"], 3),
];
let build = w.build_metrics(samples);
assert_eq!(build.distributions.len(), 2);
assert_eq!(build.counts.len(), 0);
let a = build.distributions.iter()
.find(|d| d.tags.iter().any(|t| t == "queue:a"))
.expect("missing queue:a series");
assert_eq!(a.points[0].1, vec![10.0, 20.0]);
}
#[test]
fn service_and_env_tags_appended() {
let w = make_worker();
let samples = vec![obs("foo", 1.0, vec![], 1)];
let build = w.build_metrics(samples);
let tags = &build.distributions[0].tags;
assert!(tags.contains(&"env:test".to_string()));
assert!(tags.contains(&"service:test-service".to_string()));
}
}
}
@@ -1,4 +1,5 @@
pub mod logging;
pub mod metrics;
pub mod mysql;
pub mod paging;
pub mod redis;
@@ -0,0 +1,16 @@
// ----- Environment Variable Names -----
const ENV_DATADOG_API_KEY: &str = "DATADOG_API_KEY";
const ENV_ENABLE_METRICS: &str = "ENABLE_METRICS";
// ----- Read Environment Variables -----
pub fn env_optional_datadog_api_key() -> Option<String> {
easyenv::get_env_string_optional(ENV_DATADOG_API_KEY)
}
/// Independent of API key presence — lets ops kill metrics at runtime
/// without removing the secret.
pub fn env_enable_metrics_default_false() -> bool {
easyenv::get_env_bool_or_default(ENV_ENABLE_METRICS, false)
}
@@ -65,6 +65,8 @@ markdown = { path = "../../../lib/markdown" }
media = { path = "../../../lib/files/media" }
memory_caching = { path = "../../../lib/caching/memory_caching" }
migration = { path = "../../../schema/database/migration" }
datadog_client.workspace = true
metrics.workspace = true
mimetypes.workspace = true
mysql_queries.workspace = true
opaque_cursors.workspace = true
@@ -0,0 +1,107 @@
//! Per-request metrics middleware.
//!
//! Captures `(route_pattern, method, status_code, duration_ms)` for every
//! request and hands it to the [`MetricsCollector`], which buffers it for a
//! background worker. Submission happens off the request path — this
//! middleware never awaits anything except the inner service call itself.
//!
//! Cardinality control: we use [`actix_web::dev::ServiceResponse::request()`]
//! `.match_pattern()` (the registered route template like
//! `/v1/users/{user_token}`) rather than the raw URL, so tags don't explode
//! with per-user IDs. Requests that don't match a registered route are
//! tagged `route:unmatched`.
use std::future::{Future, Ready, ready};
use std::pin::Pin;
use std::time::Instant;
use actix_web::Error;
use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
use metrics::collector::MetricsCollector;
const UNMATCHED_ROUTE: &str = "unmatched";
#[derive(Clone)]
pub struct MetricsMiddleware {
collector: MetricsCollector,
}
impl MetricsMiddleware {
pub fn new(collector: MetricsCollector) -> Self {
Self { collector }
}
}
impl<S, B> Transform<S, ServiceRequest> for MetricsMiddleware
where
S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
S::Future: 'static,
B: 'static,
{
type Response = ServiceResponse<B>;
type Error = Error;
type InitError = ();
type Transform = MetricsService<S>;
type Future = Ready<Result<Self::Transform, Self::InitError>>;
fn new_transform(&self, service: S) -> Self::Future {
ready(Ok(MetricsService {
service,
collector: self.collector.clone(),
}))
}
}
pub struct MetricsService<S> {
service: S,
collector: MetricsCollector,
}
impl<S, B> Service<ServiceRequest> for MetricsService<S>
where
S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
S::Future: 'static,
B: 'static,
{
type Response = ServiceResponse<B>;
type Error = Error;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;
actix_service::forward_ready!(service);
fn call(&self, req: ServiceRequest) -> Self::Future {
// Skip everything when metrics are disabled — keeps the no-op path
// cheap (still pays one Instant::now and an awaited service call).
if !self.collector.is_enabled() {
return Box::pin(self.service.call(req));
}
let collector = self.collector.clone();
let method = req.method().to_string();
let started = Instant::now();
let future = self.service.call(req);
Box::pin(async move {
let result = future.await;
let duration_ms = started.elapsed().as_secs_f64() * 1000.0;
match &result {
Ok(response) => {
let route = response.request()
.match_pattern()
.unwrap_or_else(|| UNMATCHED_ROUTE.to_string());
let status_code = response.status().as_u16();
collector.record_request(route, method, status_code, duration_ms);
}
Err(_) => {
// The service returned an error before producing a response, so
// we can't recover match_pattern. Surface this as 500/unmatched
// — it's rare in practice (most actix error paths still produce
// a ServiceResponse).
collector.record_request(UNMATCHED_ROUTE, method, 500, duration_ms);
}
}
result
})
}
}
@@ -0,0 +1 @@
pub mod metrics_middleware;
@@ -1,2 +1,3 @@
pub mod error_alerting_middleware;
pub mod metrics_middleware;
pub mod pushback_filter_middleware;
+40 -2
View File
@@ -45,11 +45,13 @@ use crate::billing::internal_session_cache_purge_impl::InternalSessionCachePurge
use crate::billing::stripe_internal_subscription_product_lookup_impl::StripeInternalSubscriptionProductLookupImpl;
use crate::billing::stripe_internal_user_lookup_impl::StripeInternalUserLookupImpl;
use crate::http_server::middleware::error_alerting_middleware::error_alerting_middleware::ErrorAlertingMiddleware;
use crate::http_server::middleware::metrics_middleware::metrics_middleware::MetricsMiddleware;
use crate::http_server::middleware::pushback_filter_middleware::PushbackFilter;
use crate::http_server::routes::add_routes::add_routes;
use crate::http_server::web_utils::handle_multipart_error::handle_multipart_error;
use crate::startup::build_dependencies::setup_dependencies;
use crate::startup::setup_disabled_endpoints::read_disabled_endpoints;
use crate::startup::setup_metrics::build_metrics;
use crate::state::server_state::ServerState;
use crate::threads::db_health_checker_thread::db_health_checker_thread::db_health_checker_thread;
use crate::threads::poll_ip_banlist_thread::poll_ip_bans;
@@ -144,13 +146,48 @@ async fn main() -> AnyhowResult<()> {
poll_model_token_info_thread(model_token_info_cache_clone, mysql_pool_clone).await;
});
// ==================== Metrics worker ==================== //
let (metrics_collector, maybe_metrics_worker) =
build_metrics(server_state.server_environment, &server_state.hostname);
let maybe_metrics_shutdown = maybe_metrics_worker.map(|worker| {
info!("Spawning metrics worker thread.");
let shutdown = worker.shutdown_handle();
tokio_runtime.spawn(async move {
worker.run().await;
});
shutdown
});
// ==================== Serve ==================== //
serve(server_state).await?;
serve(server_state, metrics_collector).await?;
// ==================== Graceful drain ==================== //
//
// serve() returns once actix has finished its own shutdown (SIGTERM,
// SIGINT, etc.). Tell the metrics worker to flush its last batch and
// exit, giving it a bounded grace period before we drop the runtime.
if let Some(shutdown) = maybe_metrics_shutdown {
use std::sync::atomic::Ordering;
info!("Signaling metrics worker to drain before exit.");
shutdown.store(true, Ordering::Relaxed);
// The worker drains on the next tick. One flush_interval (default 10s)
// is plenty for it to wake, drain, and POST one last batch. We don't
// join the spawned task directly because the runtime owns it.
tokio_runtime.block_on(async {
tokio::time::sleep(std::time::Duration::from_secs(12)).await;
});
}
Ok(())
}
pub async fn serve(server_state: ServerState) -> AnyhowResult<()>
pub async fn serve(
server_state: ServerState,
metrics_collector: metrics::collector::MetricsCollector,
) -> AnyhowResult<()>
{
let bind_address = server_state.env_config.bind_address.clone();
let num_workers = server_state.env_config.num_workers;
@@ -224,6 +261,7 @@ pub async fn serve(server_state: ServerState) -> AnyhowResult<()>
.memory_limit(10 * 1024 * 1024) // 10 MB
.error_handler(handle_multipart_error)
)
.wrap(MetricsMiddleware::new(metrics_collector.clone()))
.wrap(build_cors_config(old_server_environment))
.wrap(shared_array_buffer_cors())
.wrap(DefaultHeaders::new()
@@ -1,5 +1,6 @@
pub mod build_dependencies;
pub mod setup_bans;
pub mod setup_metrics;
pub mod setup_pager;
pub mod setup_disabled_endpoints;
pub mod setup_static_feature_flags;
@@ -0,0 +1,57 @@
use log::{info, warn};
use datadog_client::client::DatadogClient;
use datadog_client::creds::DatadogApiKey;
use metrics::builder::MetricsBuilder;
use metrics::collector::MetricsCollector;
use metrics::worker::MetricsWorker;
use shared_env_var_config::metrics::{env_enable_metrics_default_false, env_optional_datadog_api_key};
/// Returns `(collector, Option<worker>)`. When metrics are disabled or no
/// API key is set, returns a noop collector and no worker — wire the
/// collector into the middleware unconditionally so the request path
/// stays identical in both modes. Callers that want a graceful shutdown
/// drain can grab `worker.shutdown_handle()` before spawning the worker.
pub fn build_metrics(
server_environment: server_environment::ServerEnvironment,
hostname: &str,
) -> (MetricsCollector, Option<MetricsWorker>) {
if !env_enable_metrics_default_false() {
info!("ENABLE_METRICS=false (or unset). Metrics collection is disabled.");
return (MetricsCollector::noop(), None);
}
let Some(api_key) = env_optional_datadog_api_key() else {
warn!("ENABLE_METRICS=true but DATADOG_API_KEY is not set. Metrics will be a no-op.");
return (MetricsCollector::noop(), None);
};
let env_label = if server_environment.is_deployed_in_production() {
"production"
} else {
"development"
};
let client = DatadogClient::new(DatadogApiKey::new(api_key));
let build = MetricsBuilder::new()
.datadog_client(client)
.service_name("storyteller-web")
.env(env_label)
.host(hostname)
.build();
match build {
Ok(out) => {
info!(
"Metrics worker configured (service=storyteller-web, env={}, host={})",
env_label, hostname,
);
(out.collector, Some(out.worker))
}
Err(e) => {
warn!("Failed to build metrics worker: {}. Falling back to no-op.", e);
(MetricsCollector::noop(), None)
}
}
}