mentedb_server/
metrics.rs1use 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
26pub 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
38pub 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 #[cfg(target_os = "linux")]
49 {
50 let collector = prometheus::process_collector::ProcessCollector::for_self();
51 let _ = registry.register(Box::new(collector));
52 }
53
54 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
101pub 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
115pub 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}