Skip to main content

mentedb_server/
metrics.rs

1//! Prometheus metrics exposed at `GET /metrics` in the standard text format.
2//!
3//! Aggregate only, with no per-account or per-agent labels, so it is safe to
4//! scrape without auth. It covers process health (CPU and memory on Linux) plus
5//! MenteDB gauges (memories stored, uptime, live cluster nodes) and HTTP request
6//! rate and latency, which together answer the "is it healthy and how hard is it
7//! working" question a dashboard needs.
8
9use std::sync::Arc;
10use std::sync::OnceLock;
11use std::time::Instant;
12
13use axum::{
14    body::Body,
15    extract::{Request, State},
16    http::{StatusCode, header},
17    middleware::Next,
18    response::{IntoResponse, Response},
19};
20use prometheus::{
21    Encoder, HistogramVec, IntCounterVec, IntGauge, Registry, TextEncoder, histogram_opts, opts,
22};
23
24use crate::state::AppState;
25
26/// The metric set, held in a process-global registry.
27pub struct Metrics {
28    registry: Registry,
29    uptime: IntGauge,
30    memory_count: IntGauge,
31    live_nodes: IntGauge,
32    http_requests: IntCounterVec,
33    http_duration: HistogramVec,
34}
35
36static METRICS: OnceLock<Metrics> = OnceLock::new();
37
38/// The process-global metric set, initialized on first use.
39pub fn metrics() -> &'static Metrics {
40    METRICS.get_or_init(Metrics::new)
41}
42
43impl Metrics {
44    fn new() -> Self {
45        let registry = Registry::new();
46
47        // Process CPU and memory, populated on Linux via /proc.
48        #[cfg(target_os = "linux")]
49        {
50            let collector = prometheus::process_collector::ProcessCollector::for_self();
51            let _ = registry.register(Box::new(collector));
52        }
53
54        // `mentedb_up` is a constant 1; register it and let the registry own it.
55        let up = IntGauge::new("mentedb_up", "1 if the server is serving").unwrap();
56        up.set(1);
57        registry.register(Box::new(up)).ok();
58
59        let uptime = IntGauge::new("mentedb_uptime_seconds", "Seconds since start").unwrap();
60        let memory_count =
61            IntGauge::new("mentedb_memory_count", "Memories stored on this node").unwrap();
62        let live_nodes = IntGauge::new(
63            "mentedb_cluster_live_nodes",
64            "Live nodes in the gossip cluster (1 when sharding is off)",
65        )
66        .unwrap();
67        let http_requests = IntCounterVec::new(
68            opts!(
69                "mentedb_http_requests_total",
70                "HTTP requests by method and status"
71            ),
72            &["method", "status"],
73        )
74        .unwrap();
75        let http_duration = HistogramVec::new(
76            histogram_opts!(
77                "mentedb_http_request_duration_seconds",
78                "HTTP request latency in seconds"
79            ),
80            &["method"],
81        )
82        .unwrap();
83
84        for g in [&uptime, &memory_count, &live_nodes] {
85            registry.register(Box::new(g.clone())).ok();
86        }
87        registry.register(Box::new(http_requests.clone())).ok();
88        registry.register(Box::new(http_duration.clone())).ok();
89
90        Self {
91            registry,
92            uptime,
93            memory_count,
94            live_nodes,
95            http_requests,
96            http_duration,
97        }
98    }
99}
100
101/// Middleware: time every request and count it by method and response status.
102pub async fn track(req: Request, next: Next) -> Response {
103    let method = req.method().as_str().to_string();
104    let start = Instant::now();
105    let resp = next.run(req).await;
106    let m = metrics();
107    let status = resp.status().as_u16().to_string();
108    m.http_requests.with_label_values(&[&method, &status]).inc();
109    m.http_duration
110        .with_label_values(&[&method])
111        .observe(start.elapsed().as_secs_f64());
112    resp
113}
114
115/// `GET /metrics`: refresh the point-in-time gauges, then encode the registry.
116pub async fn handler(State(state): State<Arc<AppState>>) -> impl IntoResponse {
117    let m = metrics();
118    m.uptime.set(state.start_time.elapsed().as_secs() as i64);
119    m.memory_count.set(state.db.memory_count() as i64);
120    let live = match &state.cluster {
121        Some(c) => c.live_node_count().await as i64,
122        None => 1,
123    };
124    m.live_nodes.set(live);
125
126    let mut buf = Vec::new();
127    if TextEncoder::new()
128        .encode(&m.registry.gather(), &mut buf)
129        .is_err()
130    {
131        return (StatusCode::INTERNAL_SERVER_ERROR, "metric encode error").into_response();
132    }
133    (
134        [(header::CONTENT_TYPE, "text/plain; version=0.0.4")],
135        Body::from(buf),
136    )
137        .into_response()
138}