opentelemetry_sdk/trace/
sampler.rs

1use opentelemetry::{
2    trace::{Link, SpanKind, TraceContextExt, TraceId, TraceState},
3    Context, KeyValue,
4};
5
6#[cfg(feature = "jaeger_remote_sampler")]
7mod jaeger_remote;
8
9/// The result of sampling logic for a given span.
10#[derive(Clone, Debug, PartialEq)]
11pub struct SamplingResult {
12    /// The decision about whether or not to sample.
13    pub decision: SamplingDecision,
14
15    /// Extra attributes to be added to the span by the sampler
16    pub attributes: Vec<KeyValue>,
17
18    /// Trace state from parent context, may be modified by samplers.
19    pub trace_state: TraceState,
20}
21
22/// Decision about whether or not to sample
23#[derive(Clone, Debug, PartialEq, Eq)]
24pub enum SamplingDecision {
25    /// Span will not be recorded and all events and attributes will be dropped.
26    Drop,
27
28    /// Span data wil be recorded, but not exported.
29    RecordOnly,
30
31    /// Span data will be recorded and exported.
32    RecordAndSample,
33}
34
35#[cfg(feature = "jaeger_remote_sampler")]
36pub use jaeger_remote::{JaegerRemoteSampler, JaegerRemoteSamplerBuilder};
37#[cfg(feature = "jaeger_remote_sampler")]
38use opentelemetry_http::HttpClient;
39
40/// The [`ShouldSample`] interface allows implementations to provide samplers
41/// which will return a sampling [`SamplingResult`] based on information that
42/// is typically available just before the [`Span`] was created.
43///
44/// # Sampling
45///
46/// Sampling is a mechanism to control the noise and overhead introduced by
47/// OpenTelemetry by reducing the number of samples of traces collected and
48/// sent to the backend.
49///
50/// Sampling may be implemented on different stages of a trace collection.
51/// [OpenTelemetry SDK] defines a [`ShouldSample`] interface that can be used at
52/// instrumentation points by libraries to check the sampling [`SamplingDecision`]
53/// early and optimize the amount of telemetry that needs to be collected.
54///
55/// All other sampling algorithms may be implemented on SDK layer in exporters,
56/// or even out of process in Agent or Collector.
57///
58/// The OpenTelemetry API has two properties responsible for the data collection:
59///
60/// * [`Span::is_recording()`]. If `true` the current [`Span`] records
61///   tracing events (attributes, events, status, etc.), otherwise all tracing
62///   events are dropped. Users can use this property to determine if expensive
63///   trace events can be avoided. [`SpanProcessor`]s will receive
64///   all spans with this flag set. However, [`SpanExporter`]s will
65///   not receive them unless the `Sampled` flag was set.
66/// * `Sampled` flag in [`SpanContext::trace_flags()`]. This flag is propagated
67///   via the [`SpanContext`] to child Spans. For more details see the [W3C
68///   specification](https://w3c.github.io/trace-context/). This flag indicates
69///   that the [`Span`] has been `sampled` and will be exported. [`SpanProcessor`]s
70///   and [`SpanExporter`]s will receive spans with the `Sampled` flag set for
71///   processing.
72///
73/// The flag combination `Sampled == false` and `is_recording == true` means
74/// that the current `Span` does record information, but most likely the child
75/// `Span` will not.
76///
77/// The flag combination `Sampled == true` and `is_recording == false` could
78/// cause gaps in the distributed trace, and because of this OpenTelemetry API
79/// MUST NOT allow this combination.
80///
81/// [OpenTelemetry SDK]: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/sdk.md#sampling
82/// [`SpanContext`]: opentelemetry::trace::SpanContext
83/// [`SpanContext::trace_flags()`]: opentelemetry::trace::SpanContext#method.trace_flags
84/// [`SpanExporter`]: crate::trace::SpanExporter
85/// [`SpanProcessor`]: crate::trace::SpanProcessor
86/// [`Span`]: opentelemetry::trace::Span
87/// [`Span::is_recording()`]: opentelemetry::trace::Span#tymethod.is_recording
88pub trait ShouldSample: CloneShouldSample + Send + Sync + std::fmt::Debug {
89    /// Returns the [`SamplingDecision`] for a [`Span`] to be created.
90    ///
91    /// The [`should_sample`] function can use any of the information provided to it in order to
92    /// make a decision about whether or not a [`Span`] should or should not be sampled. However,
93    /// there are performance implications on the creation of a span
94    ///
95    /// [`Span`]: opentelemetry::trace::Span
96    /// [`should_sample`]: ShouldSample::should_sample
97    #[allow(clippy::too_many_arguments)]
98    fn should_sample(
99        &self,
100        parent_context: Option<&Context>,
101        trace_id: TraceId,
102        name: &str,
103        span_kind: &SpanKind,
104        attributes: &[KeyValue],
105        links: &[Link],
106    ) -> SamplingResult;
107}
108
109impl<T: ShouldSample + 'static> From<T> for Box<dyn ShouldSample> {
110    #[inline]
111    fn from(value: T) -> Self {
112        Box::new(value)
113    }
114}
115
116/// This trait should not be used directly instead users should use [`ShouldSample`].
117pub trait CloneShouldSample {
118    fn box_clone(&self) -> Box<dyn ShouldSample>;
119}
120
121impl<T> CloneShouldSample for T
122where
123    T: ShouldSample + Clone + 'static,
124{
125    fn box_clone(&self) -> Box<dyn ShouldSample> {
126        Box::new(self.clone())
127    }
128}
129
130impl Clone for Box<dyn ShouldSample> {
131    fn clone(&self) -> Self {
132        self.box_clone()
133    }
134}
135
136/// Default Sampling options
137///
138/// The [built-in samplers] allow for simple decisions. For more complex scenarios consider
139/// implementing your own sampler using [`ShouldSample`] trait.
140///
141/// [built-in samplers]: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/sdk.md#built-in-samplers
142#[derive(Clone, Debug)]
143#[non_exhaustive]
144pub enum Sampler {
145    /// Always sample the trace
146    AlwaysOn,
147    /// Never sample the trace
148    AlwaysOff,
149    /// Respects the parent span's sampling decision or delegates a delegate sampler for root spans.
150    ParentBased(Box<dyn ShouldSample>),
151    /// Sample a given fraction of traces. Fractions >= 1 will always sample. If the parent span is
152    /// sampled, then it's child spans will automatically be sampled. Fractions < 0 are treated as
153    /// zero, but spans may still be sampled if their parent is.
154    /// *Note:* If this is used then all Spans in a trace will become sampled assuming that the
155    /// first span is sampled as it is based on the `trace_id` not the `span_id`
156    TraceIdRatioBased(f64),
157    /// Jaeger remote sampler supports any remote service that implemented the jaeger remote sampler protocol.
158    /// The proto definition can be found [here](https://github.com/jaegertracing/jaeger-idl/blob/main/proto/api_v2/sampling.proto)
159    ///
160    /// Jaeger remote sampler allows remotely controlling the sampling configuration for the SDKs.
161    /// The sampling is typically configured at the collector and the SDKs actively poll for changes.
162    /// The sampler uses TraceIdRatioBased or rate-limited sampler under the hood.
163    /// These samplers can be configured per whole service (a.k.a default), or per span name in a
164    /// given service (a.k.a per operation).
165    #[cfg(feature = "jaeger_remote_sampler")]
166    JaegerRemote(JaegerRemoteSampler),
167}
168
169impl Sampler {
170    /// Create a jaeger remote sampler builder.
171    ///
172    /// ### Arguments
173    /// * `runtime` - A runtime to run the HTTP client.
174    /// * `http_client` - An HTTP client to query the sampling endpoint.
175    /// * `default_sampler` - A default sampler to make a sampling decision when the remote is unavailable or before the SDK receives the first response from remote.
176    /// * `service_name` - The name of the service. This is a required parameter to query the sampling endpoint.
177    ///
178    /// See [here](https://github.com/open-telemetry/opentelemetry-rust/blob/main/examples/jaeger-remote-sampler/src/main.rs) for an example.
179    #[cfg(feature = "jaeger_remote_sampler")]
180    pub fn jaeger_remote<C, Sampler, R, Svc>(
181        runtime: R,
182        http_client: C,
183        default_sampler: Sampler,
184        service_name: Svc,
185    ) -> JaegerRemoteSamplerBuilder<C, Sampler, R>
186    where
187        C: HttpClient + 'static,
188        Sampler: ShouldSample,
189        R: crate::runtime::RuntimeChannel,
190        Svc: Into<String>,
191    {
192        JaegerRemoteSamplerBuilder::new(runtime, http_client, default_sampler, service_name)
193    }
194
195    /// Create a [`ParentBasedSampler`] with `root` as the sampler used for spans with
196    /// no parent or an invalid parent.
197    ///
198    /// `Sampler::ParentBased` only lets you configure that one root case; the other 4
199    /// branches defined by the [ParentBased sampler spec] (remote/local parent,
200    /// sampled/not-sampled) are fixed to the spec defaults. Use this instead when you
201    /// need to override any of those, via `ParentBasedSampler`'s `with_*` methods.
202    ///
203    /// [ParentBased sampler spec]: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/sdk.md#parentbased
204    pub fn parent_based(root: impl Into<Box<dyn ShouldSample>>) -> ParentBasedSampler {
205        ParentBasedSampler::new(root)
206    }
207}
208
209/// A [`ShouldSample`] sampler that lets each of the 5 branches defined by the
210/// [ParentBased sampler spec] be configured independently: the root sampler (no
211/// parent or an invalid parent), and a sampler for each combination of remote/local parent and
212/// sampled/not-sampled. Any branch left unset falls back to the spec default
213/// (`Sampler::AlwaysOn` for the two `*_sampled` branches, `Sampler::AlwaysOff` for
214/// the two `*_not_sampled` branches).
215///
216/// Unlike [`Sampler::ParentBased`], which keeps only the delegate's [`SamplingDecision`]
217/// and always returns the parent's trace state with no attributes, this returns the
218/// selected delegate's [`SamplingResult`] unchanged, preserving any attributes it adds
219/// and any trace state it modifies. Worth keeping in mind when moving an existing custom
220/// root sampler over to this API.
221///
222/// Build one with [`ParentBasedSampler::new`] or [`Sampler::parent_based`].
223///
224/// [ParentBased sampler spec]: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/sdk.md#parentbased
225#[derive(Clone, Debug)]
226pub struct ParentBasedSampler {
227    root: Box<dyn ShouldSample>,
228    remote_parent_sampled: Option<Box<dyn ShouldSample>>,
229    remote_parent_not_sampled: Option<Box<dyn ShouldSample>>,
230    local_parent_sampled: Option<Box<dyn ShouldSample>>,
231    local_parent_not_sampled: Option<Box<dyn ShouldSample>>,
232}
233
234impl ParentBasedSampler {
235    /// Create a new [`ParentBasedSampler`] with `root` as the sampler used for spans with
236    /// no parent or an invalid parent.
237    ///
238    /// The other 4 branches default to the [ParentBased sampler spec] values:
239    /// - `remote_parent_sampled`: [`Sampler::AlwaysOn`]
240    /// - `remote_parent_not_sampled`: [`Sampler::AlwaysOff`]
241    /// - `local_parent_sampled`: [`Sampler::AlwaysOn`]
242    /// - `local_parent_not_sampled`: [`Sampler::AlwaysOff`]
243    ///
244    /// Each of these branches can be configured using their respective `with_*` methods.
245    pub fn new(root: impl Into<Box<dyn ShouldSample>>) -> Self {
246        Self {
247            root: root.into(),
248            remote_parent_sampled: None,
249            remote_parent_not_sampled: None,
250            local_parent_sampled: None,
251            local_parent_not_sampled: None,
252        }
253    }
254    /// Sampler used when the parent's `SpanContext` is remote and sampled.
255    /// Defaults to `Sampler::AlwaysOn`.
256    pub fn with_remote_parent_sampled(mut self, sampler: impl Into<Box<dyn ShouldSample>>) -> Self {
257        self.remote_parent_sampled = Some(sampler.into());
258        self
259    }
260
261    /// Sampler used when the parent's `SpanContext` is remote and not sampled.
262    /// Defaults to `Sampler::AlwaysOff`.
263    pub fn with_remote_parent_not_sampled(
264        mut self,
265        sampler: impl Into<Box<dyn ShouldSample>>,
266    ) -> Self {
267        self.remote_parent_not_sampled = Some(sampler.into());
268        self
269    }
270
271    /// Sampler used when the parent's `SpanContext` is local and sampled.
272    /// Defaults to `Sampler::AlwaysOn`.
273    pub fn with_local_parent_sampled(mut self, sampler: impl Into<Box<dyn ShouldSample>>) -> Self {
274        self.local_parent_sampled = Some(sampler.into());
275        self
276    }
277
278    /// Sampler used when the parent's `SpanContext` is local and not sampled.
279    /// Defaults to `Sampler::AlwaysOff`.
280    pub fn with_local_parent_not_sampled(
281        mut self,
282        sampler: impl Into<Box<dyn ShouldSample>>,
283    ) -> Self {
284        self.local_parent_not_sampled = Some(sampler.into());
285        self
286    }
287}
288
289impl ShouldSample for ParentBasedSampler {
290    fn should_sample(
291        &self,
292        parent_context: Option<&Context>,
293        trace_id: TraceId,
294        name: &str,
295        span_kind: &SpanKind,
296        attributes: &[KeyValue],
297        links: &[Link],
298    ) -> SamplingResult {
299        let delegate = |sampler: &dyn ShouldSample| {
300            sampler.should_sample(parent_context, trace_id, name, span_kind, attributes, links)
301        };
302
303        match parent_context.filter(|cx| cx.span().span_context().is_valid()) {
304            None => delegate(self.root.as_ref()),
305            Some(cx) => {
306                let span = cx.span();
307                let parent_span_context = span.span_context();
308                match (
309                    parent_span_context.is_remote(),
310                    parent_span_context.is_sampled(),
311                ) {
312                    (true, true) => match &self.remote_parent_sampled {
313                        Some(s) => delegate(s.as_ref()),
314                        None => delegate(&Sampler::AlwaysOn),
315                    },
316                    (true, false) => match &self.remote_parent_not_sampled {
317                        Some(s) => delegate(s.as_ref()),
318                        None => delegate(&Sampler::AlwaysOff),
319                    },
320                    (false, true) => match &self.local_parent_sampled {
321                        Some(s) => delegate(s.as_ref()),
322                        None => delegate(&Sampler::AlwaysOn),
323                    },
324                    (false, false) => match &self.local_parent_not_sampled {
325                        Some(s) => delegate(s.as_ref()),
326                        None => delegate(&Sampler::AlwaysOff),
327                    },
328                }
329            }
330        }
331    }
332}
333
334impl ShouldSample for Sampler {
335    fn should_sample(
336        &self,
337        parent_context: Option<&Context>,
338        trace_id: TraceId,
339        name: &str,
340        span_kind: &SpanKind,
341        attributes: &[KeyValue],
342        links: &[Link],
343    ) -> SamplingResult {
344        let decision = match self {
345            // Always sample the trace
346            Sampler::AlwaysOn => SamplingDecision::RecordAndSample,
347            // Never sample the trace
348            Sampler::AlwaysOff => SamplingDecision::Drop,
349            // The parent decision if sampled; otherwise the decision of delegate_sampler.
350            // Note: checks `has_active_span()` for backwards compatibility with existing behavior
351            // (e.g. active spans with empty/dropped contexts are treated as unsampled parents).
352            // For full spec compliance where invalid parents fall back to root, use `ParentBasedSampler`.
353            Sampler::ParentBased(delegate_sampler) => parent_context
354                .filter(|cx| cx.has_active_span())
355                .map_or_else(
356                    || {
357                        delegate_sampler
358                            .should_sample(
359                                parent_context,
360                                trace_id,
361                                name,
362                                span_kind,
363                                attributes,
364                                links,
365                            )
366                            .decision
367                    },
368                    |ctx| {
369                        let span = ctx.span();
370                        let parent_span_context = span.span_context();
371                        if parent_span_context.is_sampled() {
372                            SamplingDecision::RecordAndSample
373                        } else {
374                            SamplingDecision::Drop
375                        }
376                    },
377                ),
378            // Probabilistically sample the trace.
379            Sampler::TraceIdRatioBased(prob) => sample_based_on_probability(prob, trace_id),
380            #[cfg(feature = "jaeger_remote_sampler")]
381            Sampler::JaegerRemote(remote_sampler) => {
382                remote_sampler
383                    .should_sample(parent_context, trace_id, name, span_kind, attributes, links)
384                    .decision
385            }
386        };
387        SamplingResult {
388            decision,
389            // No extra attributes ever set by the SDK samplers.
390            attributes: Vec::new(),
391            // all sampler in SDK will not modify trace state.
392            trace_state: match parent_context {
393                Some(ctx) => ctx.span().span_context().trace_state().clone(),
394                None => TraceState::default(),
395            },
396        }
397    }
398}
399
400pub(crate) fn sample_based_on_probability(prob: &f64, trace_id: TraceId) -> SamplingDecision {
401    if *prob >= 1.0 {
402        SamplingDecision::RecordAndSample
403    } else {
404        let prob_upper_bound = (prob.max(0.0) * (1u64 << 63) as f64) as u64;
405        // TODO: update behavior when the spec definition resolves
406        // https://github.com/open-telemetry/opentelemetry-specification/issues/1413
407        let bytes = trace_id.to_bytes();
408        let (_, low) = bytes.split_at(8);
409        let trace_id_low = u64::from_be_bytes(low.try_into().unwrap());
410        let rnd_from_trace_id = trace_id_low >> 1;
411
412        if rnd_from_trace_id < prob_upper_bound {
413            SamplingDecision::RecordAndSample
414        } else {
415            SamplingDecision::Drop
416        }
417    }
418}
419
420#[cfg(all(test, feature = "testing", feature = "trace"))]
421mod tests {
422    use super::*;
423    use crate::testing::trace::TestSpan;
424    use opentelemetry::trace::{SpanContext, SpanId, TraceFlags};
425    use rand::random;
426
427    #[rustfmt::skip]
428    fn sampler_data() -> Vec<(&'static str, Sampler, f64, bool, bool)> {
429        vec![
430            // Span w/o a parent
431            ("never_sample", Sampler::AlwaysOff, 0.0, false, false),
432            ("always_sample", Sampler::AlwaysOn, 1.0, false, false),
433            ("ratio_-1", Sampler::TraceIdRatioBased(-1.0), 0.0, false, false),
434            ("ratio_.25", Sampler::TraceIdRatioBased(0.25), 0.25, false, false),
435            ("ratio_.50", Sampler::TraceIdRatioBased(0.50), 0.5, false, false),
436            ("ratio_.75", Sampler::TraceIdRatioBased(0.75), 0.75, false, false),
437            ("ratio_2.0", Sampler::TraceIdRatioBased(2.0), 1.0, false, false),
438
439            // Spans w/o a parent delegate
440            ("delegate_to_always_on", Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), 1.0, false, false),
441            ("delegate_to_always_off", Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), 0.0, false, false),
442            ("delegate_to_ratio_-1", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(-1.0))), 0.0, false, false),
443            ("delegate_to_ratio_.25", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.25))), 0.25, false, false),
444            ("delegate_to_ratio_.50", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.50))), 0.50, false, false),
445            ("delegate_to_ratio_.75", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.75))), 0.75, false, false),
446            ("delegate_to_ratio_2.0", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(2.0))), 1.0, false, false),
447
448            // Spans with a parent that is *not* sampled act like spans w/o a parent
449            ("unsampled_parent_with_ratio_-1", Sampler::TraceIdRatioBased(-1.0), 0.0, true, false),
450            ("unsampled_parent_with_ratio_.25", Sampler::TraceIdRatioBased(0.25), 0.25, true, false),
451            ("unsampled_parent_with_ratio_.50", Sampler::TraceIdRatioBased(0.50), 0.5, true, false),
452            ("unsampled_parent_with_ratio_.75", Sampler::TraceIdRatioBased(0.75), 0.75, true, false),
453            ("unsampled_parent_with_ratio_2.0", Sampler::TraceIdRatioBased(2.0), 1.0, true, false),
454            ("unsampled_parent_or_else_with_always_on", Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), 0.0, true, false),
455            ("unsampled_parent_or_else_with_always_off", Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), 0.0, true, false),
456            ("unsampled_parent_or_else_with_ratio_.25", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.25))), 0.0, true, false),
457
458            // A ratio sampler with a parent that is sampled will ignore the parent
459            ("sampled_parent_with_ratio_-1", Sampler::TraceIdRatioBased(-1.0), 0.0, true, true),
460            ("sampled_parent_with_ratio_.25", Sampler::TraceIdRatioBased(0.25), 0.25, true, true),
461            ("sampled_parent_with_ratio_2.0", Sampler::TraceIdRatioBased(2.0), 1.0, true, true),
462
463            // Spans with a parent that is sampled, will always sample, regardless of the delegate sampler
464            ("sampled_parent_or_else_with_always_on", Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), 1.0, true, true),
465            ("sampled_parent_or_else_with_always_off", Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), 1.0, true, true),
466            ("sampled_parent_or_else_with_ratio_.25", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.25))), 1.0, true, true),
467
468            // Spans with a sampled parent, but when using the NeverSample Sampler, aren't sampled
469            ("sampled_parent_span_with_never_sample", Sampler::AlwaysOff, 0.0, true, true),
470        ]
471    }
472
473    #[test]
474    fn sampling() {
475        let total = 10_000;
476        for (name, sampler, expectation, parent, sample_parent) in sampler_data() {
477            let mut sampled = 0;
478            for _ in 0..total {
479                let parent_context = if parent {
480                    let trace_flags = if sample_parent {
481                        TraceFlags::SAMPLED
482                    } else {
483                        TraceFlags::default()
484                    };
485                    let span_context = SpanContext::new(
486                        TraceId::from(1),
487                        SpanId::from(1),
488                        trace_flags,
489                        false,
490                        TraceState::default(),
491                    );
492
493                    Some(Context::current_with_span(TestSpan(span_context)))
494                } else {
495                    None
496                };
497
498                let trace_id = TraceId::from(random::<u128>());
499                if sampler
500                    .should_sample(
501                        parent_context.as_ref(),
502                        trace_id,
503                        name,
504                        &SpanKind::Internal,
505                        &[],
506                        &[],
507                    )
508                    .decision
509                    == SamplingDecision::RecordAndSample
510                {
511                    sampled += 1;
512                }
513            }
514            let mut tolerance = 0.0;
515            let got = sampled as f64 / total as f64;
516
517            if expectation > 0.0 && expectation < 1.0 {
518                // See https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval
519                let z = 4.75342; // This should succeed 99.9999% of the time
520                tolerance = z * (got * (1.0 - got) / total as f64).sqrt();
521            }
522
523            let diff = (got - expectation).abs();
524            assert!(
525                diff <= tolerance,
526                "{name} got {got:?} (diff: {diff}), expected {expectation} (w/tolerance: {tolerance})"
527            );
528        }
529    }
530
531    #[test]
532    fn clone_a_parent_sampler() {
533        let sampler = Sampler::ParentBased(Box::new(Sampler::AlwaysOn));
534        #[allow(clippy::redundant_clone)]
535        let cloned_sampler = sampler.clone();
536
537        let cx = Context::current_with_value("some_value");
538
539        let result = sampler.should_sample(
540            Some(&cx),
541            TraceId::from(1),
542            "should sample",
543            &SpanKind::Internal,
544            &[],
545            &[],
546        );
547
548        let cloned_result = cloned_sampler.should_sample(
549            Some(&cx),
550            TraceId::from(1),
551            "should sample",
552            &SpanKind::Internal,
553            &[],
554            &[],
555        );
556
557        assert_eq!(result, cloned_result);
558    }
559
560    #[test]
561    fn parent_sampler() {
562        // name, delegate, context(with or without parent), expected decision
563        let test_cases = vec![
564            (
565                "should using delegate sampler",
566                Sampler::AlwaysOn,
567                Context::new(),
568                SamplingDecision::RecordAndSample,
569            ),
570            (
571                "should use parent result, always off",
572                Sampler::AlwaysOn,
573                Context::current_with_span(TestSpan(SpanContext::new(
574                    TraceId::from(1),
575                    SpanId::from(1),
576                    TraceFlags::default(), // not sampling
577                    false,
578                    TraceState::default(),
579                ))),
580                SamplingDecision::Drop,
581            ),
582            (
583                "should use parent result, always on",
584                Sampler::AlwaysOff,
585                Context::current_with_span(TestSpan(SpanContext::new(
586                    TraceId::from(1),
587                    SpanId::from(1),
588                    TraceFlags::SAMPLED, // not sampling
589                    false,
590                    TraceState::default(),
591                ))),
592                SamplingDecision::RecordAndSample,
593            ),
594        ];
595
596        for (name, delegate, parent_cx, expected) in test_cases {
597            let sampler = Sampler::ParentBased(Box::new(delegate));
598            let result = sampler.should_sample(
599                Some(&parent_cx),
600                TraceId::from(1),
601                name,
602                &SpanKind::Internal,
603                &[],
604                &[],
605            );
606
607            assert_eq!(result.decision, expected);
608        }
609    }
610
611    /// Delegate that tags its `SamplingResult` with its own name and the span id of the
612    /// parent context it was handed, so both the branch selection and the forwarding of
613    /// the full result are observable.
614    #[derive(Clone, Debug)]
615    struct TaggedSampler(&'static str);
616
617    impl ShouldSample for TaggedSampler {
618        fn should_sample(
619            &self,
620            parent_context: Option<&Context>,
621            _trace_id: TraceId,
622            _name: &str,
623            _span_kind: &SpanKind,
624            _attributes: &[KeyValue],
625            _links: &[Link],
626        ) -> SamplingResult {
627            let parent_span_id = parent_context
628                .map(|cx| cx.span().span_context().span_id())
629                .unwrap_or(SpanId::INVALID);
630            SamplingResult {
631                decision: SamplingDecision::RecordAndSample,
632                attributes: vec![
633                    KeyValue::new("branch", self.0),
634                    KeyValue::new("parent_span_id", parent_span_id.to_string()),
635                ],
636                trace_state: TraceState::from_key_value([("branch", self.0)]).unwrap(),
637            }
638        }
639    }
640
641    fn tagged_attr(result: &SamplingResult, key: &str) -> String {
642        result
643            .attributes
644            .iter()
645            .find(|kv| kv.key.as_str() == key)
646            .map(|kv| kv.value.as_str().into_owned())
647            .expect("delegate result must be forwarded unchanged")
648    }
649
650    fn parent_cx(is_remote: bool, sampled: bool) -> Context {
651        let trace_flags = if sampled {
652            TraceFlags::SAMPLED
653        } else {
654            TraceFlags::default()
655        };
656        Context::current_with_span(TestSpan(SpanContext::new(
657            TraceId::from(1),
658            SpanId::from(1),
659            trace_flags,
660            is_remote,
661            TraceState::default(),
662        )))
663    }
664
665    #[test]
666    fn parent_based_sampler_defaults_match_spec() {
667        // No overrides: *_sampled branches default to AlwaysOn, *_not_sampled to AlwaysOff,
668        // same as the plain `Sampler::ParentBased` behavior.
669        let sampler = Sampler::parent_based(Sampler::AlwaysOff);
670
671        assert_eq!(
672            sampler
673                .should_sample(
674                    None,
675                    TraceId::from(1),
676                    "root",
677                    &SpanKind::Internal,
678                    &[],
679                    &[]
680                )
681                .decision,
682            SamplingDecision::Drop,
683            "no parent uses the root sampler"
684        );
685        for is_remote in [true, false] {
686            let cx = parent_cx(is_remote, true);
687            assert_eq!(
688                sampler
689                    .should_sample(
690                        Some(&cx),
691                        TraceId::from(1),
692                        "s",
693                        &SpanKind::Internal,
694                        &[],
695                        &[]
696                    )
697                    .decision,
698                SamplingDecision::RecordAndSample,
699                "sampled parent (remote={is_remote}) defaults to AlwaysOn"
700            );
701            let cx = parent_cx(is_remote, false);
702            assert_eq!(
703                sampler
704                    .should_sample(
705                        Some(&cx),
706                        TraceId::from(1),
707                        "s",
708                        &SpanKind::Internal,
709                        &[],
710                        &[]
711                    )
712                    .decision,
713                SamplingDecision::Drop,
714                "unsampled parent (remote={is_remote}) defaults to AlwaysOff"
715            );
716        }
717    }
718
719    #[test]
720    fn parent_based_sampler_respects_branch_overrides() {
721        // Every branch gets a distinctly tagged delegate, so mixing up e.g. the remote and
722        // local branches fails the test instead of coincidentally matching.
723        let sampler = Sampler::parent_based(TaggedSampler("root"))
724            .with_remote_parent_sampled(TaggedSampler("remote_sampled"))
725            .with_remote_parent_not_sampled(TaggedSampler("remote_not_sampled"))
726            .with_local_parent_sampled(TaggedSampler("local_sampled"))
727            .with_local_parent_not_sampled(TaggedSampler("local_not_sampled"));
728
729        let cases = [
730            (true, true, "remote_sampled"),
731            (true, false, "remote_not_sampled"),
732            (false, true, "local_sampled"),
733            (false, false, "local_not_sampled"),
734        ];
735        for (is_remote, sampled, expected) in cases {
736            let cx = parent_cx(is_remote, sampled);
737            let result = sampler.should_sample(
738                Some(&cx),
739                TraceId::from(1),
740                "s",
741                &SpanKind::Internal,
742                &[],
743                &[],
744            );
745            assert_eq!(
746                tagged_attr(&result, "branch"),
747                expected,
748                "remote={is_remote}, sampled={sampled} should honor the overridden branch"
749            );
750            assert_eq!(
751                result.trace_state.get("branch"),
752                Some(expected),
753                "the delegate's modified trace state must be forwarded"
754            );
755            assert_eq!(
756                tagged_attr(&result, "parent_span_id"),
757                SpanId::from(1).to_string(),
758                "the delegate must receive the original parent context"
759            );
760        }
761
762        // No parent at all: the root delegate runs, and its full result is forwarded too.
763        let result = sampler.should_sample(
764            None,
765            TraceId::from(1),
766            "root",
767            &SpanKind::Internal,
768            &[],
769            &[],
770        );
771        assert_eq!(tagged_attr(&result, "branch"), "root");
772        assert_eq!(result.trace_state.get("branch"), Some("root"));
773    }
774
775    #[test]
776    fn parent_based_sampler_invalid_parent_falls_back_to_root() {
777        // Root is AlwaysOn, while local_parent_not_sampled defaults to AlwaysOff.
778        let sampler = Sampler::parent_based(Sampler::AlwaysOn);
779
780        // 1. Context with an active span whose span context is invalid (zeros).
781        let cx_invalid = Context::current_with_span(TestSpan(SpanContext::empty_context()));
782        assert!(cx_invalid.has_active_span());
783        assert!(!cx_invalid.span().span_context().is_valid());
784
785        let decision = sampler
786            .should_sample(
787                Some(&cx_invalid),
788                TraceId::from(1),
789                "s",
790                &SpanKind::Internal,
791                &[],
792                &[],
793            )
794            .decision;
795        assert_eq!(
796            decision,
797            SamplingDecision::RecordAndSample,
798            "invalid parent span must fall back to the root sampler (AlwaysOn)"
799        );
800
801        // 2. Context with no active span at all.
802        let cx_empty = Context::new();
803        assert!(!cx_empty.has_active_span());
804
805        let decision = sampler
806            .should_sample(
807                Some(&cx_empty),
808                TraceId::from(1),
809                "s",
810                &SpanKind::Internal,
811                &[],
812                &[],
813            )
814            .decision;
815        assert_eq!(
816            decision,
817            SamplingDecision::RecordAndSample,
818            "context without active span must fall back to the root sampler (AlwaysOn)"
819        );
820
821        // 3. The root delegate is still handed the original parent context, not `None`.
822        // Zero trace id with a non-zero span id is invalid, but distinguishable.
823        let cx = Context::current_with_span(TestSpan(SpanContext::new(
824            TraceId::INVALID,
825            SpanId::from(7),
826            TraceFlags::SAMPLED,
827            true,
828            TraceState::default(),
829        )));
830        let result = Sampler::parent_based(TaggedSampler("root")).should_sample(
831            Some(&cx),
832            TraceId::from(1),
833            "s",
834            &SpanKind::Internal,
835            &[],
836            &[],
837        );
838        assert_eq!(tagged_attr(&result, "branch"), "root");
839        assert_eq!(
840            tagged_attr(&result, "parent_span_id"),
841            SpanId::from(7).to_string(),
842            "the root delegate must receive the original parent context"
843        );
844    }
845
846    #[test]
847    fn parent_based_sampler_new() {
848        let sampler = ParentBasedSampler::new(Sampler::AlwaysOn);
849        assert_eq!(
850            sampler
851                .should_sample(
852                    None,
853                    TraceId::from(1),
854                    "root",
855                    &SpanKind::Internal,
856                    &[],
857                    &[]
858                )
859                .decision,
860            SamplingDecision::RecordAndSample
861        );
862    }
863}
opentelemetry_sdk/trace/sampler.rs

opentelemetry_sdk/trace/
sampler.rs
