niom-turn/src/metrics.rs

use std::sync::OnceLock;
use std::sync::atomic::{AtomicU64, Ordering};

#[derive(Debug)]
pub struct Metrics {
    pub stun_messages_total: AtomicU64,
    pub channel_data_total: AtomicU64,
    pub stream_connections_total: AtomicU64,

    pub auth_challenge_total: AtomicU64,
    pub auth_stale_total: AtomicU64,
    pub auth_reject_total: AtomicU64,

    pub allocate_total: AtomicU64,
    pub allocate_success_total: AtomicU64,
    pub allocate_fail_total: AtomicU64,

    pub permissions_added_total: AtomicU64,
    pub channel_bindings_added_total: AtomicU64,

    pub allocations_active: AtomicU64,

    pub rate_limited_total: AtomicU64,
}

impl Default for Metrics {
    fn default() -> Self {
        Self {
            stun_messages_total: AtomicU64::new(0),
            channel_data_total: AtomicU64::new(0),
            stream_connections_total: AtomicU64::new(0),
            auth_challenge_total: AtomicU64::new(0),
            auth_stale_total: AtomicU64::new(0),
            auth_reject_total: AtomicU64::new(0),
            allocate_total: AtomicU64::new(0),
            allocate_success_total: AtomicU64::new(0),
            allocate_fail_total: AtomicU64::new(0),
            permissions_added_total: AtomicU64::new(0),
            channel_bindings_added_total: AtomicU64::new(0),
            allocations_active: AtomicU64::new(0),
            rate_limited_total: AtomicU64::new(0),
        }
    }
}

static METRICS: OnceLock<Metrics> = OnceLock::new();

pub fn metrics() -> &'static Metrics {
    METRICS.get_or_init(Metrics::default)
}

pub fn inc_stun_messages() {
    metrics().stun_messages_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_channel_data() {
    metrics().channel_data_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_stream_connections() {
    metrics().stream_connections_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_auth_challenge() {
    metrics().auth_challenge_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_auth_stale() {
    metrics().auth_stale_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_auth_reject() {
    metrics().auth_reject_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_allocate_total() {
    metrics().allocate_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_allocate_success() {
    metrics().allocate_success_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_allocate_fail() {
    metrics().allocate_fail_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_permission_added() {
    metrics().permissions_added_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_channel_binding_added() {
    metrics().channel_bindings_added_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_rate_limited() {
    metrics().rate_limited_total.fetch_add(1, Ordering::Relaxed);
}

pub fn inc_allocations_active() {
    metrics().allocations_active.fetch_add(1, Ordering::Relaxed);
}

pub fn dec_allocations_active() {
    // Saturating decrement to avoid underflow.
    let m = metrics();
    let mut current = m.allocations_active.load(Ordering::Relaxed);
    while current > 0 {
        match m.allocations_active.compare_exchange_weak(
            current,
            current - 1,
            Ordering::Relaxed,
            Ordering::Relaxed,
        ) {
            Ok(_) => return,
            Err(v) => current = v,
        }
    }
}

#[derive(Debug, Clone)]
pub struct MetricsSnapshot {
    pub stun_messages_total: u64,
    pub channel_data_total: u64,
    pub stream_connections_total: u64,
    pub auth_challenge_total: u64,
    pub auth_stale_total: u64,
    pub auth_reject_total: u64,
    pub allocate_total: u64,
    pub allocate_success_total: u64,
    pub allocate_fail_total: u64,
    pub permissions_added_total: u64,
    pub channel_bindings_added_total: u64,
    pub allocations_active: u64,
    pub rate_limited_total: u64,
}

pub fn snapshot() -> MetricsSnapshot {
    let m = metrics();
    MetricsSnapshot {
        stun_messages_total: m.stun_messages_total.load(Ordering::Relaxed),
        channel_data_total: m.channel_data_total.load(Ordering::Relaxed),
        stream_connections_total: m.stream_connections_total.load(Ordering::Relaxed),
        auth_challenge_total: m.auth_challenge_total.load(Ordering::Relaxed),
        auth_stale_total: m.auth_stale_total.load(Ordering::Relaxed),
        auth_reject_total: m.auth_reject_total.load(Ordering::Relaxed),
        allocate_total: m.allocate_total.load(Ordering::Relaxed),
        allocate_success_total: m.allocate_success_total.load(Ordering::Relaxed),
        allocate_fail_total: m.allocate_fail_total.load(Ordering::Relaxed),
        permissions_added_total: m.permissions_added_total.load(Ordering::Relaxed),
        channel_bindings_added_total: m.channel_bindings_added_total.load(Ordering::Relaxed),
        allocations_active: m.allocations_active.load(Ordering::Relaxed),
        rate_limited_total: m.rate_limited_total.load(Ordering::Relaxed),
    }
}