Vapora/crates/vapora-telemetry/src/tracer.rs

use opentelemetry::global;
use opentelemetry_jaeger::new_agent_pipeline;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use tracing_subscriber::{EnvFilter, Registry};

use crate::error::{Result, TelemetryError};

/// Configuration for telemetry initialization
#[derive(Debug, Clone)]
pub struct TelemetryConfig {
    /// Service name for tracing
    pub service_name: String,

    /// Jaeger agent host
    pub jaeger_host: String,

    /// Jaeger agent port (default 6831)
    pub jaeger_port: u16,

    /// Log level filter
    pub log_level: String,

    /// Enable console output
    pub console_output: bool,

    /// Enable JSON output
    pub json_output: bool,
}

impl Default for TelemetryConfig {
    fn default() -> Self {
        Self {
            service_name: "vapora".to_string(),
            jaeger_host: "localhost".to_string(),
            jaeger_port: 6831,
            log_level: "info".to_string(),
            console_output: true,
            json_output: false,
        }
    }
}

/// Telemetry initializer - sets up OpenTelemetry with Jaeger exporter
pub struct TelemetryInitializer;

impl TelemetryInitializer {
    /// Initialize tracing with OpenTelemetry and Jaeger exporter
    pub fn init(config: TelemetryConfig) -> Result<()> {
        // Create Jaeger exporter
        let tracer = new_agent_pipeline()
            .with_service_name(&config.service_name)
            .with_endpoint(format!("{}:{}", config.jaeger_host, config.jaeger_port))
            .install_simple()
            .map_err(|e| TelemetryError::JaegerError(e.to_string()))?;

        // Create OpenTelemetry layer for tracing
        let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer);

        // Create environment filter from config
        let env_filter = EnvFilter::try_from_default_env()
            .or_else(|_| EnvFilter::try_new(&config.log_level))
            .map_err(|e| TelemetryError::TracerInitFailed(e.to_string()))?;

        // Build subscriber with OpenTelemetry layer
        let registry = Registry::default().with(env_filter).with(otel_layer);

        if config.console_output {
            if config.json_output {
                registry
                    .with(tracing_subscriber::fmt::layer().json())
                    .init();
            } else {
                registry.with(tracing_subscriber::fmt::layer()).init();
            }
        } else {
            registry.init();
        }

        tracing::info!(
            service = %config.service_name,
            jaeger_endpoint = %format!("{}:{}", config.jaeger_host, config.jaeger_port),
            "Telemetry initialized successfully"
        );

        Ok(())
    }

    /// Initialize minimal tracing for testing (no Jaeger)
    pub fn init_noop() -> Result<()> {
        let env_filter = EnvFilter::try_from_default_env()
            .or_else(|_| EnvFilter::try_new("info"))
            .map_err(|e| TelemetryError::TracerInitFailed(e.to_string()))?;

        Registry::default()
            .with(env_filter)
            .with(tracing_subscriber::fmt::layer())
            .init();

        Ok(())
    }

    /// Shutdown global tracer (cleanup)
    pub fn shutdown() -> Result<()> {
        global::shutdown_tracer_provider();
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_config_default() {
        let config = TelemetryConfig::default();
        assert_eq!(config.service_name, "vapora");
        assert_eq!(config.jaeger_host, "localhost");
        assert_eq!(config.jaeger_port, 6831);
    }

    #[test]
    fn test_init_noop() {
        let result = TelemetryInitializer::init_noop();
        assert!(result.is_ok());
    }

    #[test]
    fn test_config_custom() {
        let config = TelemetryConfig {
            service_name: "test-service".to_string(),
            jaeger_host: "jaeger.example.com".to_string(),
            jaeger_port: 6832,
            log_level: "debug".to_string(),
            console_output: true,
            json_output: true,
        };

        assert_eq!(config.service_name, "test-service");
        assert_eq!(config.jaeger_host, "jaeger.example.com");
        assert_eq!(config.jaeger_port, 6832);
    }
}