dcp/observability/
tracing.rs

1//! OpenTelemetry tracing integration for DCP server.
2//!
3//! Provides distributed tracing with request ID propagation.
4
5use std::collections::HashMap;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
8
9use crate::security::{
10    is_sensitive_key, sanitize_field_key, sanitize_field_value, sanitize_text, REDACTED,
11};
12
13/// Tracing configuration
14#[derive(Debug, Clone)]
15pub struct TracingConfig {
16    /// Service name for tracing
17    pub service_name: String,
18    /// Enable tracing
19    pub enabled: bool,
20    /// Sample rate (0.0 to 1.0)
21    pub sample_rate: f64,
22    /// OTLP endpoint (if using OpenTelemetry)
23    pub otlp_endpoint: Option<String>,
24    /// Include request/response bodies in spans
25    pub include_bodies: bool,
26}
27
28impl Default for TracingConfig {
29    fn default() -> Self {
30        Self {
31            service_name: "dcp-server".to_string(),
32            enabled: true,
33            sample_rate: 1.0,
34            otlp_endpoint: None,
35            include_bodies: false,
36        }
37    }
38}
39
40/// Request ID generator
41#[derive(Debug, Default)]
42pub struct RequestIdGenerator {
43    counter: AtomicU64,
44    node_id: u16,
45}
46
47impl RequestIdGenerator {
48    /// Create a new request ID generator
49    pub fn new() -> Self {
50        Self::with_node_id(0)
51    }
52
53    /// Create with a specific node ID for distributed systems
54    pub fn with_node_id(node_id: u16) -> Self {
55        Self {
56            counter: AtomicU64::new(0),
57            node_id,
58        }
59    }
60
61    /// Generate a new request ID
62    /// Format: {timestamp_ms}-{node_id}-{counter}
63    pub fn generate(&self) -> String {
64        let timestamp = SystemTime::now()
65            .duration_since(UNIX_EPOCH)
66            .unwrap_or_default()
67            .as_millis() as u64;
68        let counter = self.counter.fetch_add(1, Ordering::Relaxed);
69        format!("{:x}-{:04x}-{:08x}", timestamp, self.node_id, counter)
70    }
71}
72
73/// Span status
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
75pub enum SpanStatus {
76    /// Operation completed successfully
77    Ok,
78    /// Operation failed with an error
79    Error,
80    /// Operation was cancelled
81    Cancelled,
82}
83
84/// Span kind
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86pub enum SpanKind {
87    /// Internal operation
88    Internal,
89    /// Server handling a request
90    Server,
91    /// Client making a request
92    Client,
93    /// Producer sending a message
94    Producer,
95    /// Consumer receiving a message
96    Consumer,
97}
98
99/// A tracing span representing a unit of work
100#[derive(Debug)]
101pub struct Span {
102    /// Span name
103    pub name: String,
104    /// Trace ID
105    pub trace_id: String,
106    /// Span ID
107    pub span_id: String,
108    /// Parent span ID
109    pub parent_span_id: Option<String>,
110    /// Span kind
111    pub kind: SpanKind,
112    /// Start time
113    pub start_time: Instant,
114    /// End time (set when span ends)
115    pub end_time: Option<Instant>,
116    /// Span status
117    pub status: SpanStatus,
118    /// Attributes
119    pub attributes: HashMap<String, SpanValue>,
120    /// Events
121    pub events: Vec<SpanEvent>,
122}
123
124/// Span attribute value
125#[derive(Debug, Clone)]
126pub enum SpanValue {
127    String(String),
128    Int(i64),
129    Float(f64),
130    Bool(bool),
131}
132
133impl From<&str> for SpanValue {
134    fn from(s: &str) -> Self {
135        SpanValue::String(s.to_string())
136    }
137}
138
139impl From<String> for SpanValue {
140    fn from(s: String) -> Self {
141        SpanValue::String(s)
142    }
143}
144
145impl From<i64> for SpanValue {
146    fn from(v: i64) -> Self {
147        SpanValue::Int(v)
148    }
149}
150
151impl From<f64> for SpanValue {
152    fn from(v: f64) -> Self {
153        SpanValue::Float(v)
154    }
155}
156
157impl From<bool> for SpanValue {
158    fn from(v: bool) -> Self {
159        SpanValue::Bool(v)
160    }
161}
162
163/// Span event
164#[derive(Debug, Clone)]
165pub struct SpanEvent {
166    /// Event name
167    pub name: String,
168    /// Event timestamp
169    pub timestamp: Instant,
170    /// Event attributes
171    pub attributes: HashMap<String, SpanValue>,
172}
173
174impl Span {
175    /// Create a new span
176    pub fn new(name: impl Into<String>, trace_id: String, span_id: String) -> Self {
177        Self {
178            name: sanitize_text(&name.into()),
179            trace_id,
180            span_id,
181            parent_span_id: None,
182            kind: SpanKind::Internal,
183            start_time: Instant::now(),
184            end_time: None,
185            status: SpanStatus::Ok,
186            attributes: HashMap::new(),
187            events: Vec::new(),
188        }
189    }
190
191    /// Set parent span ID
192    pub fn with_parent(mut self, parent_span_id: String) -> Self {
193        self.parent_span_id = Some(parent_span_id);
194        self
195    }
196
197    /// Set span kind
198    pub fn with_kind(mut self, kind: SpanKind) -> Self {
199        self.kind = kind;
200        self
201    }
202
203    /// Set an attribute
204    pub fn set_attribute(&mut self, key: impl Into<String>, value: impl Into<SpanValue>) {
205        let key = key.into();
206        let value = sanitize_span_value(&key, value.into());
207        self.attributes.insert(sanitize_field_key(&key), value);
208    }
209
210    /// Add an event
211    pub fn add_event(&mut self, name: impl Into<String>) {
212        self.events.push(SpanEvent {
213            name: sanitize_text(&name.into()),
214            timestamp: Instant::now(),
215            attributes: HashMap::new(),
216        });
217    }
218
219    /// Add an event with attributes
220    pub fn add_event_with_attributes(
221        &mut self,
222        name: impl Into<String>,
223        attributes: HashMap<String, SpanValue>,
224    ) {
225        let attributes = attributes
226            .into_iter()
227            .map(|(key, value)| {
228                let value = sanitize_span_value(&key, value);
229                (sanitize_field_key(&key), value)
230            })
231            .collect();
232
233        self.events.push(SpanEvent {
234            name: sanitize_text(&name.into()),
235            timestamp: Instant::now(),
236            attributes,
237        });
238    }
239
240    /// Set span status
241    pub fn set_status(&mut self, status: SpanStatus) {
242        self.status = status;
243    }
244
245    /// End the span
246    pub fn end(&mut self) {
247        self.end_time = Some(Instant::now());
248    }
249
250    /// Get span duration
251    pub fn duration(&self) -> Option<Duration> {
252        self.end_time.map(|end| end.duration_since(self.start_time))
253    }
254}
255
256/// Span ID generator
257fn generate_span_id() -> String {
258    use std::sync::atomic::AtomicU64;
259    static COUNTER: AtomicU64 = AtomicU64::new(0);
260    let id = COUNTER.fetch_add(1, Ordering::Relaxed);
261    let random = std::time::SystemTime::now()
262        .duration_since(UNIX_EPOCH)
263        .unwrap_or_default()
264        .as_nanos() as u64;
265    format!("{:016x}", id ^ random)
266}
267
268/// Trace ID generator
269fn generate_trace_id() -> String {
270    let timestamp = SystemTime::now()
271        .duration_since(UNIX_EPOCH)
272        .unwrap_or_default()
273        .as_nanos() as u64;
274    let random = std::time::SystemTime::now()
275        .duration_since(UNIX_EPOCH)
276        .unwrap_or_default()
277        .subsec_nanos() as u64;
278    format!("{:016x}{:016x}", timestamp, random)
279}
280
281/// Create a new span
282pub fn create_span(name: impl Into<String>) -> Span {
283    Span::new(name, generate_trace_id(), generate_span_id())
284}
285
286/// Create a child span
287pub fn create_child_span(name: impl Into<String>, parent: &Span) -> Span {
288    Span::new(name, parent.trace_id.clone(), generate_span_id()).with_parent(parent.span_id.clone())
289}
290
291/// Tracer for creating and managing spans
292#[derive(Debug)]
293pub struct Tracer {
294    config: TracingConfig,
295    request_id_generator: RequestIdGenerator,
296}
297
298impl Tracer {
299    /// Create a new tracer
300    pub fn new(config: TracingConfig) -> Self {
301        Self {
302            config,
303            request_id_generator: RequestIdGenerator::new(),
304        }
305    }
306
307    /// Create with default configuration
308    pub fn with_defaults() -> Self {
309        Self::new(TracingConfig::default())
310    }
311
312    /// Generate a request ID
313    pub fn generate_request_id(&self) -> String {
314        self.request_id_generator.generate()
315    }
316
317    /// Create a new root span
318    pub fn start_span(&self, name: impl Into<String>) -> Span {
319        create_span(name)
320    }
321
322    /// Create a child span
323    pub fn start_child_span(&self, name: impl Into<String>, parent: &Span) -> Span {
324        create_child_span(name, parent)
325    }
326
327    /// Check if tracing is enabled
328    pub fn is_enabled(&self) -> bool {
329        self.config.enabled
330    }
331
332    /// Get configuration
333    pub fn config(&self) -> &TracingConfig {
334        &self.config
335    }
336}
337
338impl Default for Tracer {
339    fn default() -> Self {
340        Self::with_defaults()
341    }
342}
343
344/// Request span for tracking HTTP/RPC requests
345pub struct RequestSpan {
346    span: Span,
347    request_id: String,
348}
349
350impl RequestSpan {
351    /// Create a new request span
352    pub fn new(method: impl Into<String>, request_id: String) -> Self {
353        let method = method.into();
354        let mut span = create_span(&method);
355        span.set_attribute("rpc.method", method);
356        span.set_attribute("request.id", request_id.clone());
357        span.kind = SpanKind::Server;
358        Self { span, request_id }
359    }
360
361    /// Get the request ID
362    pub fn request_id(&self) -> &str {
363        &self.request_id
364    }
365
366    /// Get the trace ID
367    pub fn trace_id(&self) -> &str {
368        &self.span.trace_id
369    }
370
371    /// Get the span ID
372    pub fn span_id(&self) -> &str {
373        &self.span.span_id
374    }
375
376    /// Set an attribute
377    pub fn set_attribute(&mut self, key: impl Into<String>, value: impl Into<SpanValue>) {
378        self.span.set_attribute(key, value);
379    }
380
381    /// Add an event
382    pub fn add_event(&mut self, name: impl Into<String>) {
383        self.span.add_event(name);
384    }
385
386    /// Set status to error
387    pub fn set_error(&mut self, error_message: impl Into<String>) {
388        self.span.set_status(SpanStatus::Error);
389        let error_message = error_message.into();
390        self.span
391            .set_attribute("error.message", sanitize_text(&error_message));
392    }
393
394    /// Finish the span successfully
395    pub fn finish(mut self) -> Span {
396        self.span.set_status(SpanStatus::Ok);
397        self.span.end();
398        self.span
399    }
400
401    /// Finish the span with an error
402    pub fn finish_with_error(mut self, error: impl Into<String>) -> Span {
403        self.span.set_status(SpanStatus::Error);
404        let error = error.into();
405        self.span
406            .set_attribute("error.message", sanitize_text(&error));
407        self.span.end();
408        self.span
409    }
410
411    /// Get the underlying span
412    pub fn span(&self) -> &Span {
413        &self.span
414    }
415
416    /// Get mutable access to the underlying span
417    pub fn span_mut(&mut self) -> &mut Span {
418        &mut self.span
419    }
420}
421
422fn sanitize_span_value(key: &str, value: SpanValue) -> SpanValue {
423    match value {
424        SpanValue::String(value) => SpanValue::String(sanitize_field_value(key, &value)),
425        other => {
426            if is_sensitive_key(key) {
427                SpanValue::String(REDACTED.to_string())
428            } else {
429                other
430            }
431        }
432    }
433}
434
435/// Initialize tracing with the given configuration
436pub fn init_tracing(config: TracingConfig) -> Tracer {
437    Tracer::new(config)
438}
439
440/// Span collector for testing and debugging
441#[derive(Debug, Default)]
442pub struct SpanCollector {
443    spans: std::sync::RwLock<Vec<Span>>,
444}
445
446impl SpanCollector {
447    /// Create a new span collector
448    pub fn new() -> Self {
449        Self::default()
450    }
451
452    /// Collect a span
453    pub fn collect(&self, span: Span) {
454        self.spans.write().unwrap().push(span);
455    }
456
457    /// Get all collected spans
458    pub fn spans(&self) -> Vec<Span> {
459        self.spans.read().unwrap().clone()
460    }
461
462    /// Clear collected spans
463    pub fn clear(&self) {
464        self.spans.write().unwrap().clear();
465    }
466
467    /// Get span count
468    pub fn len(&self) -> usize {
469        self.spans.read().unwrap().len()
470    }
471
472    /// Check if empty
473    pub fn is_empty(&self) -> bool {
474        self.spans.read().unwrap().is_empty()
475    }
476}
477
478impl Clone for Span {
479    fn clone(&self) -> Self {
480        Self {
481            name: self.name.clone(),
482            trace_id: self.trace_id.clone(),
483            span_id: self.span_id.clone(),
484            parent_span_id: self.parent_span_id.clone(),
485            kind: self.kind,
486            start_time: self.start_time,
487            end_time: self.end_time,
488            status: self.status,
489            attributes: self.attributes.clone(),
490            events: self.events.clone(),
491        }
492    }
493}
494
495#[cfg(test)]
496mod tests {
497    use super::*;
498
499    #[test]
500    fn test_request_id_generation() {
501        let generator = RequestIdGenerator::new();
502        let id1 = generator.generate();
503        let id2 = generator.generate();
504
505        assert_ne!(id1, id2);
506        assert!(id1.contains('-'));
507        assert!(id2.contains('-'));
508    }
509
510    #[test]
511    fn test_span_creation() {
512        let span = create_span("test_operation");
513
514        assert_eq!(span.name, "test_operation");
515        assert!(!span.trace_id.is_empty());
516        assert!(!span.span_id.is_empty());
517        assert!(span.parent_span_id.is_none());
518        assert_eq!(span.status, SpanStatus::Ok);
519    }
520
521    #[test]
522    fn test_child_span() {
523        let parent = create_span("parent");
524        let child = create_child_span("child", &parent);
525
526        assert_eq!(child.trace_id, parent.trace_id);
527        assert_eq!(child.parent_span_id, Some(parent.span_id.clone()));
528        assert_ne!(child.span_id, parent.span_id);
529    }
530
531    #[test]
532    fn test_span_attributes() {
533        let mut span = create_span("test");
534        span.set_attribute("key1", "value1");
535        span.set_attribute("key2", 42i64);
536        span.set_attribute("key3", 3.14f64);
537        span.set_attribute("key4", true);
538
539        assert_eq!(span.attributes.len(), 4);
540    }
541
542    #[test]
543    fn test_span_events() {
544        let mut span = create_span("test");
545        span.add_event("event1");
546        span.add_event("event2");
547
548        assert_eq!(span.events.len(), 2);
549        assert_eq!(span.events[0].name, "event1");
550        assert_eq!(span.events[1].name, "event2");
551    }
552
553    #[test]
554    fn test_span_duration() {
555        let mut span = create_span("test");
556        std::thread::sleep(std::time::Duration::from_millis(1));
557        span.end();
558
559        let duration = span.duration().unwrap();
560        assert!(duration.as_micros() > 0);
561    }
562
563    #[test]
564    fn test_request_span() {
565        let request_id = "test-request-123".to_string();
566        let mut span = RequestSpan::new("tools/call", request_id.clone());
567
568        assert_eq!(span.request_id(), "test-request-123");
569        span.set_attribute("tool.name", "my_tool");
570
571        let finished = span.finish();
572        assert_eq!(finished.status, SpanStatus::Ok);
573        assert!(finished.end_time.is_some());
574    }
575
576    #[test]
577    fn test_tracer() {
578        let tracer = Tracer::with_defaults();
579
580        let request_id = tracer.generate_request_id();
581        assert!(!request_id.is_empty());
582
583        let span = tracer.start_span("test");
584        assert_eq!(span.name, "test");
585    }
586
587    #[test]
588    fn test_span_collector() {
589        let collector = SpanCollector::new();
590
591        let mut span1 = create_span("span1");
592        span1.end();
593        collector.collect(span1);
594
595        let mut span2 = create_span("span2");
596        span2.end();
597        collector.collect(span2);
598
599        assert_eq!(collector.len(), 2);
600
601        let spans = collector.spans();
602        assert_eq!(spans[0].name, "span1");
603        assert_eq!(spans[1].name, "span2");
604    }
605}
dcp/observability/tracing.rs

dcp/observability/
tracing.rs
