bzzz-core 0.1.0

Bzzz core library - Declarative orchestration engine for AI Agents
Documentation
//! Metrics Export Module
//!
//! Provides Prometheus-style metrics for monitoring.

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

/// Counter metric type
#[derive(Debug, Clone)]
pub struct Counter {
    name: String,
    help: String,
    value: Arc<AtomicU64>,
}

impl Counter {
    /// Create a new counter
    pub fn new(name: &str, help: &str) -> Self {
        Counter {
            name: name.to_string(),
            help: help.to_string(),
            value: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Increment the counter by 1
    pub fn inc(&self) {
        self.value.fetch_add(1, Ordering::SeqCst);
    }

    /// Increment the counter by a value
    pub fn add(&self, delta: u64) {
        self.value.fetch_add(delta, Ordering::SeqCst);
    }

    /// Get current value
    pub fn get(&self) -> u64 {
        self.value.load(Ordering::SeqCst)
    }

    /// Export in Prometheus format
    pub fn export(&self) -> String {
        format!(
            "# HELP {} {}\n# TYPE {} counter\n{} {}\n",
            self.name,
            self.help,
            self.name,
            self.name,
            self.get()
        )
    }
}

/// Gauge metric type
#[derive(Debug, Clone)]
pub struct Gauge {
    name: String,
    help: String,
    value: Arc<AtomicU64>,
}

impl Gauge {
    /// Create a new gauge
    pub fn new(name: &str, help: &str) -> Self {
        Gauge {
            name: name.to_string(),
            help: help.to_string(),
            value: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Set the gauge value
    pub fn set(&self, value: u64) {
        self.value.store(value, Ordering::SeqCst);
    }

    /// Increment the gauge by 1
    pub fn inc(&self) {
        self.value.fetch_add(1, Ordering::SeqCst);
    }

    /// Decrement the gauge by 1
    pub fn dec(&self) {
        self.value.fetch_sub(1, Ordering::SeqCst);
    }

    /// Get current value
    pub fn get(&self) -> u64 {
        self.value.load(Ordering::SeqCst)
    }

    /// Export in Prometheus format
    pub fn export(&self) -> String {
        format!(
            "# HELP {} {}\n# TYPE {} gauge\n{} {}\n",
            self.name,
            self.help,
            self.name,
            self.name,
            self.get()
        )
    }
}

/// Metrics registry
#[derive(Debug, Clone)]
pub struct MetricsRegistry {
    counters: Vec<Counter>,
    gauges: Vec<Gauge>,
}

impl MetricsRegistry {
    /// Create a new metrics registry
    pub fn new() -> Self {
        MetricsRegistry {
            counters: Vec::new(),
            gauges: Vec::new(),
        }
    }

    /// Register a counter
    pub fn register_counter(&mut self, counter: Counter) {
        self.counters.push(counter);
    }

    /// Register a gauge
    pub fn register_gauge(&mut self, gauge: Gauge) {
        self.gauges.push(gauge);
    }

    /// Create and register a counter
    pub fn counter(&mut self, name: &str, help: &str) -> Counter {
        let counter = Counter::new(name, help);
        self.counters.push(counter.clone());
        counter
    }

    /// Create and register a gauge
    pub fn gauge(&mut self, name: &str, help: &str) -> Gauge {
        let gauge = Gauge::new(name, help);
        self.gauges.push(gauge.clone());
        gauge
    }

    /// Export all metrics in Prometheus format
    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()
    }
}

/// Bzzz runtime metrics
#[derive(Debug, Clone)]
pub struct RuntimeMetrics {
    /// Total runs started
    pub runs_total: Counter,
    /// Currently running executions
    pub runs_active: Gauge,
    /// Successful completions
    pub runs_completed: Counter,
    /// Failed executions
    pub runs_failed: Counter,
    /// Cancelled executions
    pub runs_cancelled: Counter,
    /// Execution time (ms)
    pub execution_time_ms: Counter,
}

impl RuntimeMetrics {
    /// Create runtime metrics with default names
    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",
            ),
        }
    }

    /// Register all metrics with a registry
    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());
    }

    /// Record a run start
    pub fn run_started(&self) {
        self.runs_total.inc();
        self.runs_active.inc();
    }

    /// Record a run completion
    pub fn run_completed(&self, duration_ms: u64) {
        self.runs_active.dec();
        self.runs_completed.inc();
        self.execution_time_ms.add(duration_ms);
    }

    /// Record a run failure
    pub fn run_failed(&self) {
        self.runs_active.dec();
        self.runs_failed.inc();
    }

    /// Record a run cancellation
    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);
    }
}