foxhunt/trading_engine/src/tracing.rs

//! Ultra-low latency distributed tracing for HFT systems
//!
//! This module provides OpenTelemetry-compatible distributed tracing with minimal
//! performance impact on critical trading paths. It uses lock-free data structures
//! and asynchronous export to maintain sub-microsecond latency requirements.
//!
//! ## Architecture
//!
//! ```text
//! Trading Thread (Critical Path)          Tracing Infrastructure (Async)
//! ┌─────────────────────────────┐        ┌──────────────────────────────────┐
//! │ Order Processing            │        │ SpanProcessor                    │
//! │ ┌─────────────────────────┐ │ -----> │ ┌──────────────────────────────┐ │
//! │ │ start_span_fast()       │ │ <1ns   │ │ Lock-free Span Queue         │ │
//! │ │ (atomic operations)     │ │        │ │ (Crossbeam SPSC)             │ │
//! │ └─────────────────────────┘ │        │ └──────────────────────────────┘ │
//! │ Risk Check                  │        │                                  │
//! │ Market Data Processing      │        │ Background Export Thread         │
//! │ ┌─────────────────────────┐ │        │ ┌──────────────────────────────┐ │
//! │ │ end_span_fast()         │ │ -----> │ │ Jaeger/OTLP Export           │ │
//! │ │ (atomic write)          │ │ <1ns   │ │ (Batched, Compressed)        │ │
//! │ └─────────────────────────┘ │        │ └──────────────────────────────┘ │
//! └─────────────────────────────┘        └──────────────────────────────────┘
//! ```

use anyhow::{anyhow, Result};
use crossbeam_queue::SegQueue;
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use uuid::Uuid;

/// Maximum number of spans to buffer before dropping
const MAX_SPAN_BUFFER_SIZE: usize = 100_000;

/// Span export batch size for efficiency
const SPAN_EXPORT_BATCH_SIZE: usize = 1000;

/// Ultra-lightweight span for critical path operations
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
/// FastSpan
/// 
/// TODO: Add detailed documentation for this struct
pub struct FastSpan {
    /// Unique span ID
    pub span_id: u64,
    /// Parent span ID (0 if root span)
    pub parent_id: u64,
    /// Trace ID for correlation
    pub trace_id: u128,
    /// Operation name
    pub operation_name: String,
    /// Start timestamp in nanoseconds
    pub start_time_ns: u64,
    /// End timestamp in nanoseconds (0 if not finished)
    pub end_time_ns: u64,
    /// Span tags/attributes
    pub tags: Vec<(String, String)>,
    /// Service name
    pub service_name: String,
    /// Span status
    pub status: SpanStatus,
}

/// Span status indicating success/error state
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
/// SpanStatus
/// 
/// TODO: Add detailed documentation for this enum
pub enum SpanStatus {
    // Ok variant
    Ok,
    // Error variant
    Error,
    // Timeout variant
    Timeout,
    // Cancelled variant
    Cancelled,
}

impl Default for SpanStatus {
    fn default() -> Self {
        Self::Ok
    }
}

impl FastSpan {
    /// Create new span with minimal allocation
    pub fn new(operation_name: &str, service_name: &str) -> Self {
        let now_ns = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_nanos() as u64)
            .unwrap_or(0);

        Self {
            span_id: generate_span_id(),
            parent_id: 0,
            trace_id: generate_trace_id(),
            operation_name: operation_name.to_string(),
            start_time_ns: now_ns,
            end_time_ns: 0,
            tags: Vec::new(),
            service_name: service_name.to_string(),
            status: SpanStatus::Ok,
        }
    }

    /// Create child span
    pub fn child(&self, operation_name: &str) -> Self {
        let now_ns = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_nanos() as u64)
            .unwrap_or(0);

        Self {
            span_id: generate_span_id(),
            parent_id: self.span_id,
            trace_id: self.trace_id,
            operation_name: operation_name.to_string(),
            start_time_ns: now_ns,
            end_time_ns: 0,
            tags: Vec::new(),
            service_name: self.service_name.clone(),
            status: SpanStatus::Ok,
        }
    }

    /// Add tag with minimal overhead
    #[inline(always)]
    pub fn set_tag(&mut self, key: &str, value: &str) {
        self.tags.push((key.to_string(), value.to_string()));
    }

    /// Mark span as completed
    #[inline(always)]
    pub fn finish(&mut self) {
        self.end_time_ns = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_nanos() as u64)
            .unwrap_or(0);
    }

    /// Mark span as completed with status
    #[inline(always)]
    pub fn finish_with_status(&mut self, status: SpanStatus) {
        self.status = status;
        self.finish();
    }

    /// Get span duration in nanoseconds
    pub fn duration_ns(&self) -> u64 {
        if self.end_time_ns > 0 && self.end_time_ns >= self.start_time_ns {
            self.end_time_ns - self.start_time_ns
        } else {
            0
        }
    }

    /// Get span duration in microseconds
    pub fn duration_us(&self) -> f64 {
        self.duration_ns() as f64 / 1000.0
    }

    /// Convert to Jaeger-compatible format
    pub fn to_jaeger_span(&self) -> JaegerSpan {
        JaegerSpan {
            trace_id: format!("{:032x}", self.trace_id),
            span_id: format!("{:016x}", self.span_id),
            parent_span_id: if self.parent_id > 0 {
                Some(format!("{:016x}", self.parent_id))
            } else {
    // None variant
                None
            },
            operation_name: self.operation_name.clone(),
            start_time: self.start_time_ns,
            duration: self.duration_ns(),
            tags: self.tags.iter()
                .map(|(k, v)| JaegerTag {
                    key: k.clone(),
                    value: v.clone(),
                    tag_type: "string".to_string(),
                })
                .collect(),
            process: JaegerProcess {
                service_name: self.service_name.clone(),
                tags: vec![],
            },
        }
    }
}

/// Jaeger-compatible span format for export
#[derive(Debug, Serialize)]
/// JaegerSpan
/// 
/// TODO: Add detailed documentation for this struct
pub struct JaegerSpan {
    #[serde(rename = "traceID")]
    /// Trace Id
    pub trace_id: String,
    #[serde(rename = "spanID")]
    /// Span Id
    pub span_id: String,
    #[serde(rename = "parentSpanID")]
    /// Parent Span Id
    pub parent_span_id: Option<String>,
    #[serde(rename = "operationName")]
    /// Operation Name
    pub operation_name: String,
    #[serde(rename = "startTime")]
    /// Start Time
    pub start_time: u64,
    /// Duration
    pub duration: u64,
    /// Tags
    pub tags: Vec<JaegerTag>,
    /// Process
    pub process: JaegerProcess,
}

#[derive(Debug, Serialize)]
/// JaegerTag
/// 
/// TODO: Add detailed documentation for this struct
pub struct JaegerTag {
    /// Key
    pub key: String,
    /// Value
    pub value: String,
    #[serde(rename = "type")]
    /// Tag Type
    pub tag_type: String,
}

#[derive(Debug, Serialize)]
/// JaegerProcess
/// 
/// TODO: Add detailed documentation for this struct
pub struct JaegerProcess {
    #[serde(rename = "serviceName")]
    /// Service Name
    pub service_name: String,
    /// Tags
    pub tags: Vec<JaegerTag>,
}

/// Lock-free tracing infrastructure
#[derive(Debug)]
pub struct FastTracer {
    /// Service name for all spans created by this tracer
    pub service_name: String,
    /// Lock-free queue for finished spans
    span_queue: Arc<SegQueue<FastSpan>>,
    /// Dropped spans counter
    dropped_spans: AtomicU64,
    /// Total spans created
    spans_created: AtomicU64,
    /// Total spans exported
    spans_exported: AtomicU64,
}

impl FastTracer {
    /// Create new tracer for a service
    pub fn new(service_name: &str) -> Self {
        Self {
            service_name: service_name.to_string(),
            span_queue: Arc::new(SegQueue::new()),
            dropped_spans: AtomicU64::new(0),
            spans_created: AtomicU64::new(0),
            spans_exported: AtomicU64::new(0),
        }
    }

    /// Start new span (ultra-fast path)
    #[inline(always)]
    pub fn start_span(&self, operation_name: &str) -> FastSpan {
        self.spans_created.fetch_add(1, Ordering::Relaxed);
        FastSpan::new(operation_name, &self.service_name)
    }

    /// Start child span
    #[inline(always)]
    pub fn start_child_span(&self, parent: &FastSpan, operation_name: &str) -> FastSpan {
        self.spans_created.fetch_add(1, Ordering::Relaxed);
        parent.child(operation_name)
    }

    /// Finish and submit span for export (ultra-fast path)
    #[inline(always)]
    pub fn finish_span(&self, mut span: FastSpan) {
        span.finish();
        self.submit_span(span);
    }

    /// Submit completed span to export queue
    #[inline(always)]
    pub fn submit_span(&self, span: FastSpan) {
        // Check queue size to prevent memory exhaustion
        if self.span_queue.len() < MAX_SPAN_BUFFER_SIZE {
            self.span_queue.push(span);
        } else {
            self.dropped_spans.fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Drain spans for export (called by background thread)
    pub fn drain_spans(&self, max_count: usize) -> Vec<FastSpan> {
        let mut spans = Vec::with_capacity(max_count.min(SPAN_EXPORT_BATCH_SIZE));
        
        for _ in 0..max_count {
            if let Some(span) = self.span_queue.pop() {
                spans.push(span);
            } else {
                break;
            }
        }
        
        self.spans_exported.fetch_add(spans.len() as u64, Ordering::Relaxed);
        spans
    }

    /// Get tracer statistics
    pub fn stats(&self) -> TracerStats {
        TracerStats {
            spans_created: self.spans_created.load(Ordering::Relaxed),
            spans_exported: self.spans_exported.load(Ordering::Relaxed),
            spans_dropped: self.dropped_spans.load(Ordering::Relaxed),
            queue_depth: self.span_queue.len(),
            service_name: self.service_name.clone(),
        }
    }
}

/// Tracer performance statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
/// TracerStats
/// 
/// TODO: Add detailed documentation for this struct
pub struct TracerStats {
    /// Spans Created
    pub spans_created: u64,
    /// Spans Exported
    pub spans_exported: u64,
    /// Spans Dropped
    pub spans_dropped: u64,
    /// Queue Depth
    pub queue_depth: usize,
    /// Service Name
    pub service_name: String,
}

/// Span context for correlation across service boundaries
#[derive(Debug, Clone, Serialize, Deserialize)]
/// SpanContext
/// 
/// TODO: Add detailed documentation for this struct
pub struct SpanContext {
    /// Trace Id
    pub trace_id: u128,
    /// Span Id
    pub span_id: u64,
    /// Sampled
    pub sampled: bool,
}

impl SpanContext {
    /// Create from FastSpan
    pub fn from_span(span: &FastSpan) -> Self {
        Self {
            trace_id: span.trace_id,
            span_id: span.span_id,
            sampled: true,
        }
    }

    /// Encode as HTTP header value
    pub fn to_header_value(&self) -> String {
        format!("{:032x}-{:016x}-1", self.trace_id, self.span_id)
    }

    /// Decode from HTTP header value
    pub fn from_header_value(header: &str) -> Result<Self> {
        let parts: Vec<&str> = header.split('-').collect();
        if parts.len() != 3 {
            return Err(anyhow!("Invalid trace header format"));
        }

        let trace_id = u128::from_str_radix(parts[0], 16)
            .map_err(|_| anyhow!("Invalid trace ID"))?;
        
        let span_id = u64::from_str_radix(parts[1], 16)
            .map_err(|_| anyhow!("Invalid span ID"))?;
        
        let sampled = parts[2] == "1";

        Ok(Self {
            trace_id,
            span_id,
            sampled,
        })
    }
}

/// RAII span guard for automatic span finishing
#[derive(Debug)]
pub struct SpanGuard<'a> {
    tracer: &'a FastTracer,
    span: FastSpan,
}

impl<'a> SpanGuard<'a> {
    pub fn new(tracer: &'a FastTracer, span: FastSpan) -> Self {
        Self { tracer, span }
    }

    /// Get mutable reference to the span
    pub fn span_mut(&mut self) -> &mut FastSpan {
        &mut self.span
    }

    /// Get reference to the span
    pub fn span(&self) -> &FastSpan {
        &self.span
    }

    /// Set span status
    pub fn set_status(&mut self, status: SpanStatus) {
        self.span.status = status;
    }

    /// Add tag to span
    pub fn set_tag(&mut self, key: &str, value: &str) {
        self.span.set_tag(key, value);
    }
}

impl<'a> Drop for SpanGuard<'a> {
    fn drop(&mut self) {
        self.tracer.finish_span(std::mem::take(&mut self.span));
    }
}

/// Global tracer registry
static mut GLOBAL_TRACER: Option<FastTracer> = None;
static TRACER_INIT: std::sync::Once = std::sync::Once::new();

/// Initialize global tracer
pub fn init_global_tracer(service_name: &str) {
    TRACER_INIT.call_once(|| {
        unsafe {
            GLOBAL_TRACER = Some(FastTracer::new(service_name));
        }
    });
}

/// Get global tracer reference
pub fn global_tracer() -> &'static FastTracer {
    unsafe {
        GLOBAL_TRACER.as_ref().expect("Global tracer not initialized. Call init_global_tracer() first.")
    }
}

/// Generate unique span ID using atomic counter
fn generate_span_id() -> u64 {
    static SPAN_ID_COUNTER: AtomicU64 = AtomicU64::new(1);
    SPAN_ID_COUNTER.fetch_add(1, Ordering::Relaxed)
}

/// Generate unique trace ID
fn generate_trace_id() -> u128 {
    let uuid = Uuid::new_v4();
    uuid.as_u128()
}

/// Convenience macros for span creation
#[macro_export]
macro_rules! trace_span {
    ($operation:expr) => {{
        $crate::tracing::global_tracer().start_span($operation)
    }};
}

/// Creates a span guard that automatically ends the span when dropped
/// 
/// # Arguments
/// * `operation` - The operation name for the span
/// 
/// # Example
/// ```
/// let _guard = trace_span_guard!("database_query");
/// // span automatically ends when guard is dropped
/// ```
#[macro_export]
macro_rules! trace_span_guard {
    ($operation:expr) => {{
        let tracer = $crate::tracing::global_tracer();
        let span = tracer.start_span($operation);
        $crate::tracing::SpanGuard::new(tracer, span)
    }};
}

/// Creates a child span under an existing parent span
/// 
/// # Arguments
/// * `parent` - The parent span
/// * `operation` - The operation name for the child span
/// 
/// # Example
/// ```
/// let parent_span = trace_span!("parent_operation");
/// let child_span = trace_child_span!(parent_span, "child_operation");
/// ```
#[macro_export]
macro_rules! trace_child_span {
    ($parent:expr, $operation:expr) => {{
        $crate::tracing::global_tracer().start_child_span($parent, $operation)
    }};
}

/// Span export configuration
#[derive(Debug, Clone)]
/// SpanExportConfig
/// 
/// TODO: Add detailed documentation for this struct
pub struct SpanExportConfig {
    /// Jaeger Endpoint
    pub jaeger_endpoint: String,
    /// Batch Size
    pub batch_size: usize,
    /// Export Timeout Ms
    pub export_timeout_ms: u64,
    /// Max Queue Size
    pub max_queue_size: usize,
}

impl Default for SpanExportConfig {
    fn default() -> Self {
        Self {
            jaeger_endpoint: "http://localhost:14268/api/traces".to_string(),
            batch_size: SPAN_EXPORT_BATCH_SIZE,
            export_timeout_ms: 5000,
            max_queue_size: MAX_SPAN_BUFFER_SIZE,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;
    use std::time::Duration;

    #[test]
    fn test_span_creation() {
        let span = FastSpan::new("test_operation", "test_service");
        assert_eq!(span.operation_name, "test_operation");
        assert_eq!(span.service_name, "test_service");
        assert_eq!(span.parent_id, 0);
        assert!(span.start_time_ns > 0);
    }

    #[test]
    fn test_child_span() {
        let parent = FastSpan::new("parent", "test_service");
        let child = parent.child("child");
        
        assert_eq!(child.parent_id, parent.span_id);
        assert_eq!(child.trace_id, parent.trace_id);
        assert_eq!(child.operation_name, "child");
    }

    #[test]
    fn test_span_finish() {
        let mut span = FastSpan::new("test", "service");
        thread::sleep(Duration::from_millis(1));
        span.finish();
        
        assert!(span.end_time_ns > span.start_time_ns);
        assert!(span.duration_ns() > 0);
    }

    #[test]
    fn test_tracer_operations() {
        let tracer = FastTracer::new("test_service");
        
        let span = tracer.start_span("test_op");
        tracer.finish_span(span);
        
        let stats = tracer.stats();
        assert_eq!(stats.spans_created, 1);
        assert_eq!(stats.service_name, "test_service");
    }

    #[test]
    fn test_span_context() {
        let span = FastSpan::new("test", "service");
        let context = SpanContext::from_span(&span);
        
        let header_value = context.to_header_value();
        let parsed_context = SpanContext::from_header_value(&header_value).unwrap();
        
        assert_eq!(context.trace_id, parsed_context.trace_id);
        assert_eq!(context.span_id, parsed_context.span_id);
    }

    #[test]
    fn test_span_guard() {
        let tracer = FastTracer::new("test_service");
        
        {
            let mut guard = SpanGuard::new(&tracer, tracer.start_span("test"));
            guard.set_tag("key", "value");
        } // Span should be finished and submitted here
        
        let spans = tracer.drain_spans(10);
        assert_eq!(spans.len(), 1);
        assert!(!spans[0].tags.is_empty());
    }
}