Skip to main content

sz_orm_observability/
lib.rs

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