use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant;
#[derive(Debug, Clone)]
pub struct Counter {
value: Arc<AtomicU64>,
name: String,
}
impl Counter {
pub fn new(name: impl Into<String>) -> Self {
Self {
value: Arc::new(AtomicU64::new(0)),
name: name.into(),
}
}
pub fn increment(&self) {
self.add(1);
}
pub fn add(&self, value: u64) {
let mut current = self.value.load(Ordering::Relaxed);
loop {
let new_value = current.saturating_add(value);
match self.value.compare_exchange_weak(
current,
new_value,
Ordering::Release,
Ordering::Relaxed,
) {
Ok(_) => break,
Err(actual) => current = actual,
}
}
}
pub fn value(&self) -> u64 {
self.value.load(Ordering::Acquire)
}
pub fn name(&self) -> &str {
&self.name
}
pub fn reset(&self) {
self.value.store(0, Ordering::Release);
}
}
#[derive(Debug, Clone)]
pub struct Histogram {
values: Arc<RwLock<Vec<f64>>>,
name: String,
}
impl Histogram {
pub fn new(name: impl Into<String>) -> Self {
Self {
values: Arc::new(RwLock::new(Vec::new())),
name: name.into(),
}
}
pub fn record(&self, value: f64) {
let mut v = self.values.write();
v.push(value);
}
pub fn timer(&self) -> HistogramTimer {
HistogramTimer {
histogram: self.clone(),
start: Instant::now(),
}
}
pub fn stats(&self) -> HistogramStats {
let v = self.values.read();
if v.is_empty() {
return HistogramStats {
count: 0,
sum: 0.0,
min: 0.0,
max: 0.0,
mean: 0.0,
};
}
let sum: f64 = v.iter().sum();
let count = v.len();
let mean = sum / count as f64;
let min = v.iter().copied().fold(f64::INFINITY, f64::min);
let max = v.iter().copied().fold(f64::NEG_INFINITY, f64::max);
HistogramStats {
count,
sum,
min,
max,
mean,
}
}
pub fn name(&self) -> &str {
&self.name
}
pub fn reset(&self) {
let mut v = self.values.write();
v.clear();
}
}
#[derive(Debug, Clone)]
pub struct HistogramStats {
pub count: usize,
pub sum: f64,
pub min: f64,
pub max: f64,
pub mean: f64,
}
#[derive(Debug)]
pub struct HistogramTimer {
histogram: Histogram,
start: Instant,
}
impl Drop for HistogramTimer {
fn drop(&mut self) {
let duration_ms = self.start.elapsed().as_secs_f64() * 1000.0;
self.histogram.record(duration_ms);
}
}
#[derive(Debug)]
pub struct MetricsContext {
enabled: bool,
counters: Arc<RwLock<HashMap<String, Counter>>>,
histograms: Arc<RwLock<HashMap<String, Histogram>>>,
}
impl MetricsContext {
pub fn new(enabled: bool) -> Self {
Self {
enabled,
counters: Arc::new(RwLock::new(HashMap::new())),
histograms: Arc::new(RwLock::new(HashMap::new())),
}
}
pub fn counter(&self, name: &str) -> Counter {
if !self.enabled {
return Counter::new(name);
}
let mut counters = self.counters.write();
if let Some(counter) = counters.get(name) {
counter.clone()
} else {
let counter = Counter::new(name);
counters.insert(name.to_string(), counter.clone());
counter
}
}
pub fn histogram(&self, name: &str) -> Histogram {
if !self.enabled {
return Histogram::new(name);
}
let mut histograms = self.histograms.write();
if let Some(histogram) = histograms.get(name) {
histogram.clone()
} else {
let histogram = Histogram::new(name);
histograms.insert(name.to_string(), histogram.clone());
histogram
}
}
pub fn get_counters(&self) -> HashMap<String, u64> {
self.counters
.read()
.iter()
.map(|(k, v)| (k.clone(), v.value()))
.collect()
}
pub fn get_histograms(&self) -> HashMap<String, HistogramStats> {
self.histograms
.read()
.iter()
.map(|(k, v)| (k.clone(), v.stats()))
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_counter_increment() {
let counter = Counter::new("test");
assert_eq!(counter.value(), 0);
counter.increment();
assert_eq!(counter.value(), 1);
counter.add(5);
assert_eq!(counter.value(), 6);
}
#[test]
fn test_counter_saturation() {
let counter = Counter::new("test");
counter.value.store(u64::MAX, Ordering::Release);
counter.add(100); assert_eq!(counter.value(), u64::MAX);
}
#[test]
fn test_histogram_stats() {
let histogram = Histogram::new("test");
histogram.record(1.0);
histogram.record(2.0);
histogram.record(3.0);
let stats = histogram.stats();
assert_eq!(stats.count, 3);
assert_eq!(stats.sum, 6.0);
assert_eq!(stats.min, 1.0);
assert_eq!(stats.max, 3.0);
assert_eq!(stats.mean, 2.0);
}
#[test]
fn test_histogram_timer() {
let histogram = Histogram::new("test");
{
let _timer = histogram.timer();
std::thread::sleep(std::time::Duration::from_millis(10));
}
let stats = histogram.stats();
assert_eq!(stats.count, 1);
assert!(stats.sum > 10.0); }
#[test]
fn test_metrics_context() {
let ctx = MetricsContext::new(true);
let counter = ctx.counter("requests");
counter.add(5);
let counters = ctx.get_counters();
assert_eq!(counters.get("requests"), Some(&5));
}
#[test]
fn test_metrics_context_disabled() {
let ctx = MetricsContext::new(false);
let counter = ctx.counter("requests");
counter.add(100);
let counters = ctx.get_counters();
assert!(counters.is_empty());
}
}