tracing_opentelemetry/
tracer.rs

1use opentelemetry::{
2    trace as otel,
3    trace::{
4        noop, SamplingDecision, SamplingResult, SpanBuilder, SpanContext, SpanId, SpanKind,
5        TraceContextExt, TraceFlags, TraceId, TraceState,
6    },
7    Context as OtelContext,
8};
9use opentelemetry_sdk::trace::{IdGenerator, Tracer as SdkTracer};
10
11/// An interface for authors of OpenTelemetry SDKs to build pre-sampled tracers.
12///
13/// The OpenTelemetry spec does not allow trace ids to be updated after a span
14/// has been created. In order to associate extracted parent trace ids with
15/// existing `tracing` spans, `tracing-opentelemetry` builds up otel span data
16/// using a [`SpanBuilder`] instead, and creates / exports full otel spans only
17/// when the associated `tracing` span is closed. However, in order to properly
18/// inject otel [`Context`] information to downstream requests, the sampling
19/// state must now be known _before_ the otel span has been created.
20///
21/// The logic for coming to a sampling decision and creating an injectable span
22/// context from a [`SpanBuilder`] is encapsulated in the
23/// [`PreSampledTracer::sampled_context`] method and has been implemented
24/// for the standard OpenTelemetry SDK, but this trait may be implemented by
25/// authors of alternate OpenTelemetry SDK implementations if they wish to have
26/// `tracing` compatibility.
27///
28/// See the [`OpenTelemetrySpanExt::set_parent`] and
29/// [`OpenTelemetrySpanExt::context`] methods for example usage.
30///
31/// [`Tracer`]: opentelemetry::trace::Tracer
32/// [`SpanBuilder`]: opentelemetry::trace::SpanBuilder
33/// [`PreSampledTracer::sampled_span_context`]: crate::PreSampledTracer::sampled_span_context
34/// [`OpenTelemetrySpanExt::set_parent`]: crate::OpenTelemetrySpanExt::set_parent
35/// [`OpenTelemetrySpanExt::context`]: crate::OpenTelemetrySpanExt::context
36/// [`Context`]: opentelemetry::Context
37pub trait PreSampledTracer {
38    /// Produce an otel context containing an active and pre-sampled span for
39    /// the given span builder data.
40    ///
41    /// The sampling decision, span context information, and parent context
42    /// values must match the values recorded when the tracing span is closed.
43    fn sampled_context(&self, data: &mut crate::OtelData) -> OtelContext;
44
45    /// Generate a new trace id.
46    fn new_trace_id(&self) -> otel::TraceId;
47
48    /// Generate a new span id.
49    fn new_span_id(&self) -> otel::SpanId;
50}
51
52impl PreSampledTracer for noop::NoopTracer {
53    fn sampled_context(&self, data: &mut crate::OtelData) -> OtelContext {
54        data.parent_cx.clone()
55    }
56
57    fn new_trace_id(&self) -> otel::TraceId {
58        otel::TraceId::INVALID
59    }
60
61    fn new_span_id(&self) -> otel::SpanId {
62        otel::SpanId::INVALID
63    }
64}
65
66impl PreSampledTracer for SdkTracer {
67    fn sampled_context(&self, data: &mut crate::OtelData) -> OtelContext {
68        let parent_cx = &data.parent_cx;
69        let builder = &mut data.builder;
70
71        // Gather trace state
72        let (trace_id, parent_trace_flags) =
73            current_trace_state(builder, parent_cx, self.id_generator());
74
75        // Sample or defer to existing sampling decisions
76        let (flags, trace_state) = if let Some(result) = &builder.sampling_result {
77            process_sampling_result(result, parent_trace_flags)
78        } else {
79            builder.sampling_result = Some(self.should_sample().should_sample(
80                Some(parent_cx),
81                trace_id,
82                &builder.name,
83                builder.span_kind.as_ref().unwrap_or(&SpanKind::Internal),
84                builder.attributes.as_deref().unwrap_or(&[]),
85                builder.links.as_deref().unwrap_or(&[]),
86            ));
87
88            process_sampling_result(
89                builder.sampling_result.as_ref().unwrap(),
90                parent_trace_flags,
91            )
92        }
93        .unwrap_or_default();
94
95        let span_id = builder.span_id.unwrap_or(SpanId::INVALID);
96        let span_context = SpanContext::new(trace_id, span_id, flags, false, trace_state);
97        parent_cx.with_remote_span_context(span_context)
98    }
99
100    fn new_trace_id(&self) -> otel::TraceId {
101        self.id_generator().new_trace_id()
102    }
103
104    fn new_span_id(&self) -> otel::SpanId {
105        self.id_generator().new_span_id()
106    }
107}
108
109fn current_trace_state(
110    builder: &SpanBuilder,
111    parent_cx: &OtelContext,
112    id_generator: &dyn IdGenerator,
113) -> (TraceId, TraceFlags) {
114    if parent_cx.has_active_span() {
115        let span = parent_cx.span();
116        let sc = span.span_context();
117        (sc.trace_id(), sc.trace_flags())
118    } else {
119        (
120            builder
121                .trace_id
122                .unwrap_or_else(|| id_generator.new_trace_id()),
123            Default::default(),
124        )
125    }
126}
127
128fn process_sampling_result(
129    sampling_result: &SamplingResult,
130    trace_flags: TraceFlags,
131) -> Option<(TraceFlags, TraceState)> {
132    match sampling_result {
133        SamplingResult {
134            decision: SamplingDecision::Drop,
135            ..
136        } => None,
137        SamplingResult {
138            decision: SamplingDecision::RecordOnly,
139            trace_state,
140            ..
141        } => Some((trace_flags & !TraceFlags::SAMPLED, trace_state.clone())),
142        SamplingResult {
143            decision: SamplingDecision::RecordAndSample,
144            trace_state,
145            ..
146        } => Some((trace_flags | TraceFlags::SAMPLED, trace_state.clone())),
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use crate::OtelData;
154    use opentelemetry::trace::TracerProvider as _;
155    use opentelemetry_sdk::trace::{Config, Sampler, TracerProvider};
156
157    #[test]
158    fn assigns_default_trace_id_if_missing() {
159        let provider = TracerProvider::default();
160        let tracer = provider.tracer("test");
161        let mut builder = SpanBuilder::from_name("empty".to_string());
162        builder.span_id = Some(SpanId::from(1u64));
163        builder.trace_id = None;
164        let parent_cx = OtelContext::new();
165        let cx = tracer.sampled_context(&mut OtelData { builder, parent_cx });
166        let span = cx.span();
167        let span_context = span.span_context();
168
169        assert!(span_context.is_valid());
170    }
171
172    #[rustfmt::skip]
173    fn sampler_data() -> Vec<(&'static str, Sampler, OtelContext, Option<SamplingResult>, bool)> {
174        vec![
175            // No parent samples
176            ("empty_parent_cx_always_on", Sampler::AlwaysOn, OtelContext::new(), None, true),
177            ("empty_parent_cx_always_off", Sampler::AlwaysOff, OtelContext::new(), None, false),
178
179            // Remote parent samples
180            ("remote_parent_cx_always_on", Sampler::AlwaysOn, OtelContext::new().with_remote_span_context(span_context(TraceFlags::SAMPLED, true)), None, true),
181            ("remote_parent_cx_always_off", Sampler::AlwaysOff, OtelContext::new().with_remote_span_context(span_context(TraceFlags::SAMPLED, true)), None, false),
182            ("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),
183            ("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),
184
185            // Existing sampling result defers
186            ("previous_drop_result_always_on", Sampler::AlwaysOn, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::Drop, attributes: vec![], trace_state: Default::default() }), false),
187            ("previous_record_and_sample_result_always_off", Sampler::AlwaysOff, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::RecordAndSample, attributes: vec![], trace_state: Default::default() }), true),
188
189            // Existing local parent, defers
190            ("previous_drop_result_always_on", Sampler::AlwaysOn, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::Drop, attributes: vec![], trace_state: Default::default() }), false),
191            ("previous_record_and_sample_result_always_off", Sampler::AlwaysOff, OtelContext::new(), Some(SamplingResult { decision: SamplingDecision::RecordAndSample, attributes: vec![], trace_state: Default::default() }), true),
192        ]
193    }
194
195    #[test]
196    fn sampled_context() {
197        for (name, sampler, parent_cx, previous_sampling_result, is_sampled) in sampler_data() {
198            let provider = TracerProvider::builder()
199                .with_config(Config::default().with_sampler(sampler))
200                .build();
201            let tracer = provider.tracer("test");
202            let mut builder = SpanBuilder::from_name("parent".to_string());
203            builder.sampling_result = previous_sampling_result;
204            let sampled = tracer.sampled_context(&mut OtelData { builder, parent_cx });
205
206            assert_eq!(
207                sampled.span().span_context().is_sampled(),
208                is_sampled,
209                "{}",
210                name
211            )
212        }
213    }
214
215    fn span_context(trace_flags: TraceFlags, is_remote: bool) -> SpanContext {
216        SpanContext::new(
217            TraceId::from(1u128),
218            SpanId::from(1u64),
219            trace_flags,
220            is_remote,
221            Default::default(),
222        )
223    }
224}