solidb 1.2.2

//! Prometheus metrics endpoint for SoliDB
//!
//! Exposes metrics in Prometheus text format at /metrics

use axum::{
    extract::State,
    http::{header, HeaderMap, StatusCode},
    response::IntoResponse,
};
use std::sync::atomic::Ordering;

use super::handlers::AppState;

/// jemalloc's own accounting, in bytes.
///
/// `resident` is the number to compare against RSS; `retained` is address
/// space handed back to the OS but kept mapped, which is what makes the
/// virtual size enormous (113 GB was observed on a process holding 21.7 GB
/// resident) without that being a leak. The gap between `resident` and
/// `allocated` is fragmentation plus allocator metadata — worth knowing before
/// blaming RocksDB for memory the allocator is merely holding.
///
/// Returns `None` when jemalloc is not the active allocator. The
/// `#[global_allocator]` lives in `main.rs`, so under `cargo test` — which
/// links the library without it — these mallctls are absent and this reports
/// nothing rather than zeros that look like real measurements.
#[cfg(not(target_env = "msvc"))]
fn jemalloc_stats() -> Option<[(&'static str, u64); 6]> {
    use tikv_jemalloc_ctl::{epoch, stats};

    // Statistics are cached until the epoch is advanced; without this every
    // scrape would return the values from process start.
    epoch::advance().ok()?;

    Some([
        ("allocated", stats::allocated::read().ok()? as u64),
        ("active", stats::active::read().ok()? as u64),
        ("resident", stats::resident::read().ok()? as u64),
        ("mapped", stats::mapped::read().ok()? as u64),
        ("retained", stats::retained::read().ok()? as u64),
        ("metadata", stats::metadata::read().ok()? as u64),
    ])
}

#[cfg(target_env = "msvc")]
fn jemalloc_stats() -> Option<[(&'static str, u64); 6]> {
    None
}

fn metrics_authorized(headers: &HeaderMap) -> bool {
    if std::env::var("SOLIDB_METRICS_PUBLIC")
        .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
        .unwrap_or(false)
    {
        return true;
    }
    if let Ok(expected) = std::env::var("SOLIDB_METRICS_TOKEN") {
        if !expected.is_empty() {
            let provided = headers
                .get("X-Metrics-Token")
                .and_then(|h| h.to_str().ok())
                .or_else(|| {
                    headers
                        .get("Authorization")
                        .and_then(|h| h.to_str().ok())
                        .and_then(|v| v.strip_prefix("Bearer "))
                })
                .unwrap_or("");
            return crate::server::auth::constant_time_eq(provided.as_bytes(), expected.as_bytes());
        }
    }
    if let Some(token) = headers
        .get("Authorization")
        .and_then(|h| h.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "))
    {
        if let Ok(claims) = crate::server::auth::AuthService::validate_token(token) {
            return claims
                .roles
                .as_ref()
                .is_some_and(|r| r.iter().any(|role| role == "admin"));
        }
    }
    false
}

/// Prometheus metrics handler.
///
/// Denied unless `SOLIDB_METRICS_PUBLIC=1`, a matching `SOLIDB_METRICS_TOKEN`
/// is presented, or a valid admin JWT is supplied.
pub async fn metrics_handler(
    State(state): State<AppState>,
    headers: HeaderMap,
) -> impl IntoResponse {
    if !metrics_authorized(&headers) {
        return (StatusCode::UNAUTHORIZED, "metrics require authentication").into_response();
    }
    let mut output = String::new();

    // HTTP Requests Total
    let request_count = state.request_counter.load(Ordering::Relaxed);
    output.push_str("# HELP solidb_http_requests_total Total number of HTTP requests processed\n");
    output.push_str("# TYPE solidb_http_requests_total counter\n");
    output.push_str(&format!("solidb_http_requests_total {}\n\n", request_count));

    // Uptime
    let uptime_secs = state.startup_time.elapsed().as_secs_f64();
    output.push_str("# HELP solidb_uptime_seconds Time since server started in seconds\n");
    output.push_str("# TYPE solidb_uptime_seconds gauge\n");
    output.push_str(&format!("solidb_uptime_seconds {:.3}\n\n", uptime_secs));

    // System Metrics (CPU, Memory)
    {
        let mut system = state.system_monitor.lock().unwrap();
        system.refresh_cpu_all();
        system.refresh_memory();

        // CPU Usage
        let cpu_usage = system.global_cpu_usage();
        output.push_str("# HELP solidb_cpu_usage_percent Current CPU usage percentage\n");
        output.push_str("# TYPE solidb_cpu_usage_percent gauge\n");
        output.push_str(&format!("solidb_cpu_usage_percent {:.2}\n\n", cpu_usage));

        // Memory Usage
        let total_memory = system.total_memory();
        let used_memory = system.used_memory();
        let available_memory = system.available_memory();

        output.push_str("# HELP solidb_memory_total_bytes Total system memory in bytes\n");
        output.push_str("# TYPE solidb_memory_total_bytes gauge\n");
        output.push_str(&format!("solidb_memory_total_bytes {}\n\n", total_memory));

        output.push_str("# HELP solidb_memory_used_bytes Used system memory in bytes\n");
        output.push_str("# TYPE solidb_memory_used_bytes gauge\n");
        output.push_str(&format!("solidb_memory_used_bytes {}\n\n", used_memory));

        output.push_str("# HELP solidb_memory_available_bytes Available system memory in bytes\n");
        output.push_str("# TYPE solidb_memory_available_bytes gauge\n");
        output.push_str(&format!(
            "solidb_memory_available_bytes {}\n\n",
            available_memory
        ));
    }

    // Script Stats
    let active_scripts = state.script_stats.active_scripts.load(Ordering::Relaxed);
    let active_ws = state.script_stats.active_ws.load(Ordering::Relaxed);
    let total_scripts = state
        .script_stats
        .total_scripts_executed
        .load(Ordering::Relaxed);
    let total_ws = state
        .script_stats
        .total_ws_connections
        .load(Ordering::Relaxed);

    output.push_str("# HELP solidb_active_scripts Current number of active Lua scripts\n");
    output.push_str("# TYPE solidb_active_scripts gauge\n");
    output.push_str(&format!("solidb_active_scripts {}\n\n", active_scripts));

    output.push_str(
        "# HELP solidb_active_websockets Current number of active WebSocket connections\n",
    );
    output.push_str("# TYPE solidb_active_websockets gauge\n");
    output.push_str(&format!("solidb_active_websockets {}\n\n", active_ws));

    output.push_str("# HELP solidb_scripts_executed_total Total number of Lua scripts executed\n");
    output.push_str("# TYPE solidb_scripts_executed_total counter\n");
    output.push_str(&format!(
        "solidb_scripts_executed_total {}\n\n",
        total_scripts
    ));

    output.push_str(
        "# HELP solidb_websocket_connections_total Total number of WebSocket connections\n",
    );
    output.push_str("# TYPE solidb_websocket_connections_total counter\n");
    output.push_str(&format!(
        "solidb_websocket_connections_total {}\n\n",
        total_ws
    ));

    // Database Stats
    let databases = state.storage.list_databases();
    let db_count = databases.len();
    output.push_str("# HELP solidb_databases_total Number of databases\n");
    output.push_str("# TYPE solidb_databases_total gauge\n");
    output.push_str(&format!("solidb_databases_total {}\n\n", db_count));

    // Count total collections across all databases
    let mut total_collections = 0;
    for db_name in &databases {
        if let Ok(db) = state.storage.get_database(db_name) {
            let colls = db.list_collections();
            total_collections += colls.len();
        }
    }
    output.push_str(
        "# HELP solidb_collections_total Total number of collections across all databases\n",
    );
    output.push_str("# TYPE solidb_collections_total gauge\n");
    output.push_str(&format!(
        "solidb_collections_total {}\n\n",
        total_collections
    ));

    // Memory attribution. Added after a 613-collection instance was OOM-killed
    // at 21.7 GB RSS while holding 6.3 GB of data: several consumers here have
    // no ceiling (the per-collection write buffer with `--memtable-budget`
    // unset, and index/filter blocks pinned per SST with
    // `--bounded-index-cache` off), and without these counters there is no way
    // to tell which one grew. Computed on demand: it is one property read per
    // column family, so it must not go on a timer.
    let mem = state.storage.memory_breakdown();
    for (name, help, value) in [
        (
            "solidb_memtable_bytes",
            "Live memtable memory across all column families",
            mem.memtable_bytes,
        ),
        (
            "solidb_memtable_total_bytes",
            "Live plus immutable memtables; the gap over solidb_memtable_bytes is flush backlog",
            mem.memtable_total_bytes,
        ),
        (
            "solidb_table_readers_bytes",
            "Index and filter blocks pinned per open SST, outside the block cache",
            mem.table_readers_bytes,
        ),
        (
            "solidb_block_cache_bytes",
            "Shared block cache usage",
            mem.block_cache_bytes,
        ),
        (
            "solidb_block_cache_pinned_bytes",
            "Portion of the shared block cache that cannot be evicted",
            mem.block_cache_pinned_bytes,
        ),
        (
            "solidb_sst_files",
            "SST files across all levels of all column families",
            mem.sst_files,
        ),
        (
            "solidb_column_families",
            "Open column families, including default and _meta",
            mem.column_families,
        ),
        (
            "solidb_cached_collection_handles",
            "Collection handles held in the engine's unbounded cache",
            mem.cached_collection_handles,
        ),
    ] {
        output.push_str(&format!("# HELP {} {}\n", name, help));
        output.push_str(&format!("# TYPE {} gauge\n", name));
        output.push_str(&format!("{} {}\n\n", name, value));
    }

    // Column-family churn. Every create/drop rewrites and fsyncs the entire
    // OPTIONS file under the DB mutex, so these are the counters that explain
    // both OPTIONS growth and latency that no single query accounts for.
    let cf_ops = crate::storage::cf_ops::snapshot();
    for (name, help, value) in [
        (
            "solidb_cf_ops_total",
            "Column-family creates and drops since start",
            cf_ops.ops,
        ),
        (
            "solidb_cf_reuses_total",
            "Doomed column families wiped and reused instead of dropped and recreated",
            crate::storage::cf_ops::reuses(),
        ),
        (
            "solidb_collections_autocreated_total",
            "Collections brought into existence by a write to an unknown name",
            crate::storage::cf_ops::autocreates(),
        ),
    ] {
        output.push_str(&format!("# HELP {} {}\n", name, help));
        output.push_str(&format!("# TYPE {} counter\n", name));
        output.push_str(&format!("{} {}\n\n", name, value));
    }

    output.push_str("# HELP solidb_cf_op_seconds_total Wall time spent inside column-family creates and drops\n");
    output.push_str("# TYPE solidb_cf_op_seconds_total counter\n");
    output.push_str(&format!(
        "solidb_cf_op_seconds_total {:.6}\n\n",
        cf_ops.nanos as f64 / 1e9
    ));

    // Allocator-level view, to separate live data from fragmentation and from
    // address space merely kept mapped.
    if let Some(stats) = jemalloc_stats() {
        for (field, value) in stats {
            output.push_str(&format!(
                "# HELP solidb_jemalloc_{}_bytes jemalloc stats.{}\n",
                field, field
            ));
            output.push_str(&format!("# TYPE solidb_jemalloc_{}_bytes gauge\n", field));
            output.push_str(&format!("solidb_jemalloc_{}_bytes {}\n\n", field, value));
        }
    }

    // Cluster Stats (if cluster mode is enabled)
    if let Some(ref cluster_manager) = state.cluster_manager {
        let healthy_nodes = cluster_manager.get_healthy_nodes();
        let healthy_count = healthy_nodes.len();

        output.push_str(
            "# HELP solidb_cluster_healthy_nodes Number of healthy nodes in the cluster\n",
        );
        output.push_str("# TYPE solidb_cluster_healthy_nodes gauge\n");
        output.push_str(&format!(
            "solidb_cluster_healthy_nodes {}\n\n",
            healthy_count
        ));

        // Local node ID (for identification)
        let local_node = cluster_manager.local_node_id();
        output.push_str("# HELP solidb_cluster_info Cluster information\n");
        output.push_str("# TYPE solidb_cluster_info gauge\n");
        output.push_str(&format!(
            "solidb_cluster_info{{node_id=\"{}\"}} 1\n\n",
            local_node
        ));

        // Replication lag for each peer
        let members = cluster_manager.state().get_all_members();
        let current_seq = if let Some(log) = &state.replication_log {
            log.current_sequence()
        } else {
            0
        };

        if current_seq > 0 {
            for member in members {
                if member.node.id != local_node {
                    let lag = current_seq.saturating_sub(member.last_sequence);
                    output.push_str(
                        "# HELP solidb_replication_lag_replicas Sequence lag for replicated data\n",
                    );
                    output.push_str("# TYPE solidb_replication_lag_replicas gauge\n");
                    output.push_str(&format!(
                        "solidb_replication_lag_replicas{{node_id=\"{}\",address=\"{}\"}} {}\n\n",
                        member.node.id, member.node.address, lag
                    ));
                }
            }
        }
    }

    // Shard Coordinator Stats (if sharding is enabled)
    if state.shard_coordinator.is_some() {
        output.push_str("# HELP solidb_sharding_enabled Whether sharding is enabled\n");
        output.push_str("# TYPE solidb_sharding_enabled gauge\n");
        output.push_str("solidb_sharding_enabled 1\n\n");
    }

    // Queue Stats (if queue worker is enabled)
    if state.queue_worker.is_some() {
        output.push_str("# HELP solidb_queue_worker_enabled Whether the queue worker is enabled\n");
        output.push_str("# TYPE solidb_queue_worker_enabled gauge\n");
        output.push_str("solidb_queue_worker_enabled 1\n\n");
    }

    // Return with proper content type for Prometheus
    (
        StatusCode::OK,
        [(
            header::CONTENT_TYPE,
            "text/plain; version=0.0.4; charset=utf-8",
        )],
        output,
    )
        .into_response()
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_prometheus_format() {
        // Test that output follows Prometheus format
        let output = "# HELP solidb_http_requests_total Total requests\n# TYPE solidb_http_requests_total counter\nsolidb_http_requests_total 42\n";
        assert!(output.contains("# HELP"));
        assert!(output.contains("# TYPE"));
        assert!(output.contains("counter"));
    }
}