astraea_server/
metrics.rs

1//! Observability: metrics collection and health endpoints for AstraeaDB.
2//!
3//! Provides request counters, duration histograms, and system gauges
4//! in a format compatible with Prometheus scraping.
5
6use std::collections::HashMap;
7use std::sync::RwLock;
8use std::sync::atomic::{AtomicU64, Ordering};
9use std::time::{Duration, Instant, SystemTime};
10
11/// Collected metrics for the AstraeaDB server.
12pub struct ServerMetrics {
13    /// Total requests by type.
14    request_counts: RwLock<HashMap<String, AtomicU64>>,
15    /// Total errors by type.
16    error_counts: RwLock<HashMap<String, AtomicU64>>,
17    /// Request durations in microseconds (for histogram approximation).
18    request_durations: RwLock<Vec<(String, u64)>>,
19    /// Current active connections.
20    active_connections: AtomicU64,
21    /// Total connections accepted since startup.
22    total_connections: AtomicU64,
23    /// Server start time.
24    start_time: Instant,
25    /// Start timestamp for uptime reporting.
26    start_timestamp: u64,
27    /// Maximum duration entries before compaction.
28    max_duration_entries: usize,
29}
30
31impl ServerMetrics {
32    /// Create a new metrics collector.
33    pub fn new() -> Self {
34        Self {
35            request_counts: RwLock::new(HashMap::new()),
36            error_counts: RwLock::new(HashMap::new()),
37            request_durations: RwLock::new(Vec::new()),
38            active_connections: AtomicU64::new(0),
39            total_connections: AtomicU64::new(0),
40            start_time: Instant::now(),
41            start_timestamp: SystemTime::now()
42                .duration_since(SystemTime::UNIX_EPOCH)
43                .unwrap_or_default()
44                .as_secs(),
45            max_duration_entries: 100_000,
46        }
47    }
48
49    /// Record a request of the given type.
50    pub fn record_request(&self, request_type: &str) {
51        let counts = self.request_counts.read().unwrap();
52        if let Some(counter) = counts.get(request_type) {
53            counter.fetch_add(1, Ordering::Relaxed);
54            return;
55        }
56        drop(counts);
57
58        let mut counts = self.request_counts.write().unwrap();
59        counts
60            .entry(request_type.to_string())
61            .or_insert_with(|| AtomicU64::new(0))
62            .fetch_add(1, Ordering::Relaxed);
63    }
64
65    /// Record a request error.
66    pub fn record_error(&self, request_type: &str) {
67        let counts = self.error_counts.read().unwrap();
68        if let Some(counter) = counts.get(request_type) {
69            counter.fetch_add(1, Ordering::Relaxed);
70            return;
71        }
72        drop(counts);
73
74        let mut counts = self.error_counts.write().unwrap();
75        counts
76            .entry(request_type.to_string())
77            .or_insert_with(|| AtomicU64::new(0))
78            .fetch_add(1, Ordering::Relaxed);
79    }
80
81    /// Record the duration of a request in microseconds.
82    pub fn record_duration(&self, request_type: &str, duration: Duration) {
83        let micros = duration.as_micros() as u64;
84        let mut durations = self.request_durations.write().unwrap();
85        durations.push((request_type.to_string(), micros));
86        if durations.len() > self.max_duration_entries {
87            // Keep the most recent half.
88            let keep_from = durations.len() / 2;
89            durations.drain(..keep_from);
90        }
91    }
92
93    /// Increment active connection count.
94    pub fn connection_opened(&self) {
95        self.active_connections.fetch_add(1, Ordering::Relaxed);
96        self.total_connections.fetch_add(1, Ordering::Relaxed);
97    }
98
99    /// Decrement active connection count.
100    pub fn connection_closed(&self) {
101        self.active_connections.fetch_sub(1, Ordering::Relaxed);
102    }
103
104    /// Get current active connection count.
105    pub fn active_connections(&self) -> u64 {
106        self.active_connections.load(Ordering::Relaxed)
107    }
108
109    /// Get total connections since startup.
110    pub fn total_connections(&self) -> u64 {
111        self.total_connections.load(Ordering::Relaxed)
112    }
113
114    /// Get server uptime.
115    pub fn uptime(&self) -> Duration {
116        self.start_time.elapsed()
117    }
118
119    /// Export metrics in Prometheus text exposition format.
120    pub fn to_prometheus(&self) -> String {
121        let mut output = String::new();
122
123        // Request counts
124        output.push_str("# HELP astraea_requests_total Total number of requests by type.\n");
125        output.push_str("# TYPE astraea_requests_total counter\n");
126        let counts = self.request_counts.read().unwrap();
127        for (req_type, count) in counts.iter() {
128            let val = count.load(Ordering::Relaxed);
129            output.push_str(&format!(
130                "astraea_requests_total{{type=\"{}\"}} {}\n",
131                req_type, val
132            ));
133        }
134
135        // Error counts
136        output.push_str("# HELP astraea_errors_total Total number of errors by type.\n");
137        output.push_str("# TYPE astraea_errors_total counter\n");
138        let errors = self.error_counts.read().unwrap();
139        for (req_type, count) in errors.iter() {
140            let val = count.load(Ordering::Relaxed);
141            output.push_str(&format!(
142                "astraea_errors_total{{type=\"{}\"}} {}\n",
143                req_type, val
144            ));
145        }
146
147        // Connection gauges
148        output.push_str("# HELP astraea_active_connections Current active connections.\n");
149        output.push_str("# TYPE astraea_active_connections gauge\n");
150        output.push_str(&format!(
151            "astraea_active_connections {}\n",
152            self.active_connections()
153        ));
154
155        output.push_str("# HELP astraea_connections_total Total connections since startup.\n");
156        output.push_str("# TYPE astraea_connections_total counter\n");
157        output.push_str(&format!(
158            "astraea_connections_total {}\n",
159            self.total_connections()
160        ));
161
162        // Uptime
163        output.push_str("# HELP astraea_uptime_seconds Server uptime in seconds.\n");
164        output.push_str("# TYPE astraea_uptime_seconds gauge\n");
165        output.push_str(&format!(
166            "astraea_uptime_seconds {}\n",
167            self.uptime().as_secs()
168        ));
169
170        // Duration summary (p50, p90, p99)
171        let durations = self.request_durations.read().unwrap();
172        if !durations.is_empty() {
173            output
174                .push_str("# HELP astraea_request_duration_us Request duration in microseconds.\n");
175            output.push_str("# TYPE astraea_request_duration_us summary\n");
176
177            // Aggregate by type
178            let mut by_type: HashMap<&str, Vec<u64>> = HashMap::new();
179            for (t, d) in durations.iter() {
180                by_type.entry(t.as_str()).or_default().push(*d);
181            }
182
183            for (req_type, mut vals) in by_type {
184                vals.sort_unstable();
185                let len = vals.len();
186                let p50 = vals[len / 2];
187                let p90 = vals[(len as f64 * 0.9) as usize];
188                let p99 = vals[((len as f64 * 0.99) as usize).min(len - 1)];
189                output.push_str(&format!(
190                    "astraea_request_duration_us{{type=\"{}\",quantile=\"0.5\"}} {}\n",
191                    req_type, p50
192                ));
193                output.push_str(&format!(
194                    "astraea_request_duration_us{{type=\"{}\",quantile=\"0.9\"}} {}\n",
195                    req_type, p90
196                ));
197                output.push_str(&format!(
198                    "astraea_request_duration_us{{type=\"{}\",quantile=\"0.99\"}} {}\n",
199                    req_type, p99
200                ));
201            }
202        }
203
204        output
205    }
206
207    /// Get a health check response.
208    pub fn health(&self) -> serde_json::Value {
209        serde_json::json!({
210            "status": "healthy",
211            "uptime_seconds": self.uptime().as_secs(),
212            "active_connections": self.active_connections(),
213            "total_connections": self.total_connections(),
214            "start_time": self.start_timestamp,
215        })
216    }
217}
218
219impl Default for ServerMetrics {
220    fn default() -> Self {
221        Self::new()
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use std::thread;
229
230    #[test]
231    fn record_and_count_requests() {
232        let metrics = ServerMetrics::new();
233        metrics.record_request("CreateNode");
234        metrics.record_request("CreateNode");
235        metrics.record_request("GetNode");
236
237        let prom = metrics.to_prometheus();
238        assert!(prom.contains("astraea_requests_total{type=\"CreateNode\"} 2"));
239        assert!(prom.contains("astraea_requests_total{type=\"GetNode\"} 1"));
240    }
241
242    #[test]
243    fn record_errors() {
244        let metrics = ServerMetrics::new();
245        metrics.record_error("CreateNode");
246        let prom = metrics.to_prometheus();
247        assert!(prom.contains("astraea_errors_total{type=\"CreateNode\"} 1"));
248    }
249
250    #[test]
251    fn connection_tracking() {
252        let metrics = ServerMetrics::new();
253        assert_eq!(metrics.active_connections(), 0);
254        metrics.connection_opened();
255        metrics.connection_opened();
256        assert_eq!(metrics.active_connections(), 2);
257        assert_eq!(metrics.total_connections(), 2);
258        metrics.connection_closed();
259        assert_eq!(metrics.active_connections(), 1);
260        assert_eq!(metrics.total_connections(), 2);
261    }
262
263    #[test]
264    fn duration_recording() {
265        let metrics = ServerMetrics::new();
266        for i in 0..100 {
267            metrics.record_duration("Query", Duration::from_micros(i * 10));
268        }
269        let prom = metrics.to_prometheus();
270        assert!(prom.contains("astraea_request_duration_us{type=\"Query\",quantile=\"0.5\"}"));
271        assert!(prom.contains("astraea_request_duration_us{type=\"Query\",quantile=\"0.9\"}"));
272        assert!(prom.contains("astraea_request_duration_us{type=\"Query\",quantile=\"0.99\"}"));
273    }
274
275    #[test]
276    fn health_check() {
277        let metrics = ServerMetrics::new();
278        let health = metrics.health();
279        assert_eq!(health["status"], "healthy");
280        assert!(health["uptime_seconds"].as_u64().is_some());
281    }
282
283    #[test]
284    fn uptime_increases() {
285        let metrics = ServerMetrics::new();
286        thread::sleep(Duration::from_millis(10));
287        assert!(metrics.uptime() >= Duration::from_millis(10));
288    }
289
290    #[test]
291    fn prometheus_format_valid() {
292        let metrics = ServerMetrics::new();
293        metrics.record_request("Ping");
294        let prom = metrics.to_prometheus();
295        // Should have proper Prometheus format with # HELP and # TYPE lines
296        assert!(prom.contains("# HELP astraea_requests_total"));
297        assert!(prom.contains("# TYPE astraea_requests_total counter"));
298        assert!(prom.contains("# HELP astraea_active_connections"));
299        assert!(prom.contains("# TYPE astraea_active_connections gauge"));
300    }
301}
astraea_server/metrics.rs

astraea_server/
metrics.rs
