goosefs-sdk 0.1.4

//! Client metrics registry for tracking counters and gauges.
//!
//! Provides a global, thread-safe registry for application metrics.
//! All counter/gauge mutations are atomic operations.

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

use dashmap::DashMap;

/// A simple counter metric that tracks a cumulative value.
/// Increments are atomic and safe for concurrent access.
#[derive(Default)]
pub struct Counter {
    v: AtomicI64,
}

impl Counter {
    /// Increment the counter by `n` bytes/items.
    /// This operation is atomic and will never be reordered or lost
    /// even in the presence of concurrent calls.
    #[inline]
    pub fn inc(&self, n: i64) {
        self.v.fetch_add(n, Ordering::Relaxed);
    }

    /// Get the current value of the counter.
    /// Note: This is a snapshot; the value may change immediately after.
    #[inline]
    pub fn get(&self) -> i64 {
        self.v.load(Ordering::Relaxed)
    }
}

/// A gauge metric that tracks a point-in-time value.
/// Can be set and read atomically.
#[derive(Default)]
pub struct Gauge {
    v: AtomicI64,
}

impl Gauge {
    /// Set the gauge to the given value.
    #[inline]
    pub fn set(&self, val: i64) {
        self.v.store(val, Ordering::Relaxed);
    }

    /// Get the current value of the gauge.
    #[inline]
    pub fn get(&self) -> i64 {
        self.v.load(Ordering::Relaxed)
    }
}

/// Global metrics registry.
/// Stores all named counters and gauges in thread-safe concurrent maps.
pub(crate) struct Registry {
    pub(crate) counters: DashMap<String, std::sync::Arc<Counter>>,
    pub(crate) gauges: DashMap<String, std::sync::Arc<Gauge>>,
}

impl Default for Registry {
    fn default() -> Self {
        Self {
            counters: DashMap::new(),
            gauges: DashMap::new(),
        }
    }
}

/// Global singleton registry, lazily initialized.
pub(crate) static REGISTRY: OnceLock<Registry> = OnceLock::new();

/// Initialize and return the global registry.
/// (Internal; used by counter/gauge factory functions and the reporter.)
pub(crate) fn get_registry() -> &'static Registry {
    REGISTRY.get_or_init(Registry::default)
}

/// Get or create a counter by name.
/// Returns an `Arc` that can be cheaply cloned and shared across threads.
pub fn counter(name: &str) -> std::sync::Arc<Counter> {
    let registry = get_registry();

    // Fast path: counter already exists
    if let Some(c) = registry.counters.get(name) {
        return c.value().clone();
    }

    // Slow path: create and insert
    let c = std::sync::Arc::new(Counter::default());
    // If another thread races and inserts first, we'll return theirs.
    registry
        .counters
        .entry(name.to_string())
        .or_insert(c.clone());

    // Get the final value (may be different due to race)
    registry
        .counters
        .get(name)
        .map(|entry| entry.value().clone())
        .unwrap_or(c)
}

/// Get or create a gauge by name.
/// Returns an `Arc` that can be cheaply cloned and shared across threads.
pub fn gauge(name: &str) -> std::sync::Arc<Gauge> {
    let registry = get_registry();

    // Fast path: gauge already exists
    if let Some(g) = registry.gauges.get(name) {
        return g.value().clone();
    }

    // Slow path: create and insert
    let g = std::sync::Arc::new(Gauge::default());
    // If another thread races and inserts first, we'll return theirs.
    registry.gauges.entry(name.to_string()).or_insert(g.clone());

    // Get the final value (may be different due to race)
    registry
        .gauges
        .get(name)
        .map(|entry| entry.value().clone())
        .unwrap_or(g)
}

/// Metric name constants, aligned with Java's MetricKey definitions.
/// Only metrics with `isClusterAggregated=true` should be reported in heartbeat.
pub mod name {
    // ── Throughput counters (cluster aggregated) ─────────────────────────────

    /// Local short-circuit read bytes (client reads from local Alluxio worker).
    /// Cluster aggregated: true
    pub const CLIENT_BYTES_READ_LOCAL: &str = "Client.BytesReadLocal";

    /// Local short-circuit write bytes (client writes to local Alluxio worker).
    /// Cluster aggregated: true
    pub const CLIENT_BYTES_WRITTEN_LOCAL: &str = "Client.BytesWrittenLocal";

    /// Client direct UFS write bytes (bypass Alluxio layer).
    /// Cluster aggregated: true
    pub const CLIENT_BYTES_WRITTEN_UFS: &str = "Client.BytesWrittenUfs";

    // ── RPC operation counters ───────────────────────────────────────────────

    /// Total number of file read operations (open + stream fully consumed or closed).
    pub const CLIENT_READ_OPS_TOTAL: &str = "Client.ReadOpsTotal";

    /// Total number of file write operations (create + complete).
    pub const CLIENT_WRITE_OPS_TOTAL: &str = "Client.WriteOpsTotal";

    /// Total number of getStatus RPCs to Master.
    pub const CLIENT_GET_STATUS_OPS: &str = "Client.GetStatusOps";

    /// Total number of listStatus RPCs to Master.
    pub const CLIENT_LIST_STATUS_OPS: &str = "Client.ListStatusOps";

    /// Total number of createFile RPCs to Master.
    pub const CLIENT_CREATE_FILE_OPS: &str = "Client.CreateFileOps";

    /// Total number of createDirectory RPCs to Master.
    pub const CLIENT_CREATE_DIR_OPS: &str = "Client.CreateDirOps";

    /// Total number of delete (remove) RPCs to Master.
    pub const CLIENT_DELETE_OPS: &str = "Client.DeleteOps";

    /// Total number of rename RPCs to Master.
    pub const CLIENT_RENAME_OPS: &str = "Client.RenameOps";

    // ── Error / failure counters ─────────────────────────────────────────────

    /// Total RPC failures (all types: network, timeout, server error).
    pub const CLIENT_RPC_ERRORS_TOTAL: &str = "Client.RpcErrorsTotal";

    /// RPC failures broken down by type — UNAUTHENTICATED errors.
    pub const CLIENT_RPC_AUTH_ERRORS: &str = "Client.RpcAuthErrors";

    /// RPC failures — connection refused / unavailable.
    pub const CLIENT_RPC_UNAVAILABLE_ERRORS: &str = "Client.RpcUnavailableErrors";

    /// Block read failures (stream error, incomplete, etc.).
    pub const CLIENT_READ_FAILURES: &str = "Client.ReadFailures";

    /// Block write failures.
    pub const CLIENT_WRITE_FAILURES: &str = "Client.WriteFailures";

    // ── Latency counters (cumulative microseconds, divide by ops for avg) ───

    /// Cumulative read latency in microseconds (from open stream to close/eof).
    pub const CLIENT_READ_LATENCY_US: &str = "Client.ReadLatencyUs";

    /// Cumulative write latency in microseconds (from create to complete).
    pub const CLIENT_WRITE_LATENCY_US: &str = "Client.WriteLatencyUs";

    /// Cumulative getStatus RPC latency in microseconds.
    pub const CLIENT_GET_STATUS_LATENCY_US: &str = "Client.GetStatusLatencyUs";

    /// Cumulative listStatus RPC latency in microseconds.
    pub const CLIENT_LIST_STATUS_LATENCY_US: &str = "Client.ListStatusLatencyUs";

    // ── Connection pool gauges ───────────────────────────────────────────────

    /// Number of active (cached) worker connections in the pool.
    pub const CLIENT_WORKER_CONNECTIONS_ACTIVE: &str = "Client.WorkerConnectionsActive";

    /// Total number of worker reconnects performed (counter).
    pub const CLIENT_WORKER_RECONNECTS_TOTAL: &str = "Client.WorkerReconnectsTotal";

    /// Total number of reconnects that were coalesced (deduplicated).
    pub const CLIENT_WORKER_RECONNECTS_COALESCED: &str = "Client.WorkerReconnectsCoalesced";

    // ── Block / data path gauges ─────────────────────────────────────────────

    /// Number of blocks currently being read concurrently (gauge).
    pub const CLIENT_BLOCKS_READ_IN_PROGRESS: &str = "Client.BlocksReadInProgress";

    /// Number of blocks currently being written concurrently (gauge).
    pub const CLIENT_BLOCKS_WRITTEN_IN_PROGRESS: &str = "Client.BlocksWrittenInProgress";

    /// Total blocks successfully read (counter).
    pub const CLIENT_BLOCKS_READ_TOTAL: &str = "Client.BlocksReadTotal";

    /// Total blocks successfully written (counter).
    pub const CLIENT_BLOCKS_WRITTEN_TOTAL: &str = "Client.BlocksWrittenTotal";
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn counter_inc_and_get() {
        let c = Counter::default();
        assert_eq!(c.get(), 0);

        c.inc(42);
        assert_eq!(c.get(), 42);

        c.inc(8);
        assert_eq!(c.get(), 50);

        c.inc(-10);
        assert_eq!(c.get(), 40);
    }

    #[test]
    fn counter_negative_increment() {
        let c = Counter::default();
        c.inc(-5);
        assert_eq!(c.get(), -5);
    }

    #[test]
    fn gauge_set_and_get() {
        let g = Gauge::default();
        assert_eq!(g.get(), 0);

        g.set(99);
        assert_eq!(g.get(), 99);

        g.set(-10);
        assert_eq!(g.get(), -10);
    }

    #[test]
    fn registry_counter_factory() {
        // First call creates counter
        let c1 = counter("my_counter");
        assert_eq!(c1.get(), 0);
        c1.inc(10);

        // Second call returns same Arc
        let c2 = counter("my_counter");
        assert_eq!(c2.get(), 10);

        // Different name gets different counter
        let c3 = counter("other_counter");
        assert_eq!(c3.get(), 0);
    }

    #[test]
    fn registry_gauge_factory() {
        // First call creates gauge
        let g1 = gauge("my_gauge");
        assert_eq!(g1.get(), 0);
        g1.set(55);

        // Second call returns same Arc
        let g2 = gauge("my_gauge");
        assert_eq!(g2.get(), 55);

        // Different name gets different gauge
        let g3 = gauge("other_gauge");
        assert_eq!(g3.get(), 0);
    }

    #[test]
    fn registry_counter_concurrent() {
        use std::thread;

        let c = counter("concurrent_counter");
        let mut handles = vec![];

        for _ in 0..10 {
            let c = c.clone();
            let handle = thread::spawn(move || {
                for _ in 0..1000 {
                    c.inc(1);
                }
            });
            handles.push(handle);
        }

        for h in handles {
            h.join().unwrap();
        }

        assert_eq!(c.get(), 10_000);
    }

    #[test]
    fn name_constants() {
        assert_eq!(name::CLIENT_BYTES_READ_LOCAL, "Client.BytesReadLocal");
        assert_eq!(name::CLIENT_BYTES_WRITTEN_LOCAL, "Client.BytesWrittenLocal");
        assert_eq!(name::CLIENT_BYTES_WRITTEN_UFS, "Client.BytesWrittenUfs");
    }

    #[test]
    fn name_constants_rpc_ops() {
        assert_eq!(name::CLIENT_READ_OPS_TOTAL, "Client.ReadOpsTotal");
        assert_eq!(name::CLIENT_WRITE_OPS_TOTAL, "Client.WriteOpsTotal");
        assert_eq!(name::CLIENT_GET_STATUS_OPS, "Client.GetStatusOps");
        assert_eq!(name::CLIENT_LIST_STATUS_OPS, "Client.ListStatusOps");
        assert_eq!(name::CLIENT_CREATE_FILE_OPS, "Client.CreateFileOps");
        assert_eq!(name::CLIENT_CREATE_DIR_OPS, "Client.CreateDirOps");
        assert_eq!(name::CLIENT_DELETE_OPS, "Client.DeleteOps");
        assert_eq!(name::CLIENT_RENAME_OPS, "Client.RenameOps");
    }

    #[test]
    fn name_constants_errors() {
        assert_eq!(name::CLIENT_RPC_ERRORS_TOTAL, "Client.RpcErrorsTotal");
        assert_eq!(name::CLIENT_RPC_AUTH_ERRORS, "Client.RpcAuthErrors");
        assert_eq!(
            name::CLIENT_RPC_UNAVAILABLE_ERRORS,
            "Client.RpcUnavailableErrors"
        );
        assert_eq!(name::CLIENT_READ_FAILURES, "Client.ReadFailures");
        assert_eq!(name::CLIENT_WRITE_FAILURES, "Client.WriteFailures");
    }

    #[test]
    fn name_constants_latency() {
        assert_eq!(name::CLIENT_READ_LATENCY_US, "Client.ReadLatencyUs");
        assert_eq!(name::CLIENT_WRITE_LATENCY_US, "Client.WriteLatencyUs");
        assert_eq!(
            name::CLIENT_GET_STATUS_LATENCY_US,
            "Client.GetStatusLatencyUs"
        );
        assert_eq!(
            name::CLIENT_LIST_STATUS_LATENCY_US,
            "Client.ListStatusLatencyUs"
        );
    }

    #[test]
    fn name_constants_pool_and_blocks() {
        assert_eq!(
            name::CLIENT_WORKER_CONNECTIONS_ACTIVE,
            "Client.WorkerConnectionsActive"
        );
        assert_eq!(
            name::CLIENT_WORKER_RECONNECTS_TOTAL,
            "Client.WorkerReconnectsTotal"
        );
        assert_eq!(
            name::CLIENT_WORKER_RECONNECTS_COALESCED,
            "Client.WorkerReconnectsCoalesced"
        );
        assert_eq!(
            name::CLIENT_BLOCKS_READ_IN_PROGRESS,
            "Client.BlocksReadInProgress"
        );
        assert_eq!(
            name::CLIENT_BLOCKS_WRITTEN_IN_PROGRESS,
            "Client.BlocksWrittenInProgress"
        );
        assert_eq!(name::CLIENT_BLOCKS_READ_TOTAL, "Client.BlocksReadTotal");
        assert_eq!(
            name::CLIENT_BLOCKS_WRITTEN_TOTAL,
            "Client.BlocksWrittenTotal"
        );
    }
}