1#![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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69pub enum MetricKind {
70 Counter,
72 Gauge,
74 Histogram,
76}
77
78#[derive(Debug, Clone)]
80pub struct MetricMeta {
81 pub name: String,
83 pub help: String,
85 pub kind: MetricKind,
87}
88
89pub struct Counter {
91 name: String,
92 value: Arc<RwLock<f64>>,
93 labels: HashMap<String, String>,
94}
95
96impl Counter {
97 pub fn inc(&self) {
99 self.inc_by(1.0);
100 }
101
102 pub fn inc_by(&self, delta: f64) {
104 let mut v = self.value.write();
105 *v += delta;
106 }
107
108 pub fn value(&self) -> f64 {
110 *self.value.read()
111 }
112
113 pub fn name(&self) -> &str {
115 &self.name
116 }
117
118 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
134pub struct Gauge {
136 name: String,
137 value: Arc<RwLock<f64>>,
138 labels: HashMap<String, String>,
139}
140
141impl Gauge {
142 pub fn set(&self, value: f64) {
144 *self.value.write() = value;
145 }
146
147 pub fn inc(&self) {
149 self.inc_by(1.0);
150 }
151
152 pub fn inc_by(&self, delta: f64) {
154 let mut v = self.value.write();
155 *v += delta;
156 }
157
158 pub fn dec_by(&self, delta: f64) {
160 let mut v = self.value.write();
161 *v -= delta;
162 }
163
164 pub fn value(&self) -> f64 {
166 *self.value.read()
167 }
168
169 pub fn name(&self) -> &str {
171 &self.name
172 }
173
174 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
190pub 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 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 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 pub fn count(&self) -> u64 {
220 *self.count.read()
221 }
222
223 pub fn sum(&self) -> f64 {
225 *self.sum.read()
226 }
227
228 pub fn name(&self) -> &str {
230 &self.name
231 }
232
233 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
252pub 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 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 pub fn register_counter(&self, name: &str, help: &str) -> Arc<Counter> {
279 self.register_counter_with_labels(name, help, HashMap::new())
280 }
281
282 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 pub fn register_gauge(&self, name: &str, help: &str) -> Arc<Gauge> {
312 self.register_gauge_with_labels(name, help, HashMap::new())
313 }
314
315 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 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 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 pub fn render(&self) -> String {
376 let mut output = String::new();
377
378 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 let counters = self.counters.read();
397 for c in counters.values() {
398 output.push_str(&c.render());
399 }
400
401 let gauges = self.gauges.read();
403 for g in gauges.values() {
404 output.push_str(&g.render());
405 }
406
407 let histograms = self.histograms.read();
409 for h in histograms.values() {
410 output.push_str(&h.render());
411 }
412
413 output
414 }
415}
416
417pub 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 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}