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

//! # Sampler
use crate::api;

/// Sampling options
#[derive(Clone, Debug)]
pub enum Sampler {
    /// Always sample the trace
    Always,
    /// Never sample the trace
    Never,
    /// Sample if the parent span is sampled
    Parent,
    /// Sample a given fraction of traces. Fractions >= 1 will always sample. If the parent span is
    /// sampled, then it's child spans will automatically be sampled. Fractions < 0 are treated as
    /// zero, but spans may still be sampled if their parent is.
    Probability(f64),
}

impl api::Sampler for Sampler {
    fn should_sample(
        &self,
        parent_context: Option<&api::SpanContext>,
        trace_id: api::TraceId,
        _span_id: api::SpanId,
        _name: &str,
        _span_kind: &api::SpanKind,
        _attributes: &[api::KeyValue],
        _links: &[api::Link],
    ) -> api::SamplingResult {
        let decision = match self {
            // Always sample the trace
            Sampler::Always => api::SamplingDecision::RecordAndSampled,
            // Never sample the trace
            Sampler::Never => api::SamplingDecision::NotRecord,
            // Sample if the parent span is sampled
            Sampler::Parent => {
                if parent_context.map(|ctx| ctx.is_sampled()).unwrap_or(false) {
                    api::SamplingDecision::RecordAndSampled
                } else {
                    api::SamplingDecision::NotRecord
                }
            }
            // Match parent or probabilistically sample the trace.
            Sampler::Probability(prob) => {
                if *prob >= 1.0 || parent_context.map(|ctx| ctx.is_sampled()).unwrap_or(false) {
                    api::SamplingDecision::RecordAndSampled
                } else {
                    let prob_upper_bound = (prob.max(0.0) * (1u64 << 63) as f64) as u64;
                    // The trace_id is already randomly generated, so we don't need a new one here
                    let rnd_from_trace_id = (trace_id.to_u128() as u64) >> 1;

                    if rnd_from_trace_id < prob_upper_bound {
                        api::SamplingDecision::RecordAndSampled
                    } else {
                        api::SamplingDecision::NotRecord
                    }
                }
            }
        };

        api::SamplingResult {
            decision,
            // No extra attributes ever set by the SDK samplers.
            attributes: Vec::new(),
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::api::{self, Sampler as _};
    use crate::sdk::Sampler;
    use rand::Rng;

    #[rustfmt::skip]
    fn sampler_data() -> Vec<(&'static str, Sampler, f64, bool, bool)> {
        vec![
            // Span w/o a parent
            ("never_sample",    Sampler::Never,             0.0,  false, false),
            ("always_sample",   Sampler::Always,            1.0,  false, false),
            ("probability_-1",  Sampler::Probability(-1.0), 0.0,  false, false),
            ("probability_.25", Sampler::Probability(0.25), 0.25, false, false),
            ("probability_.50", Sampler::Probability(0.50), 0.5,  false, false),
            ("probability_.75", Sampler::Probability(0.75), 0.75, false, false),
            ("probability_2.0", Sampler::Probability(2.0),  1.0,  false, false),

            // Spans with a parent that is *not* sampled act like spans w/o a parent
            ("unsampled_parent_with_probability_-1",  Sampler::Probability(-1.0), 0.0, true, false),
            ("unsampled_parent_with_probability_.25", Sampler::Probability(0.25), 0.25, true, false),
            ("unsampled_parent_with_probability_.50", Sampler::Probability(0.50), 0.5, true, false),
            ("unsampled_parent_with_probability_.75", Sampler::Probability(0.75), 0.75, true, false),
            ("unsampled_parent_with_probability_2.0", Sampler::Probability(2.0),  1.0, true, false),

            // Spans with a parent that is sampled, will always sample, regardless of the probability
            ("sampled_parent_with_probability_-1",  Sampler::Probability(-1.0), 1.0, true, true),
            ("sampled_parent_with_probability_.25", Sampler::Probability(0.25), 1.0, true, true),
            ("sampled_parent_with_probability_2.0", Sampler::Probability(2.0),  1.0, true, true),

            // Spans with a sampled parent, but when using the NeverSample Sampler, aren't sampled
            ("sampled_parent_span_with_never_sample", Sampler::Never, 0.0, true, true),
        ]
    }

    #[test]
    fn sampling() {
        let total = 10_000;
        let mut rng = rand::thread_rng();
        for (name, sampler, expectation, parent, sample_parent) in sampler_data() {
            let mut sampled = 0;
            for _ in 0..total {
                let parent_context = if parent {
                    let trace_flags = if sample_parent {
                        api::TRACE_FLAG_SAMPLED
                    } else {
                        0
                    };
                    Some(api::SpanContext::new(
                        api::TraceId::from_u128(1),
                        api::SpanId::from_u64(1),
                        trace_flags,
                        false,
                    ))
                } else {
                    None
                };
                let trace_id = api::TraceId::from_u128(rng.gen());
                if sampler
                    .should_sample(
                        parent_context.as_ref(),
                        trace_id,
                        api::SpanId::from_u64(1),
                        name,
                        &api::SpanKind::Internal,
                        &[],
                        &[],
                    )
                    .decision
                    == api::SamplingDecision::RecordAndSampled
                {
                    sampled += 1;
                }
            }
            let mut tolerance = 0.0;
            let got = sampled as f64 / total as f64;

            if expectation > 0.0 && expectation < 1.0 {
                // See https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval
                let z = 4.75342; // This should succeed 99.9999% of the time
                tolerance = z * (got * (1.0 - got) / total as f64).sqrt();
            }

            let diff = (got - expectation).abs();
            assert!(
                diff <= tolerance,
                "{} got {:?} (diff: {}), expected {} (w/tolerance: {})",
                name,
                got,
                diff,
                expectation,
                tolerance
            );
        }
    }
}