kindly_guard_server/telemetry/
distributed.rs

1// Copyright 2025 Kindly Software Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//! Distributed tracing support for `KindlyGuard`
15//!
16//! Provides context propagation, span relationships, and distributed tracing
17//! capabilities for tracking operations across system boundaries.
18
19use anyhow::Result;
20use async_trait::async_trait;
21use serde::{Deserialize, Serialize};
22use std::collections::HashMap;
23use std::sync::Arc;
24use std::time::Instant;
25use tokio::sync::RwLock;
26
27use super::{MetricValue, TelemetryContext, TelemetryMetric, TelemetryProvider};
28
29/// Distributed tracing span with full context
30#[derive(Debug, Clone)]
31pub struct DistributedSpan {
32    /// Unique span ID
33    pub span_id: String,
34
35    /// Trace ID for correlation
36    pub trace_id: String,
37
38    /// Parent span ID (if any)
39    pub parent_span_id: Option<String>,
40
41    /// Operation name
42    pub operation_name: String,
43
44    /// Start time
45    pub start_time: Instant,
46
47    /// End time (when completed)
48    pub end_time: Option<Instant>,
49
50    /// Span kind
51    pub kind: SpanKind,
52
53    /// Status
54    pub status: SpanStatus,
55
56    /// Attributes
57    pub attributes: HashMap<String, String>,
58
59    /// Events within the span
60    pub events: Vec<SpanEvent>,
61
62    /// Links to other spans
63    pub links: Vec<SpanLink>,
64}
65
66/// Span kinds following OpenTelemetry spec
67#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
68pub enum SpanKind {
69    /// Default span kind
70    Internal,
71
72    /// Span represents handling of a request
73    Server,
74
75    /// Span represents making a request
76    Client,
77
78    /// Span represents a producer of messages
79    Producer,
80
81    /// Span represents a consumer of messages
82    Consumer,
83}
84
85/// Span status
86#[derive(Debug, Clone, Serialize, Deserialize)]
87pub struct SpanStatus {
88    pub code: StatusCode,
89    pub description: Option<String>,
90}
91
92/// Status codes
93#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
94pub enum StatusCode {
95    /// Operation completed successfully
96    Ok,
97
98    /// Operation encountered an error
99    Error,
100
101    /// Status not set
102    Unset,
103}
104
105/// Event within a span
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct SpanEvent {
108    pub name: String,
109    #[serde(skip, default = "Instant::now")]
110    pub timestamp: Instant,
111    pub timestamp_utc: chrono::DateTime<chrono::Utc>,
112    pub attributes: HashMap<String, String>,
113}
114
115/// Link to another span
116#[derive(Debug, Clone, Serialize, Deserialize)]
117pub struct SpanLink {
118    pub trace_id: String,
119    pub span_id: String,
120    pub attributes: HashMap<String, String>,
121}
122
123/// Distributed tracing provider
124pub struct DistributedTracingProvider {
125    /// Base telemetry provider
126    base_provider: Arc<dyn TelemetryProvider>,
127
128    /// Active spans
129    active_spans: Arc<RwLock<HashMap<String, DistributedSpan>>>,
130
131    /// Completed spans buffer
132    completed_spans: Arc<RwLock<Vec<DistributedSpan>>>,
133
134    /// Sampling strategy
135    sampler: Arc<dyn TracingSampler>,
136
137    /// Context propagator
138    #[allow(dead_code)] // Reserved for distributed context propagation
139    propagator: Arc<dyn ContextPropagator>,
140}
141
142/// Sampling decision
143#[derive(Debug, Clone)]
144pub struct SamplingDecision {
145    pub sampled: bool,
146    pub attributes: Option<SamplingAttributes>,
147}
148
149/// Sampling attributes to add to span
150#[derive(Debug, Clone)]
151pub struct SamplingAttributes {
152    pub sampling_priority: f64,
153    pub sampling_rate: f64,
154}
155
156/// Trait for sampling strategies
157pub trait TracingSampler: Send + Sync {
158    /// Decide whether to sample a span
159    fn should_sample(
160        &self,
161        trace_id: &str,
162        parent_context: Option<&TelemetryContext>,
163        span_name: &str,
164        span_kind: SpanKind,
165        attributes: &[(String, String)],
166    ) -> SamplingDecision;
167}
168
169/// Trait for context propagation
170#[async_trait]
171pub trait ContextPropagator: Send + Sync {
172    /// Extract context from carrier (e.g., HTTP headers)
173    async fn extract(&self, carrier: &dyn ContextCarrier) -> Result<Option<TelemetryContext>>;
174
175    /// Inject context into carrier
176    async fn inject(
177        &self,
178        context: &TelemetryContext,
179        carrier: &mut dyn ContextCarrier,
180    ) -> Result<()>;
181}
182
183/// Carrier for context propagation
184pub trait ContextCarrier: Send + Sync {
185    /// Get value by key
186    fn get(&self, key: &str) -> Option<&str>;
187
188    /// Set value by key
189    fn set(&mut self, key: &str, value: String);
190
191    /// Get all keys
192    fn keys(&self) -> Vec<&str>;
193}
194
195impl DistributedTracingProvider {
196    /// Create a new distributed tracing provider
197    pub fn new(
198        base_provider: Arc<dyn TelemetryProvider>,
199        sampler: Arc<dyn TracingSampler>,
200        propagator: Arc<dyn ContextPropagator>,
201    ) -> Self {
202        Self {
203            base_provider,
204            active_spans: Arc::new(RwLock::new(HashMap::new())),
205            completed_spans: Arc::new(RwLock::new(Vec::with_capacity(1000))),
206            sampler,
207            propagator,
208        }
209    }
210
211    /// Start a new distributed span
212    pub async fn start_distributed_span(
213        &self,
214        operation_name: &str,
215        kind: SpanKind,
216        parent_context: Option<&TelemetryContext>,
217    ) -> DistributedSpan {
218        let (trace_id, parent_span_id) = if let Some(parent) = parent_context {
219            (parent.trace_id.clone(), Some(parent.span_id.clone()))
220        } else {
221            (uuid::Uuid::new_v4().to_string(), None)
222        };
223
224        let span_id = uuid::Uuid::new_v4().to_string();
225
226        // Check sampling decision
227        let sampling_decision =
228            self.sampler
229                .should_sample(&trace_id, parent_context, operation_name, kind, &[]);
230
231        let mut attributes = HashMap::new();
232
233        // Add sampling attributes if provided
234        if let Some(sampling_attrs) = sampling_decision.attributes {
235            attributes.insert(
236                "sampling.priority".to_string(),
237                sampling_attrs.sampling_priority.to_string(),
238            );
239            attributes.insert(
240                "sampling.rate".to_string(),
241                sampling_attrs.sampling_rate.to_string(),
242            );
243        }
244
245        // Add standard attributes
246        attributes.insert("span.kind".to_string(), format!("{kind:?}"));
247        attributes.insert("service.name".to_string(), "kindly-guard".to_string());
248
249        let span = DistributedSpan {
250            span_id: span_id.clone(),
251            trace_id: trace_id.clone(),
252            parent_span_id,
253            operation_name: operation_name.to_string(),
254            start_time: Instant::now(),
255            end_time: None,
256            kind,
257            status: SpanStatus {
258                code: StatusCode::Unset,
259                description: None,
260            },
261            attributes,
262            events: Vec::new(),
263            links: Vec::new(),
264        };
265
266        // Store active span if sampled
267        if sampling_decision.sampled {
268            self.active_spans
269                .write()
270                .await
271                .insert(span_id.clone(), span.clone());
272
273            // Also notify base provider
274            self.base_provider.start_span(operation_name);
275        }
276
277        span
278    }
279
280    /// Add event to span
281    pub async fn add_span_event(
282        &self,
283        span_id: &str,
284        event_name: &str,
285        attributes: Vec<(&str, &str)>,
286    ) {
287        if let Some(span) = self.active_spans.write().await.get_mut(span_id) {
288            let event = SpanEvent {
289                name: event_name.to_string(),
290                timestamp: Instant::now(),
291                timestamp_utc: chrono::Utc::now(),
292                attributes: attributes
293                    .into_iter()
294                    .map(|(k, v)| (k.to_string(), v.to_string()))
295                    .collect(),
296            };
297            span.events.push(event);
298        }
299    }
300
301    /// Add link to span
302    pub async fn add_span_link(
303        &self,
304        span_id: &str,
305        link_trace_id: &str,
306        link_span_id: &str,
307        attributes: Vec<(&str, &str)>,
308    ) {
309        if let Some(span) = self.active_spans.write().await.get_mut(span_id) {
310            let link = SpanLink {
311                trace_id: link_trace_id.to_string(),
312                span_id: link_span_id.to_string(),
313                attributes: attributes
314                    .into_iter()
315                    .map(|(k, v)| (k.to_string(), v.to_string()))
316                    .collect(),
317            };
318            span.links.push(link);
319        }
320    }
321
322    /// End a distributed span
323    pub async fn end_distributed_span(&self, span_id: &str, status: SpanStatus) {
324        if let Some(mut span) = self.active_spans.write().await.remove(span_id) {
325            span.end_time = Some(Instant::now());
326            span.status = status;
327
328            // Calculate duration
329            if let Some(end_time) = span.end_time {
330                let duration = end_time.duration_since(span.start_time);
331
332                // Record span duration metric
333                self.base_provider.record_metric(TelemetryMetric {
334                    name: format!("trace.span.duration.{}", span.operation_name),
335                    value: MetricValue::Histogram(duration.as_secs_f64() * 1000.0),
336                    labels: vec![
337                        ("span.kind".to_string(), format!("{:?}", span.kind)),
338                        ("status.code".to_string(), format!("{:?}", span.status.code)),
339                    ],
340                });
341            }
342
343            // Store completed span
344            let mut completed = self.completed_spans.write().await;
345            completed.push(span.clone());
346
347            // Limit buffer size
348            if completed.len() > 10000 {
349                completed.drain(0..1000);
350            }
351        }
352    }
353
354    /// Get trace context for a span
355    pub async fn get_span_context(&self, span_id: &str) -> Option<TelemetryContext> {
356        self.active_spans
357            .read()
358            .await
359            .get(span_id)
360            .map(|span| TelemetryContext {
361                trace_id: span.trace_id.clone(),
362                span_id: span.span_id.clone(),
363                parent_span_id: span.parent_span_id.clone(),
364                baggage: vec![],
365            })
366    }
367
368    /// Export completed spans (for background processing)
369    pub async fn export_spans(&self) -> Vec<DistributedSpan> {
370        let mut completed = self.completed_spans.write().await;
371        let spans: Vec<_> = completed.drain(..).collect();
372        spans
373    }
374}
375
376/// Probability sampler implementation
377pub struct ProbabilitySampler {
378    sampling_rate: f64,
379}
380
381impl ProbabilitySampler {
382    pub fn new(sampling_rate: f64) -> Self {
383        Self {
384            sampling_rate: sampling_rate.clamp(0.0, 1.0),
385        }
386    }
387}
388
389impl TracingSampler for ProbabilitySampler {
390    fn should_sample(
391        &self,
392        trace_id: &str,
393        parent_context: Option<&TelemetryContext>,
394        _span_name: &str,
395        _span_kind: SpanKind,
396        _attributes: &[(String, String)],
397    ) -> SamplingDecision {
398        // If parent is sampled, always sample
399        if parent_context.is_some() {
400            return SamplingDecision {
401                sampled: true,
402                attributes: Some(SamplingAttributes {
403                    sampling_priority: 1.0,
404                    sampling_rate: self.sampling_rate,
405                }),
406            };
407        }
408
409        // Use trace ID for deterministic sampling
410        // Use a better hash distribution for sequential IDs
411        use std::collections::hash_map::DefaultHasher;
412        use std::hash::{Hash, Hasher};
413
414        let mut hasher = DefaultHasher::new();
415        trace_id.hash(&mut hasher);
416        let hash = hasher.finish();
417
418        let probability = (hash as f64) / (u64::MAX as f64);
419        let sampled = probability < self.sampling_rate;
420
421        SamplingDecision {
422            sampled,
423            attributes: if sampled {
424                Some(SamplingAttributes {
425                    sampling_priority: if sampled { 1.0 } else { 0.0 },
426                    sampling_rate: self.sampling_rate,
427                })
428            } else {
429                None
430            },
431        }
432    }
433}
434
435/// W3C `TraceContext` propagator
436pub struct W3CTraceContextPropagator;
437
438#[async_trait]
439impl ContextPropagator for W3CTraceContextPropagator {
440    async fn extract(&self, carrier: &dyn ContextCarrier) -> Result<Option<TelemetryContext>> {
441        // Extract traceparent header
442        if let Some(traceparent) = carrier.get("traceparent") {
443            // Parse W3C trace context format: version-trace_id-parent_id-trace_flags
444            let parts: Vec<&str> = traceparent.split('-').collect();
445            if parts.len() >= 4 {
446                let trace_id = parts[1].to_string();
447                let span_id = parts[2].to_string();
448
449                // Extract tracestate if present
450                let baggage = if let Some(tracestate) = carrier.get("tracestate") {
451                    tracestate
452                        .split(',')
453                        .filter_map(|kv| {
454                            let parts: Vec<&str> = kv.split('=').collect();
455                            if parts.len() == 2 {
456                                Some((parts[0].to_string(), parts[1].to_string()))
457                            } else {
458                                None
459                            }
460                        })
461                        .collect()
462                } else {
463                    vec![]
464                };
465
466                return Ok(Some(TelemetryContext {
467                    trace_id,
468                    span_id: uuid::Uuid::new_v4().to_string(), // Generate new span ID
469                    parent_span_id: Some(span_id),
470                    baggage,
471                }));
472            }
473        }
474
475        Ok(None)
476    }
477
478    async fn inject(
479        &self,
480        context: &TelemetryContext,
481        carrier: &mut dyn ContextCarrier,
482    ) -> Result<()> {
483        // Create traceparent header
484        let parent_id = context.parent_span_id.as_ref().unwrap_or(&context.span_id);
485        let traceparent = format!("00-{}-{}-01", context.trace_id, parent_id);
486        carrier.set("traceparent", traceparent);
487
488        // Create tracestate header if baggage exists
489        if !context.baggage.is_empty() {
490            let tracestate = context
491                .baggage
492                .iter()
493                .map(|(k, v)| format!("{k}={v}"))
494                .collect::<Vec<_>>()
495                .join(",");
496            carrier.set("tracestate", tracestate);
497        }
498
499        Ok(())
500    }
501}
502
503/// HTTP headers carrier implementation
504pub struct HttpHeadersCarrier {
505    headers: HashMap<String, String>,
506}
507
508impl Default for HttpHeadersCarrier {
509    fn default() -> Self {
510        Self::new()
511    }
512}
513
514impl HttpHeadersCarrier {
515    pub fn new() -> Self {
516        Self {
517            headers: HashMap::new(),
518        }
519    }
520
521    pub const fn from_headers(headers: HashMap<String, String>) -> Self {
522        Self { headers }
523    }
524}
525
526impl ContextCarrier for HttpHeadersCarrier {
527    fn get(&self, key: &str) -> Option<&str> {
528        self.headers.get(key).map(std::string::String::as_str)
529    }
530
531    fn set(&mut self, key: &str, value: String) {
532        self.headers.insert(key.to_string(), value);
533    }
534
535    fn keys(&self) -> Vec<&str> {
536        self.headers
537            .keys()
538            .map(std::string::String::as_str)
539            .collect()
540    }
541}
542
543/// Span builder for convenience
544pub struct SpanBuilder<'a> {
545    provider: &'a DistributedTracingProvider,
546    operation_name: String,
547    kind: SpanKind,
548    parent_context: Option<TelemetryContext>,
549    attributes: Vec<(String, String)>,
550    links: Vec<SpanLink>,
551}
552
553impl<'a> SpanBuilder<'a> {
554    pub fn new(provider: &'a DistributedTracingProvider, operation_name: &str) -> Self {
555        Self {
556            provider,
557            operation_name: operation_name.to_string(),
558            kind: SpanKind::Internal,
559            parent_context: None,
560            attributes: Vec::new(),
561            links: Vec::new(),
562        }
563    }
564
565    pub const fn with_kind(mut self, kind: SpanKind) -> Self {
566        self.kind = kind;
567        self
568    }
569
570    pub fn with_parent(mut self, parent: TelemetryContext) -> Self {
571        self.parent_context = Some(parent);
572        self
573    }
574
575    pub fn with_attribute(mut self, key: &str, value: &str) -> Self {
576        self.attributes.push((key.to_string(), value.to_string()));
577        self
578    }
579
580    pub fn with_link(mut self, trace_id: &str, span_id: &str) -> Self {
581        self.links.push(SpanLink {
582            trace_id: trace_id.to_string(),
583            span_id: span_id.to_string(),
584            attributes: HashMap::new(),
585        });
586        self
587    }
588
589    pub async fn start(self) -> DistributedSpan {
590        let mut span = self
591            .provider
592            .start_distributed_span(
593                &self.operation_name,
594                self.kind,
595                self.parent_context.as_ref(),
596            )
597            .await;
598
599        // Add attributes
600        for (key, value) in self.attributes {
601            span.attributes.insert(key, value);
602        }
603
604        // Add links
605        span.links = self.links;
606
607        // Update the span in active_spans if it was sampled
608        if let Some(existing_span) = self
609            .provider
610            .active_spans
611            .write()
612            .await
613            .get_mut(&span.span_id)
614        {
615            existing_span.attributes = span.attributes.clone();
616            existing_span.links = span.links.clone();
617        }
618
619        span
620    }
621}
622
623#[cfg(test)]
624mod tests {
625    use super::*;
626    use crate::telemetry::standard::StandardTelemetryProvider;
627    use crate::telemetry::TelemetryConfig;
628
629    #[tokio::test]
630    async fn test_distributed_span_lifecycle() {
631        let config = TelemetryConfig::default();
632        let base_provider = Arc::new(StandardTelemetryProvider::new(config));
633        let sampler = Arc::new(ProbabilitySampler::new(1.0)); // Always sample
634        let propagator = Arc::new(W3CTraceContextPropagator);
635
636        let provider = DistributedTracingProvider::new(base_provider, sampler, propagator);
637
638        // Start a span
639        let span = provider
640            .start_distributed_span("test_operation", SpanKind::Internal, None)
641            .await;
642
643        assert!(!span.trace_id.is_empty());
644        assert!(!span.span_id.is_empty());
645        assert_eq!(span.operation_name, "test_operation");
646
647        // Add event
648        provider
649            .add_span_event(&span.span_id, "test_event", vec![("key", "value")])
650            .await;
651
652        // End span
653        provider
654            .end_distributed_span(
655                &span.span_id,
656                SpanStatus {
657                    code: StatusCode::Ok,
658                    description: None,
659                },
660            )
661            .await;
662
663        // Check completed spans
664        let exported = provider.export_spans().await;
665        assert_eq!(exported.len(), 1);
666        assert_eq!(exported[0].span_id, span.span_id);
667    }
668
669    #[tokio::test]
670    async fn test_context_propagation() {
671        let propagator = W3CTraceContextPropagator;
672        let mut carrier = HttpHeadersCarrier::new();
673
674        let context = TelemetryContext {
675            trace_id: "00112233445566778899aabbccddeeff".to_string(),
676            span_id: "0123456789abcdef".to_string(),
677            parent_span_id: Some("fedcba9876543210".to_string()),
678            baggage: vec![("user".to_string(), "test".to_string())],
679        };
680
681        // Inject context
682        propagator.inject(&context, &mut carrier).await.unwrap();
683        assert!(carrier.get("traceparent").is_some());
684
685        // Extract context
686        let extracted = propagator.extract(&carrier).await.unwrap().unwrap();
687        assert_eq!(extracted.trace_id, context.trace_id);
688        assert_eq!(
689            extracted.parent_span_id,
690            Some("fedcba9876543210".to_string())
691        );
692    }
693
694    #[test]
695    fn test_probability_sampler() {
696        let sampler = ProbabilitySampler::new(0.5);
697
698        // Test multiple trace IDs to verify sampling behavior
699        let sampled_count = (0..1000)
700            .filter(|i| {
701                let trace_id = format!("trace-{}", i);
702                let decision =
703                    sampler.should_sample(&trace_id, None, "test", SpanKind::Internal, &[]);
704                decision.sampled
705            })
706            .count();
707
708        // Should be roughly 50% sampled (with some variance)
709        // Allow for wider variance due to hash distribution
710        assert!(
711            sampled_count > 350 && sampled_count < 650,
712            "Sampled count {} is outside expected range [350, 650]",
713            sampled_count
714        );
715    }
716}
kindly_guard_server/telemetry/distributed.rs

kindly_guard_server/telemetry/
distributed.rs
