dmsc/observability/
tracing.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of DMSC.
4//! The DMSC project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18#![allow(non_snake_case)]
19
20//! # Distributed Tracing
21//!
22//! This file implements a comprehensive distributed tracing system for the DMSC framework. It provides
23//! tools for creating, managing, and propagating trace information across asynchronous operations
24//! and distributed systems. The tracing system follows the W3C Trace Context standard and integrates
25//! with tokio's context propagation mechanism.
26//!
27//! ## Key Components
28//!
29//! - **DMSCSpanId**: Unique identifier for a span
30//! - **DMSCTraceId**: Unique identifier for a trace
31//! - **DMSCSpanKind**: Enumeration of span types (Server, Client, Producer, Consumer, Internal)
32//! - **DMSCSpanStatus**: Status of a span (Ok, Error, Unset)
33//! - **DMSCSpan**: A single distributed tracing span with attributes, events, and status
34//! - **DMSCSpanEvent**: Timed events within a span
35//! - **DMSCTracingContext**: Thread-local tracing context for propagation
36//! - **DMSCTracer**: Distributed tracer for creating and managing spans
37//! - **DMSCTracerManager**: Manager for multiple tracer instances
38//! - **DefaultTracerManager**: Global tracer manager instance
39//!
40//! ## Design Principles
41//!
42//! 1. **W3C Trace Context Compliance**: Follows the W3C Trace Context standard for interoperability
43//! 2. **Async Context Propagation**: Integrates with tokio's context propagation mechanism
44//! 3. **Thread Safety**: Uses Arc and RwLock for safe concurrent access
45//! 4. **Sampling Support**: Configurable sampling rate to control overhead
46//! 5. **Hierarchical Spans**: Supports parent-child span relationships
47//! 6. **Baggage Support**: Allows carrying contextual information across spans
48//! 7. **Extensible**: Easy to add new span kinds and attributes
49//! 8. **Low Overhead**: Efficient implementation with minimal performance impact
50//! 9. **Global Access**: Provides a global tracer manager for easy access
51//! 10. **Serialization Support**: All tracing components are serializable for export
52//!
53//! ## Usage
54//!
55//! ```rust
56//! use dmsc::observability::{init_tracer, tracer, DMSCSpanKind, DMSCSpanStatus};
57//! use dmsc::core::DMSCResult;
58//!
59//! async fn example() -> DMSCResult<()> {
60//!     // Initialize the global tracer with 100% sampling rate
61//!     init_tracer(1.0);
62//!     
63//!     // Get the global tracer
64//!     let tracer = tracer();
65//!     
66//!     // Start a new trace
67//!     let trace_id = tracer.start_trace("example_trace").unwrap();
68//!     
69//!     // Start a child span
70//!     let span_id = tracer.start_span_from_context("child_span", DMSCSpanKind::Internal).unwrap();
71//!     
72//!     // Add an attribute to the span
73//!     tracer.span_mut(&span_id, |span| {
74//!         span.set_attribute("key".to_string(), "value".to_string());
75//!     })?;
76//!     
77//!     // Add an event to the span
78//!     tracer.span_mut(&span_id, |span| {
79//!         let mut attributes = std::collections::HashMap::new();
80//!         attributes.insert("event_key".to_string(), "event_value".to_string());
81//!         span.add_event("example_event".to_string(), attributes);
82//!     })?;
83//!     
84//!     // End the child span with OK status
85//!     tracer.end_span(&span_id, DMSCSpanStatus::Ok)?;
86//!     
87//!     Ok(())
88//! }
89//! ```
90
91use serde::{Deserialize, Serialize};
92use std::cell::RefCell;
93use std::collections::HashMap;
94use std::sync::{Arc, RwLock};
95use std::time::{Duration, SystemTime, UNIX_EPOCH};
96use uuid::Uuid;
97
98use crate::core::DMSCResult;
99use crate::core::DMSCError;
100use crate::core::lock::RwLockExtensions;
101
102#[cfg(feature = "pyo3")]
103use pyo3::prelude::*;
104
105/// Distributed tracing span ID
106#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
107pub struct DMSCSpanId(String);
108
109impl Default for DMSCSpanId {
110    fn default() -> Self {
111        Self::new()
112    }
113}
114
115impl DMSCSpanId {
116    pub fn new() -> Self {
117        Self(Uuid::new_v4().to_string())
118    }
119
120    pub fn from_string(s: String) -> Self {
121        Self(s)
122    }
123
124    pub fn as_str(&self) -> &str {
125        &self.0
126    }
127}
128
129/// Distributed tracing trace ID
130#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
131pub struct DMSCTraceId(String);
132
133impl Default for DMSCTraceId {
134    fn default() -> Self {
135        Self::new()
136    }
137}
138
139impl DMSCTraceId {
140    pub fn new() -> Self {
141        Self(Uuid::new_v4().to_string())
142    }
143
144    pub fn from_string(s: String) -> Self {
145        Self(s)
146    }
147
148    pub fn as_str(&self) -> &str {
149        &self.0
150    }
151}
152
153/// Span kind enumeration
154#[derive(Debug, Clone, Serialize, Deserialize)]
155pub enum DMSCSpanKind {
156    Server,
157    Client,
158    Producer,
159    Consumer,
160    Internal,
161}
162
163/// Span status
164#[derive(Debug, Clone, Serialize, Deserialize)]
165pub enum DMSCSpanStatus {
166    Ok,
167    Error(String),
168    Unset,
169}
170
171/// A distributed tracing span
172#[derive(Debug, Clone, Serialize, Deserialize)]
173pub struct DMSCSpan {
174    pub trace_id: DMSCTraceId,
175    pub span_id: DMSCSpanId,
176    pub parent_span_id: Option<DMSCSpanId>,
177    pub name: String,
178    pub kind: DMSCSpanKind,
179    pub start_time: u64, // microseconds since epoch
180    pub end_time: Option<u64>,
181    pub attributes: HashMap<String, String>,
182    pub events: Vec<DMSCSpanEvent>,
183    pub status: DMSCSpanStatus,
184}
185
186impl DMSCSpan {
187    pub fn new(
188        trace_id: DMSCTraceId,
189        parent_span_id: Option<DMSCSpanId>,
190        name: String,
191        kind: DMSCSpanKind,
192    ) -> Self {
193        let start_time = SystemTime::now()
194            .duration_since(UNIX_EPOCH)
195            .unwrap_or(Duration::from_secs(0))
196            .as_micros() as u64;
197
198        Self {
199            trace_id,
200            span_id: DMSCSpanId::new(),
201            parent_span_id,
202            name,
203            kind,
204            start_time,
205            end_time: None,
206            attributes: HashMap::new(),
207            events: Vec::new(),
208            status: DMSCSpanStatus::Unset,
209        }
210    }
211
212    pub fn set_attribute(&mut self, key: String, value: String) {
213        self.attributes.insert(key, value);
214    }
215
216    pub fn add_event(&mut self, name: String, attributes: HashMap<String, String>) {
217        let timestamp = SystemTime::now()
218            .duration_since(UNIX_EPOCH)
219            .unwrap_or(Duration::from_secs(0))
220            .as_micros() as u64;
221
222        self.events.push(DMSCSpanEvent {
223            name,
224            timestamp,
225            attributes,
226        });
227    }
228
229    pub fn end(&mut self, status: DMSCSpanStatus) {
230        let end_time = SystemTime::now()
231            .duration_since(UNIX_EPOCH)
232            .unwrap_or(Duration::from_secs(0))
233            .as_micros() as u64;
234
235        self.end_time = Some(end_time);
236        self.status = status;
237    }
238
239    pub fn duration(&self) -> Option<Duration> {
240        if let Some(end_time) = self.end_time {
241            let duration_micros = end_time.saturating_sub(self.start_time);
242            Some(Duration::from_micros(duration_micros))
243        } else {
244            None
245        }
246    }
247}
248
249/// Span event for recording timed occurrences
250#[derive(Debug, Clone, Serialize, Deserialize)]
251pub struct DMSCSpanEvent {
252    pub name: String,
253    pub timestamp: u64, // microseconds since epoch
254    pub attributes: HashMap<String, String>,
255}
256
257/// Thread-local tracing context
258#[derive(Debug, Clone)]
259pub struct DMSCTracingContext {
260    current_trace_id: Option<DMSCTraceId>,
261    current_span_id: Option<DMSCSpanId>,
262    baggage: HashMap<String, String>,
263}
264
265// Thread-local storage for tracing context
266thread_local! {
267    static CURRENTONTEXT: RefCell<Option<DMSCTracingContext>> = const { RefCell::new(None) };
268}
269
270impl Default for DMSCTracingContext {
271    fn default() -> Self {
272        Self::new()
273    }
274}
275
276impl DMSCTracingContext {
277    pub fn new() -> Self {
278        Self {
279            current_trace_id: None,
280            current_span_id: None,
281            baggage: HashMap::new(),
282        }
283    }
284
285    pub fn with_trace_id(mut self, trace_id: DMSCTraceId) -> Self {
286        self.current_trace_id = Some(trace_id);
287        self
288    }
289
290    pub fn with_span_id(mut self, span_id: DMSCSpanId) -> Self {
291        self.current_span_id = Some(span_id);
292        self
293    }
294
295    pub fn set_baggage(&mut self, key: String, value: String) {
296        self.baggage.insert(key, value);
297    }
298
299    pub fn get_baggage(&self, key: &str) -> Option<&String> {
300        self.baggage.get(key)
301    }
302
303    pub fn trace_id(&self) -> Option<&DMSCTraceId> {
304        self.current_trace_id.as_ref()
305    }
306
307    pub fn span_id(&self) -> Option<&DMSCSpanId> {
308        self.current_span_id.as_ref()
309    }
310
311    /// Set this context as the current thread-local context
312    pub fn set_as_current(&self) {
313        CURRENTONTEXT.with(|ctx| {
314            *ctx.borrow_mut() = Some(self.clone());
315        });
316    }
317
318    /// Get the current tracing context from thread-local storage
319    pub fn current() -> Option<Self> {
320        CURRENTONTEXT.with(|ctx| {
321            ctx.borrow().clone()
322        })
323    }
324
325    /// Create a new context with the same trace ID but new span ID
326    pub fn new_child(&self, span_id: DMSCSpanId) -> Self {
327        Self {
328            current_trace_id: self.current_trace_id.clone(),
329            current_span_id: Some(span_id),
330            baggage: self.baggage.clone(),
331        }
332    }
333}
334
335/// Sampling strategy enumeration
336#[derive(Debug, Clone, Serialize, Deserialize)]
337pub enum DMSCSamplingStrategy {
338    /// Fixed rate sampling (0.0 to 1.0)
339    Rate(f64),
340    /// Trace ID-based deterministic sampling
341    Deterministic(f64),
342    /// Adaptive sampling that adjusts based on load
343    Adaptive(f64),
344}
345
346/// Distributed tracer
347#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
348pub struct DMSCTracer {
349    spans: Arc<RwLock<HashMap<DMSCTraceId, Vec<DMSCSpan>>>>,
350    active_spans: Arc<RwLock<HashMap<DMSCSpanId, DMSCSpan>>>,
351    sampling_strategy: DMSCSamplingStrategy,
352    adaptive_window: Arc<RwLock<Vec<u64>>>,
353    max_adaptive_window: usize,
354}
355
356impl DMSCTracer {
357    pub fn new(sampling_rate: f64) -> Self {
358        Self {
359            spans: Arc::new(RwLock::new(HashMap::new())),
360            active_spans: Arc::new(RwLock::new(HashMap::new())),
361            sampling_strategy: DMSCSamplingStrategy::Rate(sampling_rate.clamp(0.0, 1.0)),
362            adaptive_window: Arc::new(RwLock::new(Vec::new())),
363            max_adaptive_window: 100,
364        }
365    }
366    
367    /// Create a new tracer with a custom sampling strategy
368    pub fn with_strategy(strategy: DMSCSamplingStrategy) -> Self {
369        Self {
370            spans: Arc::new(RwLock::new(HashMap::new())),
371            active_spans: Arc::new(RwLock::new(HashMap::new())),
372            sampling_strategy: strategy,
373            adaptive_window: Arc::new(RwLock::new(Vec::new())),
374            max_adaptive_window: 100,
375        }
376    }
377
378    /// Start a new trace and set it as current context
379    pub fn start_trace(&self, name: String) -> Option<DMSCTraceId> {
380        if !self.should_sample() {
381            return None;
382        }
383
384        let trace_id = DMSCTraceId::new();
385        let span = DMSCSpan::new(trace_id.clone(), None, name, DMSCSpanKind::Server);
386
387        let span_id = span.span_id.clone();
388        {
389            let mut active_spans = self.active_spans.write_safe("active spans for new trace").ok()?;
390            active_spans.insert(span_id.clone(), span);
391        }
392        {
393            let mut spans = self.spans.write_safe("spans for new trace").ok()?;
394            spans.insert(trace_id.clone(), Vec::new());
395        }
396
397        // Set current context
398        let context = DMSCTracingContext::new()
399            .with_trace_id(trace_id.clone())
400            .with_span_id(span_id);
401        context.set_as_current();
402
403        Some(trace_id)
404    }
405
406    /// Start a new span in existing trace, using current context if available
407    pub fn start_span(
408        &self,
409        trace_id: Option<&DMSCTraceId>,
410        parent_span_id: Option<DMSCSpanId>,
411        name: String,
412        kind: DMSCSpanKind,
413    ) -> Option<DMSCSpanId> {
414        // Try to get trace_id from current context if not provided
415        let resolved_trace_id = match trace_id {
416            Some(id) => id.clone(),
417            None => {
418                if let Some(context) = DMSCTracingContext::current() {
419                    if let Some(id) = context.trace_id() {
420                        id.clone()
421                    } else {
422                        return None;
423                    }
424                } else {
425                    return None;
426                }
427            }
428        };
429
430        // Try to get parent_span_id from current context if not provided
431        let resolved_parent_span_id = match parent_span_id {
432            Some(id) => Some(id.clone()),
433            None => DMSCTracingContext::current().and_then(|context| context.span_id().cloned()),
434        };
435
436        let spans = match self.spans.read_safe("spans for span check") {
437            Ok(s) => s,
438            Err(_) => return None,
439        };
440        if !spans.contains_key(&resolved_trace_id) {
441            return None;
442        }
443
444        let span = DMSCSpan::new(
445            resolved_trace_id.clone(),
446            resolved_parent_span_id,
447            name,
448            kind,
449        );
450
451        let span_id = span.span_id.clone();
452        {
453            let mut active_spans = self.active_spans.write_safe("active spans for new span").ok()?;
454            active_spans.insert(span_id.clone(), span);
455        }
456
457        // Update current context with new span
458        if let Some(context) = DMSCTracingContext::current() {
459            let new_context = context.new_child(span_id.clone());
460            new_context.set_as_current();
461        } else {
462            // Create new context if none exists
463            let context = DMSCTracingContext::new()
464                .with_trace_id(resolved_trace_id)
465                .with_span_id(span_id.clone());
466            context.set_as_current();
467        }
468
469        Some(span_id)
470    }
471
472    /// Start a new span using current context
473    pub fn start_span_from_context(&self, name: String, kind: DMSCSpanKind) -> Option<DMSCSpanId> {
474        self.start_span(None, None, name, kind)
475    }
476
477    /// End a span and restore parent span context if available
478    pub fn end_span(&self, span_id: &DMSCSpanId, status: DMSCSpanStatus) -> DMSCResult<()> {
479        let mut active_spans = self.active_spans.write_safe("active spans for end span")?;
480
481        if let Some(mut span) = active_spans.remove(span_id) {
482            span.end(status);
483
484            let trace_id = span.trace_id.clone();
485            let parent_span_id = span.parent_span_id.clone();
486            drop(active_spans);
487
488            {
489                let mut spans = self.spans.write_safe("spans for end span")?;
490                if let Some(spans_list) = spans.get_mut(&trace_id) {
491                    spans_list.push(span);
492                }
493            }
494
495            // Restore parent span context if available
496            if let Some(parent_span_id) = parent_span_id {
497                // Try to find parent span in active spans
498                let active_spans = self.active_spans.read_safe("active spans for parent check")?;
499                if active_spans.get(&parent_span_id).is_some() {
500                    let context = DMSCTracingContext::new()
501                        .with_trace_id(trace_id)
502                        .with_span_id(parent_span_id);
503                    context.set_as_current();
504                }
505            } else {
506                // No parent span, clear context
507                let context = DMSCTracingContext::new();
508                context.set_as_current();
509            }
510        }
511
512        Ok(())
513    }
514
515    /// Get span for modification
516    pub fn span_mut<F>(&self, span_id: &DMSCSpanId, f: F) -> DMSCResult<()>
517    where
518        F: FnOnce(&mut DMSCSpan),
519    {
520        let mut active_spans = self.active_spans.write_safe("active spans for span_mut")?;
521
522        if let Some(span) = active_spans.get_mut(span_id) {
523            f(span);
524            Ok(())
525        } else {
526            Err(crate::core::DMSCError::Other("Span not found".to_string()))
527        }
528    }
529
530    /// Export completed traces
531    pub fn export_traces(&self) -> HashMap<DMSCTraceId, Vec<DMSCSpan>> {
532        match self.spans.read_safe("spans for export") {
533            Ok(spans) => spans.clone(),
534            Err(_) => HashMap::new(),
535        }
536    }
537
538    /// Get active traces count
539    pub fn active_trace_count(&self) -> usize {
540        match self.spans.read_safe("spans for count") {
541            Ok(spans) => spans.len(),
542            Err(_) => 0,
543        }
544    }
545
546    /// Get active span count
547    pub fn active_span_count(&self) -> usize {
548        match self.active_spans.read_safe("active spans for count") {
549            Ok(active_spans) => active_spans.len(),
550            Err(_) => 0,
551        }
552    }
553
554    fn should_sample(&self) -> bool {
555        match &self.sampling_strategy {
556            DMSCSamplingStrategy::Rate(rate) => {
557                if *rate >= 1.0 {
558                    true
559                } else if *rate <= 0.0 {
560                    false
561                } else {
562                    use rand::Rng;
563                    let mut rng = rand::thread_rng();
564                    rng.gen::<f64>() < *rate
565                }
566            }
567            DMSCSamplingStrategy::Deterministic(rate) => {
568                if *rate >= 1.0 {
569                    true
570                } else if *rate <= 0.0 {
571                    false
572                } else {
573                    // Create a deterministic hash based on current time and thread ID
574                    let now = SystemTime::now()
575                        .duration_since(UNIX_EPOCH)
576                        .unwrap_or(Duration::from_secs(0))
577                        .as_nanos();
578                    // Get a numeric representation of the thread ID using hash
579                    let thread_id = format!("{:?}", std::thread::current().id())
580                        .as_bytes()
581                        .iter()
582                        .fold(0u64, |acc, &b| acc.wrapping_mul(31).wrapping_add(b as u64));
583                    let combined = now.wrapping_add(thread_id as u128);
584                    
585                    // Simple hash function
586                    let hash = (combined as u64).wrapping_mul(0x517cc1b727220a95);
587                    let hash_f64 = (hash as f64) / (u64::MAX as f64);
588                    
589                    hash_f64 < *rate
590                }
591            }
592            DMSCSamplingStrategy::Adaptive(target_rate) => {
593                if *target_rate >= 1.0 {
594                    true
595                } else if *target_rate <= 0.0 {
596                    false
597                } else {
598                    // Calculate current load based on active spans
599                    let active_count = match self.active_spans.read_safe("active spans for sampling") {
600                        Ok(active_spans) => active_spans.len() as f64,
601                        Err(_) => 0.0,
602                    };
603                    
604                    let mut window = match self.adaptive_window.write_safe("adaptive window for sampling") {
605                        Ok(w) => w,
606                        Err(_) => {
607                            // If we can't acquire the lock, default to high load (low sampling rate)
608                            return false;
609                        }
610                    };
611                    
612                    // Add current active count to window
613                    window.push(active_count as u64);
614                    if window.len() > self.max_adaptive_window {
615                        window.remove(0);
616                    }
617                    
618                    // Calculate average load over window
619                    let avg_load = if window.is_empty() {
620                        0.0
621                    } else {
622                        window.iter().sum::<u64>() as f64 / window.len() as f64
623                    };
624                    
625                    // Adaptive sampling: lower rate when load is high, higher when load is low
626                    const BASE_LOAD: f64 = 100.0;
627                    let adjusted_rate = target_rate * (1.0 + (BASE_LOAD - avg_load) / BASE_LOAD);
628                    let clamped_rate = adjusted_rate.clamp(0.01, 1.0);
629                    
630                    use rand::Rng;
631                    let mut rng = rand::thread_rng();
632                    rng.gen::<f64>() < clamped_rate
633                }
634            }
635        }
636    }
637
638}
639
640#[cfg(feature = "pyo3")]
641#[pyo3::prelude::pymethods]
642impl DMSCTracer {
643    /// Create a new tracer from Python with a sampling rate
644    #[new]
645    fn py_new(sampling_rate: f64) -> Self {
646        Self::new(sampling_rate)
647    }
648
649    /// Start a new trace from Python
650    #[pyo3(name = "start_trace")]
651    fn start_trace_impl(&self, name: String) -> PyResult<Option<String>> {
652        match self.start_trace(name) {
653            Some(trace_id) => Ok(Some(trace_id.as_str().to_string())),
654            None => Ok(None),
655        }
656    }
657
658    /// Start a new span from Python using current context
659    #[pyo3(name = "start_span_from_context")]
660    fn start_span_from_context_impl(&self, name: String, kind: String) -> PyResult<Option<String>> {
661        let span_kind = match kind.as_str() {
662            "Server" => DMSCSpanKind::Server,
663            "Client" => DMSCSpanKind::Client,
664            "Producer" => DMSCSpanKind::Producer,
665            "Consumer" => DMSCSpanKind::Consumer,
666            _ => DMSCSpanKind::Internal,
667        };
668
669        match self.start_span_from_context(name, span_kind) {
670            Some(span_id) => Ok(Some(span_id.as_str().to_string())),
671            None => Ok(None),
672        }
673    }
674
675    /// End a span from Python
676    #[pyo3(name = "end_span")]
677    fn end_span_impl(&self, span_id: String, status: String) -> PyResult<()> {
678        let span_id_obj = DMSCSpanId::from_string(span_id);
679        let span_status = match status.as_str() {
680            "Ok" => DMSCSpanStatus::Ok,
681            "Error" => DMSCSpanStatus::Error("Python error".to_string()),
682            _ => DMSCSpanStatus::Unset,
683        };
684
685        self.end_span(&span_id_obj, span_status)
686            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to end span: {e}")))
687    }
688
689    /// Set span attribute from Python
690    #[pyo3(name = "span_set_attribute")]
691    fn span_set_attribute_impl(&self, span_id: String, key: String, value: String) -> PyResult<()> {
692        let span_id_obj = DMSCSpanId::from_string(span_id);
693        self.span_mut(&span_id_obj, |span| {
694            span.set_attribute(key, value);
695        })
696        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to set attribute: {e}")))
697    }
698
699    /// Add span event from Python
700    #[pyo3(name = "span_add_event")]
701    fn span_add_event_impl(&self, span_id: String, name: String, attributes: HashMap<String, String>) -> PyResult<()> {
702        let span_id_obj = DMSCSpanId::from_string(span_id);
703        self.span_mut(&span_id_obj, |span| {
704            span.add_event(name, attributes);
705        })
706        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to add event: {e}")))
707    }
708
709    /// Export traces from Python
710    #[pyo3(name = "export_traces")]
711    fn export_traces_impl(&self, py: pyo3::Python<'_>) -> PyResult<HashMap<String, Vec<pyo3::Py<pyo3::PyAny>>>> {
712        let traces = self.export_traces();
713        let mut result = HashMap::with_capacity(traces.len());
714
715        for (trace_id, spans) in traces {
716            let mut span_list = Vec::with_capacity(spans.len());
717            for span in spans {
718                let span_dict = pyo3::types::PyDict::new(py);
719                span_dict.set_item("trace_id", span.trace_id.as_str())?;
720                span_dict.set_item("span_id", span.span_id.as_str())?;
721                if let Some(parent_id) = &span.parent_span_id {
722                    span_dict.set_item("parent_span_id", parent_id.as_str())?;
723                }
724                span_dict.set_item("name", &span.name)?;
725                span_dict.set_item("kind", format!("{:?}", span.kind))?;
726                span_dict.set_item("start_time", span.start_time)?;
727                span_dict.set_item("end_time", span.end_time)?;
728                span_dict.set_item("attributes", span.attributes.clone())?;
729                span_dict.set_item("events", span.events.len())?;
730                span_dict.set_item("status", format!("{:?}", span.status))?;
731                span_list.push(span_dict.into());
732            }
733            result.insert(trace_id.as_str().to_string(), span_list);
734        }
735        Ok(result)
736    }
737
738    /// Get active trace count from Python
739    #[pyo3(name = "active_trace_count")]
740    fn active_trace_count_impl(&self) -> usize {
741        self.active_trace_count()
742    }
743
744    /// Get active span count from Python
745    #[pyo3(name = "active_span_count")]
746    fn active_span_count_impl(&self) -> usize {
747        self.active_span_count()
748    }
749}
750
751/// Tracer manager for managing multiple tracer instances
752pub struct DMSCTracerManager {
753    tracers: HashMap<String, Arc<DMSCTracer>>,
754    default_tracer: Option<String>,
755}
756
757impl Default for DMSCTracerManager {
758    fn default() -> Self {
759        Self::new()
760    }
761}
762
763impl DMSCTracerManager {
764    pub fn new() -> Self {
765        Self {
766            tracers: HashMap::new(),
767            default_tracer: None,
768        }
769    }
770
771    pub fn register_tracer(&mut self, name: &str, tracer: Arc<DMSCTracer>) {
772        self.tracers.insert(name.to_string(), tracer);
773        if self.default_tracer.is_none() {
774            self.default_tracer = Some(name.to_string());
775        }
776    }
777
778    #[allow(dead_code)]
779    pub fn get_tracer(&self, name: &str) -> Option<&Arc<DMSCTracer>> {
780        self.tracers.get(name)
781    }
782
783    pub fn get_default_tracer(&self) -> Option<&Arc<DMSCTracer>> {
784        if let Some(default_name) = &self.default_tracer {
785            self.tracers.get(default_name)
786        } else {
787            None
788        }
789    }
790
791    #[allow(dead_code)]
792    pub fn set_default_tracer(&mut self, name: &str) -> bool {
793        if self.tracers.contains_key(name) {
794            self.default_tracer = Some(name.to_string());
795            true
796        } else {
797            false
798        }
799    }
800
801    #[allow(dead_code)]
802    pub fn remove_tracer(&mut self, name: &str) -> bool {
803        let removed = self.tracers.remove(name).is_some();
804        if let Some(default_name) = &self.default_tracer {
805            if default_name == name {
806                self.default_tracer = None;
807            }
808        }
809        removed
810    }
811}
812
813/// Default tracer manager instance
814pub struct DefaultTracerManager {
815    inner: Arc<RwLock<DMSCTracerManager>>,
816}
817
818impl Default for DefaultTracerManager {
819    fn default() -> Self {
820        Self {
821            inner: Arc::new(RwLock::new(DMSCTracerManager::new())),
822        }
823    }
824}
825
826impl DefaultTracerManager {
827    #[allow(dead_code)]
828    pub fn new() -> Self {
829        Default::default()
830    }
831
832    pub async fn register_tracer(&self, name: &str, sampling_rate: f64) -> DMSCResult<()> {
833        let tracer = Arc::new(DMSCTracer::new(sampling_rate));
834        let mut manager = self.inner.write_safe("tracer manager for register")?;
835        manager.register_tracer(name, tracer);
836        Ok(())
837    }
838    
839    pub async fn register_tracer_with_strategy(&self, name: &str, strategy: DMSCSamplingStrategy) -> DMSCResult<()> {
840        let tracer = Arc::new(DMSCTracer::with_strategy(strategy));
841        let mut manager = self.inner.write_safe("tracer manager for register with strategy")?;
842        manager.register_tracer(name, tracer);
843        Ok(())
844    }
845
846    #[allow(dead_code)]
847    pub async fn get_tracer(&self, name: &str) -> DMSCResult<Option<Arc<DMSCTracer>>> {
848        let manager = self.inner.read_safe("tracer manager for get")?;
849        Ok(manager.get_tracer(name).cloned())
850    }
851
852    pub async fn get_default_tracer(&self) -> DMSCResult<Option<Arc<DMSCTracer>>> {
853        let manager = self.inner.read_safe("tracer manager for get default")?;
854        Ok(manager.get_default_tracer().cloned())
855    }
856
857    #[allow(dead_code)]
858    pub async fn set_default_tracer(&self, name: &str) -> DMSCResult<bool> {
859        let mut manager = self.inner.write_safe("tracer manager for set default")?;
860        Ok(manager.set_default_tracer(name))
861    }
862
863    #[allow(dead_code)]
864    pub async fn remove_tracer(&self, name: &str) -> DMSCResult<bool> {
865        let mut manager = self.inner.write_safe("tracer manager for remove")?;
866        Ok(manager.remove_tracer(name))
867    }
868}
869
870/// Global tracer manager instance
871pub static DEFAULT_TRACER_MANAGER: std::sync::LazyLock<DefaultTracerManager> = std::sync::LazyLock::new(DefaultTracerManager::default);
872
873/// Initialize global tracer with fixed rate (backward compatibility)
874pub fn init_tracer(sampling_rate: f64) {
875    let runtime = match tokio::runtime::Builder::new_current_thread()
876        .enable_all()
877        .build()
878    {
879        Ok(r) => r,
880        Err(e) => {
881            eprintln!("Failed to create tokio runtime: {}", e);
882            return;
883        }
884    };
885    
886    runtime.block_on(async {
887        if let Err(e) = DEFAULT_TRACER_MANAGER.register_tracer("default", sampling_rate).await {
888            eprintln!("Failed to register tracer: {}", e);
889        }
890    });
891}
892
893/// Initialize global tracer with custom sampling strategy
894pub fn init_tracer_with_strategy(strategy: DMSCSamplingStrategy) {
895    let rate = match strategy {
896        DMSCSamplingStrategy::Rate(rate) => rate,
897        DMSCSamplingStrategy::Deterministic(rate) => rate,
898        DMSCSamplingStrategy::Adaptive(rate) => rate,
899    };
900    
901    let runtime = match tokio::runtime::Builder::new_current_thread()
902        .enable_all()
903        .build()
904    {
905        Ok(r) => r,
906        Err(e) => {
907            eprintln!("Failed to create tokio runtime: {}", e);
908            return;
909        }
910    };
911    
912    runtime.block_on(async {
913        if let Err(e) = DEFAULT_TRACER_MANAGER.register_tracer("default", rate).await {
914            eprintln!("Failed to register tracer: {}", e);
915        }
916    });
917}
918
919/// Get global tracer (backward compatibility)
920pub fn tracer() -> Result<Arc<DMSCTracer>, Box<DMSCError>> {
921    let runtime = match tokio::runtime::Builder::new_current_thread()
922        .enable_all()
923        .build()
924    {
925        Ok(r) => r,
926        Err(e) => {
927            return Err(Box::new(DMSCError::Other(format!(
928                "Failed to create tokio runtime for tracer: {}",
929                e
930            ))));
931        }
932    };
933
934    runtime.block_on(async {
935        match DEFAULT_TRACER_MANAGER.get_default_tracer().await {
936            Ok(Some(tracer)) => Ok(tracer),
937            Ok(None) => {
938                Err(Box::new(DMSCError::Other(
939                    "Tracer not initialized".to_string(),
940                )))
941            }
942            Err(e) => Err(Box::new(DMSCError::Other(format!(
943                "Failed to get tracer: {}",
944                e
945            )))),
946        }
947    })
948}
dmsc/observability/tracing.rs

dmsc/observability/
tracing.rs
