sz_orm_observability/
lib.rs

1//! SZ-ORM 可观测性模块
2//!
3//! 提供 Prometheus exporter、SLO 燃烧率监控等能力。
4//! OTLP 导出由 `sz-orm-tracing` 的 `otlp` feature 提供（本模块不含 OTLP）。
5//!
6//! # 核心能力
7//!
8//! ## 1. MetricsRegistry（默认启用）
9//!
10//! 统一的指标注册中心，支持 Counter / Gauge / Histogram 三种类型，
11//! 内置线程安全（`RwLock`），可通过 `render()` 输出 Prometheus 文本格式。
12//!
13//! ## 2. Prometheus exporter
14//!
15//! 通过 `start_metrics_server` 在指定端口暴露 `/metrics` HTTP 端点，
16//! 供 Prometheus 拉取。基于 `tokio::net::TcpListener` 实现（不依赖 hyper）。
17//!
18//! ## 3. SLO 燃烧率
19//!
20//! 基于 5m / 1h 两个窗口计算 SLO 燃烧率，支持多窗口告警。
21//!
22//! # 快速入门
23//!
24//! ```no_run
25//! use sz_orm_observability::{MetricsRegistry, MetricKind};
26//! use std::time::Duration;
27//!
28//! // 创建指标注册中心
29//! let registry = MetricsRegistry::new();
30//!
31//! // 注册指标
32//! let counter = registry.register_counter("sz_orm_pool_acquires_total", "Total pool acquire calls");
33//! let gauge = registry.register_gauge("sz_orm_pool_active_connections", "Current active connections");
34//! let histogram = registry.register_histogram(
35//!     "sz_orm_query_duration_seconds",
36//!     "Query duration in seconds",
37//!     vec![0.001, 0.01, 0.1, 1.0, 10.0],
38//! );
39//!
40//! // 更新指标
41//! counter.inc();
42//! gauge.set(5.0);
43//! histogram.observe(0.025);
44//!
45//! // 输出 Prometheus 文本格式
46//! let output = registry.render();
47//! println!("{}", output);
48//! ```
49
50#![warn(missing_docs)]
51
52use parking_lot::RwLock;
53use std::collections::HashMap;
54use std::sync::Arc;
55
56pub mod slo;
57pub mod summary;
58
59pub use slo::{SloBurnRate, SloConfig, SloMonitor};
60pub use summary::{
61    LabeledHistogram, PushSnapshot, PushgatewayConfig, PushgatewayExporter, Summary,
62};
63
64/// 指标类型
65#[derive(Debug, Clone, Copy, PartialEq, Eq)]
66pub enum MetricKind {
67    /// 单调递增计数器（如总请求数）
68    Counter,
69    /// 可增可减的瞬时值（如当前连接数）
70    Gauge,
71    /// 直方图（如请求延迟分布）
72    Histogram,
73}
74
75/// 指标元数据
76#[derive(Debug, Clone)]
77pub struct MetricMeta {
78    /// 指标名（如 `sz_orm_pool_acquires_total`）
79    pub name: String,
80    /// 帮助文本
81    pub help: String,
82    /// 指标类型
83    pub kind: MetricKind,
84}
85
86/// 计数器（单调递增）
87pub struct Counter {
88    name: String,
89    value: Arc<RwLock<f64>>,
90    labels: HashMap<String, String>,
91}
92
93impl Counter {
94    /// 递增 1
95    pub fn inc(&self) {
96        self.inc_by(1.0);
97    }
98
99    /// 递增指定值
100    pub fn inc_by(&self, delta: f64) {
101        let mut v = self.value.write();
102        *v += delta;
103    }
104
105    /// 当前值
106    pub fn value(&self) -> f64 {
107        *self.value.read()
108    }
109
110    /// 指标名
111    pub fn name(&self) -> &str {
112        &self.name
113    }
114
115    /// 渲染为 Prometheus 文本格式
116    pub fn render(&self) -> String {
117        let v = self.value.read();
118        if self.labels.is_empty() {
119            format!("{} {}\n", self.name, v)
120        } else {
121            let labels: Vec<String> = self
122                .labels
123                .iter()
124                .map(|(k, val)| format!("{}=\"{}\"", k, val.replace('"', "\\\"")))
125                .collect();
126            format!("{}{{{}}} {}\n", self.name, labels.join(","), v)
127        }
128    }
129}
130
131/// Gauge（可增可减）
132pub struct Gauge {
133    name: String,
134    value: Arc<RwLock<f64>>,
135    labels: HashMap<String, String>,
136}
137
138impl Gauge {
139    /// 设置值
140    pub fn set(&self, value: f64) {
141        *self.value.write() = value;
142    }
143
144    /// 递增
145    pub fn inc(&self) {
146        self.inc_by(1.0);
147    }
148
149    /// 递增指定值
150    pub fn inc_by(&self, delta: f64) {
151        let mut v = self.value.write();
152        *v += delta;
153    }
154
155    /// 递减指定值
156    pub fn dec_by(&self, delta: f64) {
157        let mut v = self.value.write();
158        *v -= delta;
159    }
160
161    /// 当前值
162    pub fn value(&self) -> f64 {
163        *self.value.read()
164    }
165
166    /// 指标名
167    pub fn name(&self) -> &str {
168        &self.name
169    }
170
171    /// 渲染为 Prometheus 文本格式
172    pub fn render(&self) -> String {
173        let v = self.value.read();
174        if self.labels.is_empty() {
175            format!("{} {}\n", self.name, v)
176        } else {
177            let labels: Vec<String> = self
178                .labels
179                .iter()
180                .map(|(k, val)| format!("{}=\"{}\"", k, val.replace('"', "\\\"")))
181                .collect();
182            format!("{}{{{}}} {}\n", self.name, labels.join(","), v)
183        }
184    }
185}
186
187/// 直方图（延迟分布等）
188pub struct Histogram {
189    name: String,
190    buckets: Vec<f64>,
191    counts: Arc<RwLock<Vec<u64>>>,
192    sum: Arc<RwLock<f64>>,
193    count: Arc<RwLock<u64>>,
194}
195
196impl Histogram {
197    /// 观察一个值
198    pub fn observe(&self, value: f64) {
199        let mut counts = self.counts.write();
200        for (i, bucket) in self.buckets.iter().enumerate() {
201            if value <= *bucket {
202                counts[i] += 1;
203            }
204        }
205        // 最后一个 bucket 是 +Inf，必须递增
206        let last = counts.len() - 1;
207        counts[last] += 1;
208
209        let mut sum = self.sum.write();
210        *sum += value;
211        let mut count = self.count.write();
212        *count += 1;
213    }
214
215    /// 总观察次数
216    pub fn count(&self) -> u64 {
217        *self.count.read()
218    }
219
220    /// 所有观察值之和
221    pub fn sum(&self) -> f64 {
222        *self.sum.read()
223    }
224
225    /// 指标名
226    pub fn name(&self) -> &str {
227        &self.name
228    }
229
230    /// 渲染为 Prometheus 文本格式
231    pub fn render(&self) -> String {
232        let counts = self.counts.read();
233        let sum = self.sum.read();
234        let count = self.count.read();
235
236        let mut output = String::new();
237        for (i, bucket) in self.buckets.iter().enumerate() {
238            output.push_str(&format!(
239                "{}_bucket{{le=\"{}\"}} {}\n",
240                self.name, bucket, counts[i]
241            ));
242        }
243        output.push_str(&format!("{}_sum {}\n", self.name, sum));
244        output.push_str(&format!("{}_count {}\n", self.name, count));
245        output
246    }
247}
248
249/// 指标注册中心
250pub struct MetricsRegistry {
251    counters: RwLock<HashMap<String, Arc<Counter>>>,
252    gauges: RwLock<HashMap<String, Arc<Gauge>>>,
253    histograms: RwLock<HashMap<String, Arc<Histogram>>>,
254    metas: RwLock<Vec<MetricMeta>>,
255}
256
257impl Default for MetricsRegistry {
258    fn default() -> Self {
259        Self::new()
260    }
261}
262
263impl MetricsRegistry {
264    /// 创建空注册中心
265    pub fn new() -> Self {
266        Self {
267            counters: RwLock::new(HashMap::new()),
268            gauges: RwLock::new(HashMap::new()),
269            histograms: RwLock::new(HashMap::new()),
270            metas: RwLock::new(Vec::new()),
271        }
272    }
273
274    /// 注册 Counter
275    pub fn register_counter(&self, name: &str, help: &str) -> Arc<Counter> {
276        self.register_counter_with_labels(name, help, HashMap::new())
277    }
278
279    /// 注册带标签的 Counter
280    pub fn register_counter_with_labels(
281        &self,
282        name: &str,
283        help: &str,
284        labels: HashMap<String, String>,
285    ) -> Arc<Counter> {
286        let mut counters = self.counters.write();
287        let key = format!("{}_{:?}", name, labels);
288        if let Some(c) = counters.get(&key) {
289            return c.clone();
290        }
291        let counter = Arc::new(Counter {
292            name: name.to_string(),
293            value: Arc::new(RwLock::new(0.0)),
294            labels,
295        });
296        counters.insert(key, counter.clone());
297
298        let mut metas = self.metas.write();
299        metas.push(MetricMeta {
300            name: name.to_string(),
301            help: help.to_string(),
302            kind: MetricKind::Counter,
303        });
304        counter
305    }
306
307    /// 注册 Gauge
308    pub fn register_gauge(&self, name: &str, help: &str) -> Arc<Gauge> {
309        self.register_gauge_with_labels(name, help, HashMap::new())
310    }
311
312    /// 注册带标签的 Gauge
313    pub fn register_gauge_with_labels(
314        &self,
315        name: &str,
316        help: &str,
317        labels: HashMap<String, String>,
318    ) -> Arc<Gauge> {
319        let mut gauges = self.gauges.write();
320        let key = format!("{}_{:?}", name, labels);
321        if let Some(g) = gauges.get(&key) {
322            return g.clone();
323        }
324        let gauge = Arc::new(Gauge {
325            name: name.to_string(),
326            value: Arc::new(RwLock::new(0.0)),
327            labels,
328        });
329        gauges.insert(key, gauge.clone());
330
331        let mut metas = self.metas.write();
332        metas.push(MetricMeta {
333            name: name.to_string(),
334            help: help.to_string(),
335            kind: MetricKind::Gauge,
336        });
337        gauge
338    }
339
340    /// 注册 Histogram
341    pub fn register_histogram(&self, name: &str, help: &str, buckets: Vec<f64>) -> Arc<Histogram> {
342        let mut histograms = self.histograms.write();
343        if let Some(h) = histograms.get(name) {
344            return h.clone();
345        }
346        // 最后一个 bucket 必须是 +Inf
347        let mut all_buckets = buckets;
348        all_buckets.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
349        if !all_buckets.contains(&f64::INFINITY) {
350            all_buckets.push(f64::INFINITY);
351        }
352        let count = all_buckets.len();
353        let histogram = Arc::new(Histogram {
354            name: name.to_string(),
355            buckets: all_buckets,
356            counts: Arc::new(RwLock::new(vec![0; count])),
357            sum: Arc::new(RwLock::new(0.0)),
358            count: Arc::new(RwLock::new(0)),
359        });
360        histograms.insert(name.to_string(), histogram.clone());
361
362        let mut metas = self.metas.write();
363        metas.push(MetricMeta {
364            name: name.to_string(),
365            help: help.to_string(),
366            kind: MetricKind::Histogram,
367        });
368        histogram
369    }
370
371    /// 渲染所有指标为 Prometheus 文本格式
372    pub fn render(&self) -> String {
373        let mut output = String::new();
374
375        // 输出 HELP/TYPE 头
376        let metas = self.metas.read();
377        let mut seen = std::collections::HashSet::new();
378        for meta in metas.iter() {
379            if seen.contains(&meta.name) {
380                continue;
381            }
382            seen.insert(meta.name.clone());
383            output.push_str(&format!("# HELP {} {}\n", meta.name, meta.help));
384            let type_str = match meta.kind {
385                MetricKind::Counter => "counter",
386                MetricKind::Gauge => "gauge",
387                MetricKind::Histogram => "histogram",
388            };
389            output.push_str(&format!("# TYPE {} {}\n", meta.name, type_str));
390        }
391
392        // 输出 Counter 值
393        let counters = self.counters.read();
394        for c in counters.values() {
395            output.push_str(&c.render());
396        }
397
398        // 输出 Gauge 值
399        let gauges = self.gauges.read();
400        for g in gauges.values() {
401            output.push_str(&g.render());
402        }
403
404        // 输出 Histogram 值
405        let histograms = self.histograms.read();
406        for h in histograms.values() {
407            output.push_str(&h.render());
408        }
409
410        output
411    }
412}
413
414/// 启动 Prometheus metrics HTTP server
415///
416/// 在指定地址暴露 `/metrics` 端点，返回 Prometheus 文本格式的指标数据。
417/// 每个连接在独立 tokio task 中处理。
418pub async fn start_metrics_server(
419    registry: Arc<MetricsRegistry>,
420    addr: std::net::SocketAddr,
421) -> Result<(), std::io::Error> {
422    use tokio::io::AsyncWriteExt;
423
424    let listener = tokio::net::TcpListener::bind(addr).await?;
425    loop {
426        let (mut stream, _) = listener.accept().await?;
427        let registry = registry.clone();
428        tokio::spawn(async move {
429            let metrics = registry.render();
430            let response = format!(
431                "HTTP/1.1 200 OK\r\nContent-Type: text/plain; version=0.0.4\r\nContent-Length: {}\r\n\r\n{}",
432                metrics.len(),
433                metrics
434            );
435            let _ = stream.write_all(response.as_bytes()).await;
436        });
437    }
438}
439
440#[cfg(test)]
441mod tests {
442    use super::*;
443
444    #[test]
445    fn test_counter_basic() {
446        let registry = MetricsRegistry::new();
447        let counter = registry.register_counter("test_counter", "Test counter");
448        counter.inc();
449        counter.inc_by(2.5);
450        assert_eq!(counter.value(), 3.5);
451    }
452
453    #[test]
454    fn test_gauge_basic() {
455        let registry = MetricsRegistry::new();
456        let gauge = registry.register_gauge("test_gauge", "Test gauge");
457        gauge.set(10.0);
458        gauge.inc();
459        gauge.dec_by(3.0);
460        assert_eq!(gauge.value(), 8.0);
461    }
462
463    #[test]
464    fn test_histogram_basic() {
465        let registry = MetricsRegistry::new();
466        let histogram =
467            registry.register_histogram("test_histogram", "Test histogram", vec![0.1, 0.5, 1.0]);
468        histogram.observe(0.05);
469        histogram.observe(0.2);
470        histogram.observe(0.6);
471        histogram.observe(1.5);
472
473        assert_eq!(histogram.count(), 4);
474        assert!((histogram.sum() - 2.35).abs() < 1e-9);
475    }
476
477    #[test]
478    fn test_render_prometheus_format() {
479        let registry = MetricsRegistry::new();
480        let counter = registry.register_counter("ops_total", "Total operations");
481        let gauge = registry.register_gauge("conn_active", "Active connections");
482        let histogram =
483            registry.register_histogram("latency_seconds", "Latency in seconds", vec![0.01, 0.1]);
484
485        counter.inc_by(10.0);
486        gauge.set(5.0);
487        histogram.observe(0.005);
488        histogram.observe(0.05);
489        histogram.observe(0.5);
490
491        let output = registry.render();
492        assert!(output.contains("# HELP ops_total Total operations"));
493        assert!(output.contains("# TYPE ops_total counter"));
494        assert!(output.contains("ops_total 10"));
495        assert!(output.contains("conn_active 5"));
496        assert!(output.contains("latency_seconds_bucket{le=\"0.01\"} 1"));
497        assert!(output.contains("latency_seconds_bucket{le=\"0.1\"} 2"));
498        assert!(output.contains("latency_seconds_sum 0.555"));
499        assert!(output.contains("latency_seconds_count 3"));
500    }
501
502    #[test]
503    fn test_counter_with_labels() {
504        let registry = MetricsRegistry::new();
505        let mut labels = HashMap::new();
506        labels.insert("method".to_string(), "GET".to_string());
507        labels.insert("status".to_string(), "200".to_string());
508
509        let counter =
510            registry.register_counter_with_labels("http_requests_total", "HTTP requests", labels);
511        counter.inc();
512        let output = registry.render();
513        // HashMap 顺序未定义，分别验证各标签
514        assert!(output.contains("http_requests_total{"));
515        assert!(output.contains("method=\"GET\""));
516        assert!(output.contains("status=\"200\""));
517        assert!(output.contains("} 1"));
518    }
519}
sz_orm_observability/lib.rs

sz_orm_observability/
lib.rs
