rialo-telemetry 0.2.0-alpha.2

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Prometheus-specific functionality for metrics collection.

use anyhow::Result;
use telemetry_subscribers::span_latency_prom::PrometheusSpanLatencyLayer;

/// Default number of buckets for span latency histograms.
///
/// This value provides a good balance between granularity and memory usage
/// for latency distribution measurements in Prometheus metrics.
pub const DEFAULT_SPAN_LATENCY_BUCKETS: usize = 15;

/// Prometheus-specific configuration
#[derive(Debug, Clone)]
pub struct PrometheusConfig {
    /// Prometheus registry for metrics export
    pub registry: ::prometheus::Registry,
    /// Number of buckets for span latency histograms (default: [`DEFAULT_SPAN_LATENCY_BUCKETS`])
    pub span_latency_buckets: usize,
    /// Whether to enable span latency metrics (default: true)
    pub enable_span_latency: bool,
}

impl PrometheusConfig {
    /// Create a new PrometheusConfig with default values.
    pub fn new(registry: ::prometheus::Registry) -> Self {
        Self {
            registry,
            span_latency_buckets: DEFAULT_SPAN_LATENCY_BUCKETS,
            enable_span_latency: true,
        }
    }

    /// Set the number of buckets for span latency histograms.
    pub fn with_span_latency_buckets(mut self, buckets: usize) -> Self {
        self.span_latency_buckets = buckets;
        self
    }

    /// Enable or disable span latency metrics.
    pub fn with_span_latency_enabled(mut self, enabled: bool) -> Self {
        self.enable_span_latency = enabled;
        self
    }
}

/// Initialize Prometheus metrics support.
///
/// This function validates the Prometheus registry and optionally creates
/// the PrometheusSpanLatencyLayer for tracking span latencies.
pub fn init_prometheus(
    prometheus_config: &PrometheusConfig,
) -> Result<Option<PrometheusSpanLatencyLayer>> {
    validate_prometheus_registry(&prometheus_config.registry);

    if prometheus_config.enable_span_latency {
        let span_lat_layer = PrometheusSpanLatencyLayer::try_new(
            &prometheus_config.registry,
            prometheus_config.span_latency_buckets,
        )
        .map_err(|e| anyhow::anyhow!("Could not initialize span latency layer: {:?}", e))?;
        Ok(Some(span_lat_layer))
    } else {
        Ok(None)
    }
}

/// Helper function to validate Prometheus registry (just stores it for future use)
fn validate_prometheus_registry(_registry: &::prometheus::Registry) -> bool {
    // For now, just validate that it's not null - the existing Prometheus
    // infrastructure in the node will continue to work as before
    true
}