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
11pub trait PreSampledTracer {
38 fn sampled_context(&self, data: &mut crate::OtelData) -> OtelContext;
44
45 fn new_trace_id(&self) -> otel::TraceId;
47
48 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 let (trace_id, parent_trace_flags) =
73 current_trace_state(builder, parent_cx, self.id_generator());
74
75 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 ("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_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 ("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 ("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}