apollo-router 1.61.13

A configurable, high-performance routing runtime for Apollo Federation 🚀
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

use opentelemetry_api::Context;
use opentelemetry_api::trace::SpanContext;
use opentelemetry_api::trace::TraceResult;
use opentelemetry_sdk::export::trace::SpanData;
use opentelemetry_sdk::trace::Span;
use opentelemetry_sdk::trace::SpanProcessor;

/// When using the Datadog agent we need spans to always be exported. However, the batch span processor will only export spans that are sampled.
/// This wrapper will override the trace flags to always sample.
/// THe datadog exporter itself will look at the `sampling.priority` trace context attribute to determine if the span should be sampled.
#[derive(Debug)]
pub(crate) struct DatadogSpanProcessor<T: SpanProcessor> {
    delegate: T,
}

impl<T: SpanProcessor> DatadogSpanProcessor<T> {
    pub(crate) fn new(delegate: T) -> Self {
        Self { delegate }
    }
}

impl<T: SpanProcessor> SpanProcessor for DatadogSpanProcessor<T> {
    fn on_start(&self, span: &mut Span, cx: &Context) {
        self.delegate.on_start(span, cx)
    }

    fn on_end(&self, mut span: SpanData) {
        // Note that the trace state for measuring and sampling priority is handled in the AgentSampler
        // The only purpose of this span processor is to ensure that a span can pass through a batch processor.
        let new_trace_flags = span.span_context.trace_flags().with_sampled(true);
        span.span_context = SpanContext::new(
            span.span_context.trace_id(),
            span.span_context.span_id(),
            new_trace_flags,
            span.span_context.is_remote(),
            span.span_context.trace_state().clone(),
        );
        self.delegate.on_end(span)
    }

    fn force_flush(&self) -> TraceResult<()> {
        self.delegate.force_flush()
    }

    fn shutdown(&mut self) -> TraceResult<()> {
        self.delegate.shutdown()
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::Mutex;
    use std::time::SystemTime;

    use opentelemetry_api::Context;
    use opentelemetry_api::trace::SpanId;
    use opentelemetry_api::trace::SpanKind;
    use opentelemetry_api::trace::TraceFlags;
    use opentelemetry_api::trace::TraceId;
    use opentelemetry_sdk::trace::EvictedHashMap;
    use opentelemetry_sdk::trace::EvictedQueue;
    use opentelemetry_sdk::trace::SpanProcessor;

    use super::*;

    #[derive(Debug, Clone)]
    struct MockSpanProcessor {
        spans: Arc<Mutex<Vec<SpanData>>>,
    }

    impl MockSpanProcessor {
        fn new() -> Self {
            Self {
                spans: Default::default(),
            }
        }
    }

    impl SpanProcessor for MockSpanProcessor {
        fn on_start(&self, _span: &mut Span, _cx: &Context) {}

        fn on_end(&self, span: SpanData) {
            self.spans.lock().unwrap().push(span);
        }

        fn force_flush(&self) -> TraceResult<()> {
            Ok(())
        }

        fn shutdown(&mut self) -> TraceResult<()> {
            Ok(())
        }
    }

    #[test]
    fn test_on_end_updates_trace_flags() {
        let mock_processor = MockSpanProcessor::new();
        let processor = DatadogSpanProcessor::new(mock_processor.clone());
        let span_context = SpanContext::new(
            TraceId::from_u128(1),
            SpanId::from_u64(1),
            TraceFlags::default(),
            false,
            Default::default(),
        );
        let span_data = SpanData {
            span_context,
            parent_span_id: SpanId::from_u64(1),
            span_kind: SpanKind::Client,
            name: Default::default(),
            start_time: SystemTime::now(),
            end_time: SystemTime::now(),
            attributes: EvictedHashMap::new(32, 32),
            events: EvictedQueue::new(32),
            links: EvictedQueue::new(32),
            status: Default::default(),
            resource: Default::default(),
            instrumentation_lib: Default::default(),
        };

        processor.on_end(span_data.clone());

        // Verify that the trace flags are updated to sampled
        let updated_trace_flags = span_data.span_context.trace_flags().with_sampled(true);
        let stored_spans = mock_processor.spans.lock().unwrap();
        assert_eq!(stored_spans.len(), 1);
        assert_eq!(
            stored_spans[0].span_context.trace_flags(),
            updated_trace_flags
        );
    }
}