use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
use std::time::Duration;
#[derive(Clone, Copy)]
pub enum Outcome {
Allowed,
Blocked,
RateLimited,
BadRequest,
UpstreamError,
InternalError,
}
impl Outcome {
const ALL: [Outcome; 6] = [
Outcome::Allowed,
Outcome::Blocked,
Outcome::RateLimited,
Outcome::BadRequest,
Outcome::UpstreamError,
Outcome::InternalError,
];
fn label(self) -> &'static str {
match self {
Outcome::Allowed => "allowed",
Outcome::Blocked => "blocked",
Outcome::RateLimited => "rate_limited",
Outcome::BadRequest => "bad_request",
Outcome::UpstreamError => "upstream_error",
Outcome::InternalError => "internal_error",
}
}
}
const BUCKETS_SECONDS: [f64; 14] =
[0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0];
pub struct Metrics {
requests: [AtomicU64; Outcome::ALL.len()],
hist_buckets: [AtomicU64; BUCKETS_SECONDS.len()],
hist_count: AtomicU64,
hist_sum_nanos: AtomicU64,
}
impl Metrics {
pub fn new() -> Self {
Self {
requests: std::array::from_fn(|_| AtomicU64::new(0)),
hist_buckets: std::array::from_fn(|_| AtomicU64::new(0)),
hist_count: AtomicU64::new(0),
hist_sum_nanos: AtomicU64::new(0),
}
}
pub fn record(&self, outcome: Outcome, elapsed: Duration) {
self.requests[outcome as usize].fetch_add(1, Relaxed);
let secs = elapsed.as_secs_f64();
let mut i = 0;
while i < BUCKETS_SECONDS.len() && secs > BUCKETS_SECONDS[i] {
i += 1;
}
if i < BUCKETS_SECONDS.len() {
self.hist_buckets[i].fetch_add(1, Relaxed);
}
self.hist_count.fetch_add(1, Relaxed);
self.hist_sum_nanos.fetch_add(elapsed.as_nanos() as u64, Relaxed);
}
pub fn render(&self) -> String {
let mut out = String::with_capacity(1024);
out.push_str("# HELP waf_up 1 if the WAF process is running.\n");
out.push_str("# TYPE waf_up gauge\n");
out.push_str("waf_up 1\n");
out.push_str("# HELP waf_build_info Build information.\n");
out.push_str("# TYPE waf_build_info gauge\n");
out.push_str(&format!(
"waf_build_info{{version=\"{}\"}} 1\n",
env!("CARGO_PKG_VERSION")
));
out.push_str("# HELP waf_requests_total Total requests by final decision.\n");
out.push_str("# TYPE waf_requests_total counter\n");
for outcome in Outcome::ALL {
let n = self.requests[outcome as usize].load(Relaxed);
out.push_str(&format!("waf_requests_total{{decision=\"{}\"}} {}\n", outcome.label(), n));
}
out.push_str("# HELP waf_request_duration_seconds Request handling latency in seconds.\n");
out.push_str("# TYPE waf_request_duration_seconds histogram\n");
let mut cumulative = 0u64;
for (i, bound) in BUCKETS_SECONDS.iter().enumerate() {
cumulative += self.hist_buckets[i].load(Relaxed);
out.push_str(&format!(
"waf_request_duration_seconds_bucket{{le=\"{}\"}} {}\n",
bound, cumulative
));
}
let count = self.hist_count.load(Relaxed);
out.push_str(&format!(
"waf_request_duration_seconds_bucket{{le=\"+Inf\"}} {}\n",
count
));
let sum_seconds = self.hist_sum_nanos.load(Relaxed) as f64 / 1e9;
out.push_str(&format!("waf_request_duration_seconds_sum {}\n", sum_seconds));
out.push_str(&format!("waf_request_duration_seconds_count {}\n", count));
out
}
}
impl Default for Metrics {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn render_has_all_decision_labels_and_histogram_shape() {
let m = Metrics::new();
m.record(Outcome::Allowed, Duration::from_millis(2));
m.record(Outcome::Blocked, Duration::from_micros(300));
m.record(Outcome::RateLimited, Duration::from_secs(20));
let text = m.render();
for label in ["allowed", "blocked", "rate_limited", "bad_request", "upstream_error", "internal_error"] {
assert!(text.contains(&format!("decision=\"{label}\"")), "missing {label}");
}
assert!(text.contains("waf_requests_total{decision=\"allowed\"} 1"));
assert!(text.contains("waf_requests_total{decision=\"bad_request\"} 0"));
assert!(text.contains("le=\"+Inf\"} 3"));
assert!(text.contains("waf_request_duration_seconds_count 3"));
assert!(text.contains("waf_up 1"));
}
#[test]
fn buckets_are_cumulative_and_monotonic() {
let m = Metrics::new();
for _ in 0..3 {
m.record(Outcome::Allowed, Duration::from_micros(800)); }
let text = m.render();
assert!(text.contains("le=\"0.0005\"} 0"));
assert!(text.contains("le=\"0.001\"} 3"));
assert!(text.contains("le=\"0.01\"} 3"));
assert!(text.contains("le=\"+Inf\"} 3"));
}
}