use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
#[derive(Debug, Clone)]
pub struct Counter {
name: String,
help: String,
value: Arc<AtomicU64>,
}
impl Counter {
pub fn new(name: &str, help: &str) -> Self {
Counter {
name: name.to_string(),
help: help.to_string(),
value: Arc::new(AtomicU64::new(0)),
}
}
pub fn inc(&self) {
self.value.fetch_add(1, Ordering::SeqCst);
}
pub fn add(&self, delta: u64) {
self.value.fetch_add(delta, Ordering::SeqCst);
}
pub fn get(&self) -> u64 {
self.value.load(Ordering::SeqCst)
}
pub fn export(&self) -> String {
format!(
"# HELP {} {}\n# TYPE {} counter\n{} {}\n",
self.name,
self.help,
self.name,
self.name,
self.get()
)
}
}
#[derive(Debug, Clone)]
pub struct Gauge {
name: String,
help: String,
value: Arc<AtomicU64>,
}
impl Gauge {
pub fn new(name: &str, help: &str) -> Self {
Gauge {
name: name.to_string(),
help: help.to_string(),
value: Arc::new(AtomicU64::new(0)),
}
}
pub fn set(&self, value: u64) {
self.value.store(value, Ordering::SeqCst);
}
pub fn inc(&self) {
self.value.fetch_add(1, Ordering::SeqCst);
}
pub fn dec(&self) {
self.value.fetch_sub(1, Ordering::SeqCst);
}
pub fn get(&self) -> u64 {
self.value.load(Ordering::SeqCst)
}
pub fn export(&self) -> String {
format!(
"# HELP {} {}\n# TYPE {} gauge\n{} {}\n",
self.name,
self.help,
self.name,
self.name,
self.get()
)
}
}
#[derive(Debug, Clone)]
pub struct MetricsRegistry {
counters: Vec<Counter>,
gauges: Vec<Gauge>,
}
impl MetricsRegistry {
pub fn new() -> Self {
MetricsRegistry {
counters: Vec::new(),
gauges: Vec::new(),
}
}
pub fn register_counter(&mut self, counter: Counter) {
self.counters.push(counter);
}
pub fn register_gauge(&mut self, gauge: Gauge) {
self.gauges.push(gauge);
}
pub fn counter(&mut self, name: &str, help: &str) -> Counter {
let counter = Counter::new(name, help);
self.counters.push(counter.clone());
counter
}
pub fn gauge(&mut self, name: &str, help: &str) -> Gauge {
let gauge = Gauge::new(name, help);
self.gauges.push(gauge.clone());
gauge
}
pub fn export(&self) -> String {
let mut output = String::new();
for counter in &self.counters {
output.push_str(&counter.export());
}
for gauge in &self.gauges {
output.push_str(&gauge.export());
}
output
}
}
impl Default for MetricsRegistry {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct RuntimeMetrics {
pub runs_total: Counter,
pub runs_active: Gauge,
pub runs_completed: Counter,
pub runs_failed: Counter,
pub runs_cancelled: Counter,
pub execution_time_ms: Counter,
}
impl RuntimeMetrics {
pub fn new() -> Self {
RuntimeMetrics {
runs_total: Counter::new("bzzz_runs_total", "Total number of runs started"),
runs_active: Gauge::new("bzzz_runs_active", "Number of currently active runs"),
runs_completed: Counter::new("bzzz_runs_completed", "Number of completed runs"),
runs_failed: Counter::new("bzzz_runs_failed", "Number of failed runs"),
runs_cancelled: Counter::new("bzzz_runs_cancelled", "Number of cancelled runs"),
execution_time_ms: Counter::new(
"bzzz_execution_time_ms",
"Total execution time in milliseconds",
),
}
}
pub fn register(&self, registry: &mut MetricsRegistry) {
registry.register_counter(self.runs_total.clone());
registry.register_gauge(self.runs_active.clone());
registry.register_counter(self.runs_completed.clone());
registry.register_counter(self.runs_failed.clone());
registry.register_counter(self.runs_cancelled.clone());
registry.register_counter(self.execution_time_ms.clone());
}
pub fn run_started(&self) {
self.runs_total.inc();
self.runs_active.inc();
}
pub fn run_completed(&self, duration_ms: u64) {
self.runs_active.dec();
self.runs_completed.inc();
self.execution_time_ms.add(duration_ms);
}
pub fn run_failed(&self) {
self.runs_active.dec();
self.runs_failed.inc();
}
pub fn run_cancelled(&self) {
self.runs_active.dec();
self.runs_cancelled.inc();
}
}
impl Default for RuntimeMetrics {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_counter() {
let counter = Counter::new("test_counter", "A test counter");
assert_eq!(counter.get(), 0);
counter.inc();
assert_eq!(counter.get(), 1);
counter.add(5);
assert_eq!(counter.get(), 6);
}
#[test]
fn test_gauge() {
let gauge = Gauge::new("test_gauge", "A test gauge");
assert_eq!(gauge.get(), 0);
gauge.set(10);
assert_eq!(gauge.get(), 10);
gauge.inc();
assert_eq!(gauge.get(), 11);
gauge.dec();
assert_eq!(gauge.get(), 10);
}
#[test]
fn test_metrics_registry() {
let mut registry = MetricsRegistry::new();
let counter = registry.counter("test_counter", "Test counter");
let gauge = registry.gauge("test_gauge", "Test gauge");
counter.inc();
gauge.set(5);
let output = registry.export();
assert!(output.contains("test_counter 1"));
assert!(output.contains("test_gauge 5"));
}
#[test]
fn test_runtime_metrics() {
let metrics = RuntimeMetrics::new();
metrics.run_started();
assert_eq!(metrics.runs_total.get(), 1);
assert_eq!(metrics.runs_active.get(), 1);
metrics.run_completed(100);
assert_eq!(metrics.runs_active.get(), 0);
assert_eq!(metrics.runs_completed.get(), 1);
assert_eq!(metrics.execution_time_ms.get(), 100);
}
}