#![warn(missing_docs)]
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
pub mod slo;
pub use slo::{SloBurnRate, SloConfig, SloMonitor};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MetricKind {
Counter,
Gauge,
Histogram,
}
#[derive(Debug, Clone)]
pub struct MetricMeta {
pub name: String,
pub help: String,
pub kind: MetricKind,
}
pub struct Counter {
name: String,
value: Arc<RwLock<f64>>,
labels: HashMap<String, String>,
}
impl Counter {
pub fn inc(&self) {
self.inc_by(1.0);
}
pub fn inc_by(&self, delta: f64) {
let mut v = self.value.write();
*v += delta;
}
pub fn value(&self) -> f64 {
*self.value.read()
}
pub fn name(&self) -> &str {
&self.name
}
pub fn render(&self) -> String {
let v = self.value.read();
if self.labels.is_empty() {
format!("{} {}\n", self.name, v)
} else {
let labels: Vec<String> = self
.labels
.iter()
.map(|(k, val)| format!("{}=\"{}\"", k, val.replace('"', "\\\"")))
.collect();
format!("{}{{{}}} {}\n", self.name, labels.join(","), v)
}
}
}
pub struct Gauge {
name: String,
value: Arc<RwLock<f64>>,
labels: HashMap<String, String>,
}
impl Gauge {
pub fn set(&self, value: f64) {
*self.value.write() = value;
}
pub fn inc(&self) {
self.inc_by(1.0);
}
pub fn inc_by(&self, delta: f64) {
let mut v = self.value.write();
*v += delta;
}
pub fn dec_by(&self, delta: f64) {
let mut v = self.value.write();
*v -= delta;
}
pub fn value(&self) -> f64 {
*self.value.read()
}
pub fn name(&self) -> &str {
&self.name
}
pub fn render(&self) -> String {
let v = self.value.read();
if self.labels.is_empty() {
format!("{} {}\n", self.name, v)
} else {
let labels: Vec<String> = self
.labels
.iter()
.map(|(k, val)| format!("{}=\"{}\"", k, val.replace('"', "\\\"")))
.collect();
format!("{}{{{}}} {}\n", self.name, labels.join(","), v)
}
}
}
pub struct Histogram {
name: String,
buckets: Vec<f64>,
counts: Arc<RwLock<Vec<u64>>>,
sum: Arc<RwLock<f64>>,
count: Arc<RwLock<u64>>,
}
impl Histogram {
pub fn observe(&self, value: f64) {
let mut counts = self.counts.write();
for (i, bucket) in self.buckets.iter().enumerate() {
if value <= *bucket {
counts[i] += 1;
}
}
let last = counts.len() - 1;
counts[last] += 1;
let mut sum = self.sum.write();
*sum += value;
let mut count = self.count.write();
*count += 1;
}
pub fn count(&self) -> u64 {
*self.count.read()
}
pub fn sum(&self) -> f64 {
*self.sum.read()
}
pub fn name(&self) -> &str {
&self.name
}
pub fn render(&self) -> String {
let counts = self.counts.read();
let sum = self.sum.read();
let count = self.count.read();
let mut output = String::new();
for (i, bucket) in self.buckets.iter().enumerate() {
output.push_str(&format!(
"{}_bucket{{le=\"{}\"}} {}\n",
self.name, bucket, counts[i]
));
}
output.push_str(&format!("{}_sum {}\n", self.name, sum));
output.push_str(&format!("{}_count {}\n", self.name, count));
output
}
}
pub struct MetricsRegistry {
counters: RwLock<HashMap<String, Arc<Counter>>>,
gauges: RwLock<HashMap<String, Arc<Gauge>>>,
histograms: RwLock<HashMap<String, Arc<Histogram>>>,
metas: RwLock<Vec<MetricMeta>>,
}
impl Default for MetricsRegistry {
fn default() -> Self {
Self::new()
}
}
impl MetricsRegistry {
pub fn new() -> Self {
Self {
counters: RwLock::new(HashMap::new()),
gauges: RwLock::new(HashMap::new()),
histograms: RwLock::new(HashMap::new()),
metas: RwLock::new(Vec::new()),
}
}
pub fn register_counter(&self, name: &str, help: &str) -> Arc<Counter> {
self.register_counter_with_labels(name, help, HashMap::new())
}
pub fn register_counter_with_labels(
&self,
name: &str,
help: &str,
labels: HashMap<String, String>,
) -> Arc<Counter> {
let mut counters = self.counters.write();
let key = format!("{}_{:?}", name, labels);
if let Some(c) = counters.get(&key) {
return c.clone();
}
let counter = Arc::new(Counter {
name: name.to_string(),
value: Arc::new(RwLock::new(0.0)),
labels,
});
counters.insert(key, counter.clone());
let mut metas = self.metas.write();
metas.push(MetricMeta {
name: name.to_string(),
help: help.to_string(),
kind: MetricKind::Counter,
});
counter
}
pub fn register_gauge(&self, name: &str, help: &str) -> Arc<Gauge> {
self.register_gauge_with_labels(name, help, HashMap::new())
}
pub fn register_gauge_with_labels(
&self,
name: &str,
help: &str,
labels: HashMap<String, String>,
) -> Arc<Gauge> {
let mut gauges = self.gauges.write();
let key = format!("{}_{:?}", name, labels);
if let Some(g) = gauges.get(&key) {
return g.clone();
}
let gauge = Arc::new(Gauge {
name: name.to_string(),
value: Arc::new(RwLock::new(0.0)),
labels,
});
gauges.insert(key, gauge.clone());
let mut metas = self.metas.write();
metas.push(MetricMeta {
name: name.to_string(),
help: help.to_string(),
kind: MetricKind::Gauge,
});
gauge
}
pub fn register_histogram(&self, name: &str, help: &str, buckets: Vec<f64>) -> Arc<Histogram> {
let mut histograms = self.histograms.write();
if let Some(h) = histograms.get(name) {
return h.clone();
}
let mut all_buckets = buckets;
all_buckets.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
if !all_buckets.contains(&f64::INFINITY) {
all_buckets.push(f64::INFINITY);
}
let count = all_buckets.len();
let histogram = Arc::new(Histogram {
name: name.to_string(),
buckets: all_buckets,
counts: Arc::new(RwLock::new(vec![0; count])),
sum: Arc::new(RwLock::new(0.0)),
count: Arc::new(RwLock::new(0)),
});
histograms.insert(name.to_string(), histogram.clone());
let mut metas = self.metas.write();
metas.push(MetricMeta {
name: name.to_string(),
help: help.to_string(),
kind: MetricKind::Histogram,
});
histogram
}
pub fn render(&self) -> String {
let mut output = String::new();
let metas = self.metas.read();
let mut seen = std::collections::HashSet::new();
for meta in metas.iter() {
if seen.contains(&meta.name) {
continue;
}
seen.insert(meta.name.clone());
output.push_str(&format!("# HELP {} {}\n", meta.name, meta.help));
let type_str = match meta.kind {
MetricKind::Counter => "counter",
MetricKind::Gauge => "gauge",
MetricKind::Histogram => "histogram",
};
output.push_str(&format!("# TYPE {} {}\n", meta.name, type_str));
}
let counters = self.counters.read();
for c in counters.values() {
output.push_str(&c.render());
}
let gauges = self.gauges.read();
for g in gauges.values() {
output.push_str(&g.render());
}
let histograms = self.histograms.read();
for h in histograms.values() {
output.push_str(&h.render());
}
output
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_counter_basic() {
let registry = MetricsRegistry::new();
let counter = registry.register_counter("test_counter", "Test counter");
counter.inc();
counter.inc_by(2.5);
assert_eq!(counter.value(), 3.5);
}
#[test]
fn test_gauge_basic() {
let registry = MetricsRegistry::new();
let gauge = registry.register_gauge("test_gauge", "Test gauge");
gauge.set(10.0);
gauge.inc();
gauge.dec_by(3.0);
assert_eq!(gauge.value(), 8.0);
}
#[test]
fn test_histogram_basic() {
let registry = MetricsRegistry::new();
let histogram =
registry.register_histogram("test_histogram", "Test histogram", vec![0.1, 0.5, 1.0]);
histogram.observe(0.05);
histogram.observe(0.2);
histogram.observe(0.6);
histogram.observe(1.5);
assert_eq!(histogram.count(), 4);
assert!((histogram.sum() - 2.35).abs() < 1e-9);
}
#[test]
fn test_render_prometheus_format() {
let registry = MetricsRegistry::new();
let counter = registry.register_counter("ops_total", "Total operations");
let gauge = registry.register_gauge("conn_active", "Active connections");
let histogram =
registry.register_histogram("latency_seconds", "Latency in seconds", vec![0.01, 0.1]);
counter.inc_by(10.0);
gauge.set(5.0);
histogram.observe(0.005);
histogram.observe(0.05);
histogram.observe(0.5);
let output = registry.render();
assert!(output.contains("# HELP ops_total Total operations"));
assert!(output.contains("# TYPE ops_total counter"));
assert!(output.contains("ops_total 10"));
assert!(output.contains("conn_active 5"));
assert!(output.contains("latency_seconds_bucket{le=\"0.01\"} 1"));
assert!(output.contains("latency_seconds_bucket{le=\"0.1\"} 2"));
assert!(output.contains("latency_seconds_sum 0.555"));
assert!(output.contains("latency_seconds_count 3"));
}
#[test]
fn test_counter_with_labels() {
let registry = MetricsRegistry::new();
let mut labels = HashMap::new();
labels.insert("method".to_string(), "GET".to_string());
labels.insert("status".to_string(), "200".to_string());
let counter =
registry.register_counter_with_labels("http_requests_total", "HTTP requests", labels);
counter.inc();
let output = registry.render();
assert!(output.contains("http_requests_total{"));
assert!(output.contains("method=\"GET\""));
assert!(output.contains("status=\"200\""));
assert!(output.contains("} 1"));
}
}