serverless_fn/

telemetry.rs

//! Telemetry and monitoring module.
//!
//! Provides metrics collection, logging, and distributed tracing capabilities.

use std::sync::OnceLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant;

use crate::error::ServerlessError;

/// Metrics collector for serverless function invocations.
pub struct MetricsCollector {
    invocation_count: AtomicU64,
    total_duration_ms: AtomicU64,
}

impl MetricsCollector {
    /// Creates a new metrics collector.
    #[must_use]
    pub fn new() -> Self {
        Self {
            invocation_count: AtomicU64::new(0),
            total_duration_ms: AtomicU64::new(0),
        }
    }

    /// Records a function invocation.
    ///
    /// # Arguments
    ///
    /// * `function_name` - Name of the function being invoked
    /// * `duration_ms` - Duration of the invocation in milliseconds
    /// * `success` - Whether the invocation was successful
    pub fn record_invocation(&self, function_name: &str, duration_ms: u64, success: bool) {
        if cfg!(feature = "telemetry") {
            self.invocation_count.fetch_add(1, Ordering::SeqCst);
            self.total_duration_ms
                .fetch_add(duration_ms, Ordering::SeqCst);

            tracing::info!(
                "Function invocation: {} - Success: {}, Duration: {}ms",
                function_name,
                success,
                duration_ms
            );
        }
    }

    /// Returns the average invocation duration in milliseconds.
    #[must_use]
    pub fn avg_duration_ms(&self) -> f64 {
        let count = self.invocation_count.load(Ordering::SeqCst);
        if count == 0 {
            return 0.0;
        }
        let total = self.total_duration_ms.load(Ordering::SeqCst);
        total as f64 / count as f64
    }

    /// Returns the total invocation count.
    #[must_use]
    pub fn invocation_count(&self) -> u64 {
        self.invocation_count.load(Ordering::SeqCst)
    }
}

impl Default for MetricsCollector {
    fn default() -> Self {
        Self::new()
    }
}

/// Telemetry context for tracking function execution.
pub struct TelemetryContext {
    start_time: Instant,
    function_name: String,
    trace_id: String,
}

impl TelemetryContext {
    /// Creates a new telemetry context.
    #[must_use]
    pub fn new(function_name: &str) -> Self {
        Self {
            start_time: Instant::now(),
            function_name: function_name.to_string(),
            trace_id: uuid::Uuid::new_v4().to_string(),
        }
    }

    /// Records the completion of a function execution.
    ///
    /// # Arguments
    ///
    /// * `success` - Whether the execution was successful
    /// * `error` - Optional error information if the execution failed
    pub fn record_completion(&self, success: bool, error: Option<&ServerlessError>) {
        if cfg!(feature = "telemetry") {
            let duration_ms = self.start_time.elapsed().as_millis() as u64;

            if let Some(err) = error {
                tracing::error!(
                    trace_id = %self.trace_id,
                    function = %self.function_name,
                    duration_ms = duration_ms,
                    error = %err,
                    "Function execution failed"
                );
            } else {
                tracing::info!(
                    trace_id = %self.trace_id,
                    function = %self.function_name,
                    duration_ms = duration_ms,
                    "Function execution completed"
                );
            }

            get_metrics_collector().record_invocation(&self.function_name, duration_ms, success);
        }
    }
}

/// Global metrics collector instance.
static METRICS_COLLECTOR: OnceLock<MetricsCollector> = OnceLock::new();

/// Returns a reference to the global metrics collector.
#[must_use]
pub fn get_metrics_collector() -> &'static MetricsCollector {
    METRICS_COLLECTOR.get_or_init(MetricsCollector::new)
}

/// Initializes the telemetry system.
pub fn init_telemetry() {
    if cfg!(feature = "telemetry") {
        tracing::info!("Telemetry system initialized");
    }
}

/// Middleware for telemetry tracking.
///
/// # Errors
///
/// Returns an error if the operation fails.
pub async fn telemetry_middleware<F, R>(
    function_name: &str,
    operation: F,
) -> Result<R, ServerlessError>
where
    F: std::future::Future<Output = Result<R, ServerlessError>>,
{
    if cfg!(feature = "telemetry") {
        let ctx = TelemetryContext::new(function_name);
        match operation.await {
            Ok(result) => {
                ctx.record_completion(true, None);
                Ok(result)
            }
            Err(err) => {
                ctx.record_completion(false, Some(&err));
                Err(err)
            }
        }
    } else {
        operation.await
    }
}

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

    #[test]
    fn test_metrics_collector_initial_state() {
        let collector = MetricsCollector::new();
        assert_eq!(collector.invocation_count(), 0);
        assert_eq!(collector.avg_duration_ms(), 0.0);
    }

    #[test]
    fn test_metrics_collector_records_invocation() {
        let collector = MetricsCollector::new();

        // Note: record_invocation only updates metrics when telemetry feature is enabled
        collector.record_invocation("test_func", 100, true);
        collector.record_invocation("test_func", 200, true);

        let count = collector.invocation_count();
        let avg_duration = collector.avg_duration_ms();

        if cfg!(feature = "telemetry") {
            assert_eq!(count, 2);
            assert_eq!(avg_duration, 150.0);
        } else {
            assert_eq!(count, 0);
            assert_eq!(avg_duration, 0.0);
        }
    }
}