tracing_opentelemetry/
tracer.rs

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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
use opentelemetry::{
    trace as otel,
    trace::{
        noop, SamplingDecision, SamplingResult, SpanBuilder, SpanContext, SpanId, SpanKind,
        TraceContextExt, TraceFlags, TraceId, TraceState,
    },
    Context as OtelContext,
};
use opentelemetry_sdk::trace::{IdGenerator, Tracer as SdkTracer};

/// An interface for authors of OpenTelemetry SDKs to build pre-sampled tracers.
///
/// The OpenTelemetry spec does not allow trace ids to be updated after a span
/// has been created. In order to associate extracted parent trace ids with
/// existing `tracing` spans, `tracing-opentelemetry` builds up otel span data
/// using a [`SpanBuilder`] instead, and creates / exports full otel spans only
/// when the associated `tracing` span is closed. However, in order to properly
/// inject otel [`Context`] information to downstream requests, the sampling
/// state must now be known _before_ the otel span has been created.
///
/// The logic for coming to a sampling decision and creating an injectable span
/// context from a [`SpanBuilder`] is encapsulated in the
/// [`PreSampledTracer::sampled_context`] method and has been implemented
/// for the standard OpenTelemetry SDK, but this trait may be implemented by
/// authors of alternate OpenTelemetry SDK implementations if they wish to have
/// `tracing` compatibility.
///
/// See the [`OpenTelemetrySpanExt::set_parent`] and
/// [`OpenTelemetrySpanExt::context`] methods for example usage.
///
/// [`Tracer`]: opentelemetry::trace::Tracer
/// [`SpanBuilder`]: opentelemetry::trace::SpanBuilder
/// [`PreSampledTracer::sampled_span_context`]: crate::PreSampledTracer::sampled_span_context
/// [`OpenTelemetrySpanExt::set_parent`]: crate::OpenTelemetrySpanExt::set_parent
/// [`OpenTelemetrySpanExt::context`]: crate::OpenTelemetrySpanExt::context
/// [`Context`]: opentelemetry::Context
pub trait PreSampledTracer {
    /// Produce an otel context containing an active and pre-sampled span for
    /// the given span builder data.
    ///
    /// The sampling decision, span context information, and parent context
    /// values must match the values recorded when the tracing span is closed.
    fn sampled_context(&self, data: &mut crate::OtelData) -> OtelContext;

    /// Generate a new trace id.
    fn new_trace_id(&self) -> otel::TraceId;

    /// Generate a new span id.
    fn new_span_id(&self) -> otel::SpanId;
}

impl PreSampledTracer for noop::NoopTracer {
    fn sampled_context(&self, data: &mut crate::OtelData) -> OtelContext {
        data.parent_cx.clone()
    }

    fn new_trace_id(&self) -> otel::TraceId {
        otel::TraceId::INVALID
    }

    fn new_span_id(&self) -> otel::SpanId {
        otel::SpanId::INVALID
    }
}

impl PreSampledTracer for SdkTracer {
    fn sampled_context(&self, data: &mut crate::OtelData) -> OtelContext {
        let parent_cx = &data.parent_cx;
        let builder = &mut data.builder;

        // Gather trace state
        let (trace_id, parent_trace_flags) =
            current_trace_state(builder, parent_cx, self.id_generator());

        // Sample or defer to existing sampling decisions
        let (flags, trace_state) = if let Some(result) = &builder.sampling_result {
            process_sampling_result(result, parent_trace_flags)
        } else {
            builder.sampling_result = Some(self.should_sample().should_sample(
                Some(parent_cx),
                trace_id,
                &builder.name,
                builder.span_kind.as_ref().unwrap_or(&SpanKind::Internal),
                builder.attributes.as_deref().unwrap_or(&[]),
                builder.links.as_deref().unwrap_or(&[]),
            ));

            process_sampling_result(
                builder.sampling_result.as_ref().unwrap(),
                parent_trace_flags,
            )
        }
        .unwrap_or_default();

        let span_id = builder.span_id.unwrap_or(SpanId::INVALID);
        let span_context = SpanContext::new(trace_id, span_id, flags, false, trace_state);
        parent_cx.with_remote_span_context(span_context)
    }

    fn new_trace_id(&self) -> otel::TraceId {
        self.id_generator().new_trace_id()
    }

    fn new_span_id(&self) -> otel::SpanId {
        self.id_generator().new_span_id()
    }
}

fn current_trace_state(
    builder: &SpanBuilder,
    parent_cx: &OtelContext,
    id_generator: &dyn IdGenerator,
) -> (TraceId, TraceFlags) {
    if parent_cx.has_active_span() {
        let span = parent_cx.span();
        let sc = span.span_context();
        (sc.trace_id(), sc.trace_flags())
    } else {
        (
            builder
                .trace_id
                .unwrap_or_else(|| id_generator.new_trace_id()),
            Default::default(),
        )
    }
}

fn process_sampling_result(
    sampling_result: &SamplingResult,
    trace_flags: TraceFlags,
) -> Option<(TraceFlags, TraceState)> {
    match sampling_result {
        SamplingResult {
            decision: SamplingDecision::Drop,
            ..
        } => None,
        SamplingResult {
            decision: SamplingDecision::RecordOnly,
            trace_state,
            ..
        } => Some((trace_flags & !TraceFlags::SAMPLED, trace_state.clone())),
        SamplingResult {
            decision: SamplingDecision::RecordAndSample,
            trace_state,
            ..
        } => Some((trace_flags | TraceFlags::SAMPLED, trace_state.clone())),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::OtelData;
    use opentelemetry::trace::TracerProvider as _;
    use opentelemetry_sdk::trace::{Config, Sampler, TracerProvider};

    #[test]
    fn assigns_default_trace_id_if_missing() {
        let provider = TracerProvider::default();
        let tracer = provider.tracer("test");
        let mut builder = SpanBuilder::from_name("empty".to_string());
        builder.span_id = Some(SpanId::from(1u64));
        builder.trace_id = None;
        let parent_cx = OtelContext::new();
        let cx = tracer.sampled_context(&mut OtelData { builder, parent_cx });
        let span = cx.span();
        let span_context = span.span_context();

        assert!(span_context.is_valid());
    }

    #[rustfmt::skip]
    fn sampler_data() -> Vec<(&'static str, Sampler, OtelContext, Option<SamplingResult>, bool)> {
        vec![
            // No parent samples
            ("empty_parent_cx_always_on", Sampler::AlwaysOn, OtelContext::new(), None, true),
            ("empty_parent_cx_always_off", Sampler::AlwaysOff, OtelContext::new(), None, false),

            // Remote parent samples
            ("remote_parent_cx_always_on", Sampler::AlwaysOn, OtelContext::new().with_remote_span_context(span_context(TraceFlags::SAMPLED, true)), None, true),
            ("remote_parent_cx_always_off", Sampler::AlwaysOff, OtelContext::new().with_remote_span_context(span_context(TraceFlags::SAMPLED, true)), None, false),
            ("sampled_remote_parent_cx_parent_based", Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), OtelContext::new().with_remote_span_context(span_context(TraceFlags::SAMPLED, true)), None, true),
            ("unsampled_remote_parent_cx_parent_based", Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), OtelContext::new().with_remote_span_context(span_context(TraceFlags::default(), true)), None, false),

            // Existing sampling result defers
            ("previous_drop_result_always_on", Sampler::AlwaysOn, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::Drop, attributes: vec![], trace_state: Default::default() }), false),
            ("previous_record_and_sample_result_always_off", Sampler::AlwaysOff, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::RecordAndSample, attributes: vec![], trace_state: Default::default() }), true),

            // Existing local parent, defers
            ("previous_drop_result_always_on", Sampler::AlwaysOn, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::Drop, attributes: vec![], trace_state: Default::default() }), false),
            ("previous_record_and_sample_result_always_off", Sampler::AlwaysOff, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::RecordAndSample, attributes: vec![], trace_state: Default::default() }), true),
        ]
    }

    #[test]
    fn sampled_context() {
        for (name, sampler, parent_cx, previous_sampling_result, is_sampled) in sampler_data() {
            let provider = TracerProvider::builder()
                .with_config(Config::default().with_sampler(sampler))
                .build();
            let tracer = provider.tracer("test");
            let mut builder = SpanBuilder::from_name("parent".to_string());
            builder.sampling_result = previous_sampling_result;
            let sampled = tracer.sampled_context(&mut OtelData { builder, parent_cx });

            assert_eq!(
                sampled.span().span_context().is_sampled(),
                is_sampled,
                "{}",
                name
            )
        }
    }

    fn span_context(trace_flags: TraceFlags, is_remote: bool) -> SpanContext {
        SpanContext::new(
            TraceId::from(1u128),
            SpanId::from(1u64),
            trace_flags,
            is_remote,
            Default::default(),
        )
    }
}