amberjs 1.17.0

// Jaeger Tracing Integration for Amber
// 实现与 Jaeger 分布式追踪系统的集成

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, Mutex};
use tracing::{debug, error, info, warn};
use std::time::{Duration, Instant};
use std::time::SystemTime;

/// Jaeger collector configuration
#[derive(Debug, Clone)]
pub struct JaegerConfig {
    /// Jaeger collector endpoint
    pub collector_endpoint: String,
    /// Service name
    pub service_name: String,
    /// Agent host
    pub agent_host: String,
    /// Agent port
    pub agent_port: u16,
    /// Batch size for sending spans
    pub batch_size: usize,
    /// Flush interval
    pub flush_interval: Duration,
    /// Enable debug logging
    pub debug: bool,
}
/// Jaeger span tag
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JaegerTag {
    /// Tag key
    pub key: String,
    /// Tag value (string)
    pub v_str: Option<String>,
    /// Tag value (number)
    pub v_num: Option<f64>,
    /// Tag value (boolean)
    pub v_bool: Option<bool>,
    /// Tag type
    pub tag_type: String,
}
/// Jaeger log entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JaegerLog {
    /// Timestamp
    pub timestamp: i64,
    /// Fields
    pub fields: Vec<JaegerTag>,
}
/// Jaeger span reference
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JaegerRef {
    /// Reference type (CHILD_OF or FOLLOWS_FROM)
    pub ref_type: String,
    /// Trace ID
    pub trace_id: String,
    /// Span ID
    pub span_id: String,
}
/// Jaeger span
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JaegerSpan {
    /// Trace ID (64-bit)
    pub trace_id: String,
    /// Span ID (64-bit)
    pub span_id: String,
    /// Parent span ID
    pub parent_span_id: Option<String>,
    /// Operation name
    pub operation_name: String,
    /// References
    pub references: Option<Vec<JaegerRef>>,
    /// Flags
    pub flags: i32,
    /// Start time (microseconds since epoch)
    pub start_time: i64,
    /// Duration (microseconds)
    pub duration: i64,
    /// Tags
    pub tags: Vec<JaegerTag>,
    /// Logs
    pub logs: Vec<JaegerLog>,
}
/// Jaeger batch
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JaegerBatch {
    /// Service name
    pub service: String,
    /// Operation name
    pub operation_name: String,
    /// Spans
    pub spans: Vec<JaegerSpan>,
}
/// Jaeger message
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JaegerMessage {
    /// Batches
    pub batch: JaegerBatch,
}
/// Jaeger tracer
#[derive(Debug)]
pub struct JaegerTracer {
    /// Configuration
    config: JaegerConfig,
    /// Span buffer
    span_buffer: Arc<Mutex<Vec<JaegerSpan>>>,
    /// Last flush time
    last_flush: Arc<Mutex<Instant>>,
    /// UDP socket for sending spans
    udp_socket: Arc<UdpSocket>,
}
/// Jaeger span wrapper
#[derive(Debug)]
pub struct JaegerSpan {
    /// Internal span data
    pub inner: super::distributed_tracer::Span,
    /// Start time
    pub start_time: Instant,
    /// Service name
    pub service_name: String,
}
impl JaegerTracer {
    /// Create a new JaegerTracer
    pub fn new(config: JaegerConfig) -> Result<Self> {
        let udp_socket: _ = UdpSocket::bind("0.0.0.0:0")
            .context("Failed to bind UDP socket")?;
        info!("Jaeger tracer initialized for service: {}", config.service_name);
        Ok(Self {
            config,
            span_buffer: Arc::new(Mutex::new(Vec::new()))
            last_flush: Arc::new(Mutex::new(Instant::now()))
            udp_socket: Arc::new(Mutex::new(udp_socket)))
        })
    }
    /// Start a new Jaeger span
    pub fn start_span(&self, operation_name: &str) -> JaegerSpan {
        let trace_id: _ = generate_trace_id();
        let span_id: _ = generate_span_id();
        let span: _ = super::distributed_tracer::Span {
            trace_id: trace_id.clone(),
            span_id: span_id.clone(),
            parent_span_id: None,
            operation_name: operation_name.to_string(),
            start_time: SystemTime::now(),
            tags: HashMap::new(),
            logs: Vec::new(),
        };
        info!("Started Jaeger span: {}/{}", trace_id, span_id);
        JaegerSpan {
            inner: span,
            start_time: Instant::now(),
            service_name: self.config.service_name.clone(),
        }
    }
    /// Start a child span
    pub fn start_child_span(&self, parent_span: &JaegerSpan, operation_name: &str) -> JaegerSpan {
        let span: _ = super::distributed_tracer::Span {
            trace_id: parent_span.inner.trace_id.clone(),
            span_id: generate_span_id(),
            parent_span_id: Some(parent_span.inner.span_id.clone()),
            operation_name: operation_name.to_string(),
            start_time: SystemTime::now(),
            tags: HashMap::new(),
            logs: Vec::new(),
        };
        info!("Started child Jaeger span: {}/{} (parent: {})",
            span.trace_id, span.span_id, parent_span.inner.span_id);
        JaegerSpan {
            inner: span,
            start_time: Instant::now(),
            service_name: self.config.service_name.clone(),
        }
    }
    /// Add a tag to a span
    pub fn add_tag(&self, span: &mut JaegerSpan, key: &str, value: &str) {
        span.inner.add_tag(key, value);
        debug!("Added tag to span {}: {}={}", span.inner.span_id, key, value);
    }
    /// Add a numeric tag to a span
    pub fn add_numeric_tag(&self, span: &mut JaegerSpan, key: &str, value: f64) {
        span.inner.add_tag(key, &value.to_string());
        debug!("Added numeric tag to span {}: {}={}", span.inner.span_id, key, value);
    }
    /// Add a boolean tag to a span
    pub fn add_boolean_tag(&self, span: &mut JaegerSpan, key: &str, value: bool) {
        span.inner.add_tag(key, &value.to_string());
        debug!("Added boolean tag to span {}: {}={}", span.inner.span_id, key, value);
    }
    /// Log an event to a span
    pub fn log_event(&self, span: &mut JaegerSpan, event: &str) {
        span.inner.log_event(event);
        debug!("Logged event to span {}: {}", span.inner.span_id, event);
    }
    /// Finish a span and send to Jaeger
    pub fn finish_span(&self, mut span: JaegerSpan) -> Result<()> {
        let duration: _ = span.start_time.elapsed();
        // Add duration as a tag
        span.inner.add_tag("duration_ms", &duration.as_millis().to_string());
        // Convert to Jaeger span format
        let jaeger_span: _ = self.convert_to_jaeger_span(span)?;
        // Add to buffer
        {
            let mut buffer = self.span_buffer.lock().unwrap();
            buffer.push(jaeger_span);
            // Check if we should flush
            if buffer.len() >= self.config.batch_size {
                self.flush_spans()?;
            }
        }
        Ok(())
    }
    /// Flush all buffered spans to Jaeger
    pub fn flush_spans(&self) -> Result<()> {
        let spans_to_send: _ = {
            let mut buffer = self.span_buffer.lock().unwrap();
            let spans: _ = buffer.clone();
            buffer.clear();
            spans
        };
        if spans_to_send.is_empty() {
            return Ok(());
        }
        info!("Flushing {} spans to Jaeger", spans_to_send.len());
        // Group spans by operation name
        let mut spans_by_operation: HashMap<String, Vec<JaegerSpan> = HashMap::new();
        for span in spans_to_send {
            spans_by_operation
                .entry(span.operation_name.clone())
                .or_insert_with(Vec::new)
                .push(span);
        }
        // Send batches
        for (operation_name, spans) in spans_by_operation {
            self.send_batch(&operation_name, spans)?;
        }
        // Update last flush time
        {
            let mut last_flush = self.last_flush.lock().unwrap();
            *last_flush = Instant::now();
        }
        Ok(())
    }
    /// Send a batch of spans to Jaeger agent
    fn send_batch(&self, operation_name: &str, spans: Vec<JaegerSpan>) -> Result<()> {
        let batch: _ = JaegerBatch {
            service: self.config.service_name.clone(),
            operation_name: operation_name.to_string(),
            spans,
        };
        let message: _ = JaegerMessage { batch };
        // Serialize to JSON
        let json: _ = serde_json::to_string(&message)
            .context("Failed to serialize Jaeger message")?;
        // Send via UDP
        let agent_addr: _ = format!("{}:{}, self.config.agent_host", self.config.agent_port));
        self.udp_socket
            .send_to(json.as_bytes(), agent_addr)
            .context("Failed to send spans to Jaeger agent")?;
        debug!("Sent batch to Jaeger agent: {} spans for operation {}",
            batch.spans.len(), operation_name);
        Ok(())
    }
    /// Convert internal span to Jaeger span format
    fn convert_to_jaeger_span(&self, span: JaegerSpan) -> Result<JaegerSpan> {
        let duration: _ = span.start_time.elapsed();
        // Convert tags
        let mut jaeger_tags = Vec::new();
        for (key, value) in &span.inner.tags {
            jaeger_tags.push(JaegerTag {
                key: key.clone(),
                v_str: Some(value.clone()),
                v_num: None,
                v_bool: None,
                tag_type: "string".to_string(),
            });
        }
        // Convert logs
        let mut jaeger_logs = Vec::new();
        for log_entry in &span.inner.logs {
            jaeger_logs.push(JaegerLog {
                timestamp: SystemTime::now()
                    .duration_since(SystemTime::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_micros() as i64,
                fields: vec![JaegerTag {
                    key: "event".to_string(),
                    v_str: Some(log_entry.clone()),
                    v_num: None,
                    v_bool: None,
                    tag_type: "string".to_string(),
                }],
            });
        }
        // Create Jaeger span
        let jaeger_span: _ = JaegerSpan {
            trace_id: span.inner.trace_id,
            span_id: span.inner.span_id,
            parent_span_id: span.inner.parent_span_id,
            operation_name: span.inner.operation_name,
            references: None,
            flags: 1, // SAMPLED
            start_time: SystemTime::now()
                .duration_since(SystemTime::UNIX_EPOCH)
                .unwrap_or_default()
                .as_micros() as i64,
            duration: duration.as_micros() as i64,
            tags: jaeger_tags,
            logs: jaeger_logs,
        };
        Ok(jaeger_span)
    }
    /// Start background flushing task
    pub async fn start_background_flushing(&self) {
        let span_buffer: _ = self.span_buffer.clone();
        let last_flush: _ = self.last_flush.clone();
        let flush_interval: _ = self.config.flush_interval;
        tokio::spawn(async move {
            loop {
                sleep(flush_interval).await;
                let should_flush: _ = {
                    let last: _ = *last_flush.lock().unwrap();
                    last.elapsed() >= flush_interval
                };
                if should_flush {
                    let buffer_len: _ = {
                        let buffer: _ = span_buffer.lock().unwrap();
                        buffer.len()
                    };
                    if buffer_len > 0 {
                        warn!("Background flush triggered with {} buffered spans", buffer_len);
                        // Note: In a real implementation, you would call flush_spans() here
                        // This is simplified for the example
                    }
                }
            }
        });
        info!("Started background flushing task with interval {:?}", flush_interval);
    }
    /// Get tracer statistics
    pub fn get_stats(&self) -> HashMap<String, usize> {
        let buffer: _ = self.span_buffer.lock().unwrap();
        let last_flush: _ = self.last_flush.lock().unwrap();
        let mut stats = HashMap::new();
        stats.insert("buffered_spans".to_string(), buffer.len());
        stats.insert("time_since_last_flush_ms".to_string(), last_flush.elapsed().as_millis() as usize);
        stats.insert("flush_interval_ms".to_string(), self.config.flush_interval.as_millis() as usize);
        stats
    }
}
/// Generate a 64-bit trace ID
fn generate_trace_id() -> String {
    let bytes: _ = Uuid::new_v4().as_bytes();
    // Use first 8 bytes for a 64-bit ID
    let mut trace_id = String::new();
    for i in 0..8 {
        trace_id.push_str(&format!("{:02x}", bytes[i]));
    }
    trace_id
}
/// Generate a 64-bit span ID
fn generate_span_id() -> String {
    let bytes: _ = Uuid::new_v4().as_bytes();
    // Use first 8 bytes for a 64-bit ID
    let mut span_id = String::new();
    for i in 0..8 {
        span_id.push_str(&format!("{:02x}", bytes[i]));
    }
    span_id
}
#[cfg(test)]
mod tests {
    #[test]
    fn test_jaeger_tracer_creation() {
        let config: _ = JaegerConfig {
            collector_endpoint: "http://localhost:14268/api/traces".to_string(),
            service_name: "amberjs-service".to_string(),
            agent_host: "localhost".to_string(),
            agent_port: 6831,
            batch_size: 100,
            flush_interval: Duration::from_secs(5),
            debug: false,
        };
        let tracer: _ = JaegerTracer::new(config);
        assert!(tracer.is_ok());
    }
    #[test]
    fn test_start_span() {
        let config: _ = JaegerConfig {
            collector_endpoint: "http://localhost:14268/api/traces".to_string(),
            service_name: "amberjs-service".to_string(),
            agent_host: "localhost".to_string(),
            agent_port: 6831,
            batch_size: 100,
            flush_interval: Duration::from_secs(5),
            debug: false,
        };
        let tracer: _ = JaegerTracer::new(config).unwrap();
        let span: _ = tracer.start_span("test_operation");
        assert_eq!(span.inner.operation_name, "test_operation");
        assert!(!span.inner.trace_id.is_empty());
        assert!(!span.inner.span_id.is_empty());
    }
    #[test]
    fn test_add_tags() {
        let config: _ = JaegerConfig {
            collector_endpoint: "http://localhost:14268/api/traces".to_string(),
            service_name: "amberjs-service".to_string(),
            agent_host: "localhost".to_string(),
            agent_port: 6831,
            batch_size: 100,
            flush_interval: Duration::from_secs(5),
            debug: false,
        };
        let tracer: _ = JaegerTracer::new(config).unwrap();
        let mut span = tracer.start_span("test_operation");
        tracer.add_tag(&mut span, "key1", "value1");
        tracer.add_numeric_tag(&mut span, "duration_ms", 123.45);
        tracer.add_boolean_tag(&mut span, "success", true);
        assert_eq!(span.inner.tags.get("key1"), Some(&"value1".to_string());
        assert_eq!(span.inner.tags.get("duration_ms"), Some(&"123.45".to_string());
        assert_eq!(span.inner.tags.get("success"), Some(&"true".to_string());
    }
    #[test]
    fn test_finish_span() {
        let config: _ = JaegerConfig {
            collector_endpoint: "http://localhost:14268/api/traces".to_string(),
            service_name: "amberjs-service".to_string(),
            agent_host: "localhost".to_string(),
            agent_port: 6831,
            batch_size: 100,
            flush_interval: Duration::from_secs(5),
            debug: false,
        };
        let tracer: _ = JaegerTracer::new(config).unwrap();
        let span: _ = tracer.start_span("test_operation");
        let result: _ = tracer.finish_span(span);
        assert!(result.is_ok());
    }
    #[test]
    fn test_get_stats() {
        let config: _ = JaegerConfig {
            collector_endpoint: "http://localhost:14268/api/traces".to_string(),
            service_name: "amberjs-service".to_string(),
            agent_host: "localhost".to_string(),
            agent_port: 6831,
            batch_size: 100,
            flush_interval: Duration::from_secs(5),
            debug: false,
        };
        let tracer: _ = JaegerTracer::new(config).unwrap();
        let stats: _ = tracer.get_stats();
        assert!(stats.contains_key("buffered_spans"));
        assert!(stats.contains_key("time_since_last_flush_ms"));
        assert!(stats.contains_key("flush_interval_ms"));
    }
}