uvb_tracing/

lib.rs

/*!
 * UVB Tracing - OpenTelemetry Distributed Tracing
 *
 * Provides distributed tracing instrumentation for UVB authentication system
 * using OpenTelemetry and the tracing ecosystem.
 */

use opentelemetry::{global, KeyValue};
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::{
    propagation::TraceContextPropagator,
    runtime,
    trace::{self},
    Resource,
};
use opentelemetry_semantic_conventions as semconv;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Registry};

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

    /// Service version
    pub service_version: String,

    /// Environment (production, staging, development)
    pub environment: String,

    /// OTLP endpoint (e.g., "http://localhost:4317")
    pub otlp_endpoint: String,

    /// Enable JSON logging
    pub json_logging: bool,

    /// Log level filter (e.g., "info,uvb=debug")
    pub log_level: String,
}

impl Default for TracingConfig {
    fn default() -> Self {
        Self {
            service_name: "uvb-api".to_string(),
            service_version: env!("CARGO_PKG_VERSION").to_string(),
            environment: "development".to_string(),
            otlp_endpoint: "http://localhost:4317".to_string(),
            json_logging: false,
            log_level: "info".to_string(),
        }
    }
}

/// Initialize OpenTelemetry tracing with OTLP exporter
///
/// # Example
///
/// ```no_run
/// use uvb_tracing::{TracingConfig, init_tracing};
///
/// #[tokio::main]
/// async fn main() {
///     let config = TracingConfig {
///         service_name: "uvb-api".to_string(),
///         otlp_endpoint: "http://jaeger:4317".to_string(),
///         ..Default::default()
///     };
///
///     init_tracing(config).expect("Failed to initialize tracing");
/// }
/// ```
pub fn init_tracing(config: TracingConfig) -> Result<(), Box<dyn std::error::Error>> {
    // Set up trace context propagation
    global::set_text_map_propagator(TraceContextPropagator::new());

    // Create OTLP trace exporter
    let tracer = opentelemetry_otlp::new_pipeline()
        .tracing()
        .with_exporter(
            opentelemetry_otlp::new_exporter()
                .tonic()
                .with_endpoint(&config.otlp_endpoint),
        )
        .with_trace_config(trace::config().with_resource(Resource::new(vec![
            KeyValue::new(semconv::resource::SERVICE_NAME, config.service_name.clone()),
            KeyValue::new(semconv::resource::SERVICE_VERSION, config.service_version),
            KeyValue::new(
                semconv::resource::DEPLOYMENT_ENVIRONMENT,
                config.environment,
            ),
        ])))
        .install_batch(runtime::Tokio)?;

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

    // Create env filter layer
    let env_filter =
        EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(&config.log_level));

    // Create subscriber
    let subscriber = Registry::default().with(env_filter).with(telemetry_layer);

    // Add JSON formatting if enabled
    if config.json_logging {
        let fmt_layer = tracing_subscriber::fmt::layer()
            .json()
            .with_current_span(true)
            .with_span_list(true);

        subscriber.with(fmt_layer).try_init()?;
    } else {
        let fmt_layer = tracing_subscriber::fmt::layer()
            .with_target(true)
            .with_level(true)
            .with_thread_ids(true);

        subscriber.with(fmt_layer).try_init()?;
    }

    Ok(())
}

/// Shutdown OpenTelemetry gracefully
///
/// Call this before application exit to ensure all spans are flushed
pub fn shutdown_tracing() {
    global::shutdown_tracer_provider();
}

/// Span attribute constants for UVB-specific tracing
pub mod attributes {
    /// Tenant ID attribute
    pub const TENANT_ID: &str = "uvb.tenant_id";

    /// User ID attribute
    pub const USER_ID: &str = "uvb.user_id";

    /// Application ID attribute
    pub const APPLICATION_ID: &str = "uvb.application_id";

    /// Transaction ID attribute
    pub const TRANSACTION_ID: &str = "uvb.transaction_id";

    /// Factor ID attribute
    pub const FACTOR_ID: &str = "uvb.factor_id";

    /// Challenge ID attribute
    pub const CHALLENGE_ID: &str = "uvb.challenge_id";

    /// Enrollment ID attribute
    pub const ENROLLMENT_ID: &str = "uvb.enrollment_id";

    /// Session ID attribute
    pub const SESSION_ID: &str = "uvb.session_id";

    /// Intent attribute (login, payment, etc.)
    pub const INTENT: &str = "uvb.intent";

    /// Verification status attribute
    pub const VERIFICATION_STATUS: &str = "uvb.verification_status";

    /// Assurance level attribute
    pub const ASSURANCE_LEVEL: &str = "uvb.assurance_level";

    /// Policy ID attribute
    pub const POLICY_ID: &str = "uvb.policy_id";
}

/// Helper macros for common tracing patterns
#[macro_export]
macro_rules! span_with_tenant {
    ($name:expr, $tenant_id:expr) => {
        tracing::info_span!(
            $name,
            { $crate::attributes::TENANT_ID } = %$tenant_id
        )
    };
}

#[macro_export]
macro_rules! span_with_transaction {
    ($name:expr, $transaction_id:expr, $tenant_id:expr) => {
        tracing::info_span!(
            $name,
            { $crate::attributes::TRANSACTION_ID } = %$transaction_id,
            { $crate::attributes::TENANT_ID } = %$tenant_id
        )
    };
}

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

    #[test]
    fn test_default_config() {
        let config = TracingConfig::default();
        assert_eq!(config.service_name, "uvb-api");
        assert_eq!(config.environment, "development");
    }

    #[test]
    fn test_config_builder() {
        let config = TracingConfig {
            service_name: "test-service".to_string(),
            environment: "test".to_string(),
            ..Default::default()
        };

        assert_eq!(config.service_name, "test-service");
        assert_eq!(config.environment, "test");
    }
}