use crate::api;
use crate::api::trace::TraceState;
pub trait ShouldSample: Send + Sync + std::fmt::Debug {
#[allow(clippy::too_many_arguments)]
fn should_sample(
&self,
parent_context: Option<&api::trace::SpanReference>,
trace_id: api::trace::TraceId,
name: &str,
span_kind: &api::trace::SpanKind,
attributes: &[api::KeyValue],
links: &[api::trace::Link],
) -> SamplingResult;
}
#[derive(Clone, Debug, PartialEq)]
pub struct SamplingResult {
pub decision: SamplingDecision,
pub attributes: Vec<api::KeyValue>,
pub trace_state: TraceState,
}
#[derive(Clone, Debug, PartialEq)]
pub enum SamplingDecision {
Drop,
RecordOnly,
RecordAndSample,
}
#[derive(Clone, Debug)]
pub enum Sampler {
AlwaysOn,
AlwaysOff,
ParentBased(Box<Sampler>),
TraceIdRatioBased(f64),
}
impl ShouldSample for Sampler {
fn should_sample(
&self,
parent_context: Option<&api::trace::SpanReference>,
trace_id: api::trace::TraceId,
name: &str,
span_kind: &api::trace::SpanKind,
attributes: &[api::KeyValue],
links: &[api::trace::Link],
) -> SamplingResult {
let decision = match self {
Sampler::AlwaysOn => SamplingDecision::RecordAndSample,
Sampler::AlwaysOff => SamplingDecision::Drop,
Sampler::ParentBased(delegate_sampler) => parent_context.map_or(
delegate_sampler
.should_sample(parent_context, trace_id, name, span_kind, attributes, links)
.decision,
|ctx| {
if ctx.is_sampled() {
SamplingDecision::RecordAndSample
} else {
SamplingDecision::Drop
}
},
),
Sampler::TraceIdRatioBased(prob) => {
if *prob >= 1.0 {
SamplingDecision::RecordAndSample
} else {
let prob_upper_bound = (prob.max(0.0) * (1u64 << 63) as f64) as u64;
let rnd_from_trace_id = (trace_id.to_u128() as u64) >> 1;
if rnd_from_trace_id < prob_upper_bound {
SamplingDecision::RecordAndSample
} else {
SamplingDecision::Drop
}
}
}
};
SamplingResult {
decision,
attributes: Vec::new(),
trace_state: match parent_context {
Some(ctx) => ctx.trace_state().clone(),
None => TraceState::default(),
},
}
}
}
#[cfg(test)]
mod tests {
use crate::api;
use crate::api::trace::TraceState;
use crate::sdk::trace::{Sampler, SamplingDecision, ShouldSample};
use rand::Rng;
#[rustfmt::skip]
fn sampler_data() -> Vec<(&'static str, Sampler, f64, bool, bool)> {
vec![
("never_sample", Sampler::AlwaysOff, 0.0, false, false),
("always_sample", Sampler::AlwaysOn, 1.0, false, false),
("ratio_-1", Sampler::TraceIdRatioBased(-1.0), 0.0, false, false),
("ratio_.25", Sampler::TraceIdRatioBased(0.25), 0.25, false, false),
("ratio_.50", Sampler::TraceIdRatioBased(0.50), 0.5, false, false),
("ratio_.75", Sampler::TraceIdRatioBased(0.75), 0.75, false, false),
("ratio_2.0", Sampler::TraceIdRatioBased(2.0), 1.0, false, false),
("delegate_to_always_on", Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), 1.0, false, false),
("delegate_to_always_off", Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), 0.0, false, false),
("delegate_to_ratio_-1", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(-1.0))), 0.0, false, false),
("delegate_to_ratio_.25", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.25))), 0.25, false, false),
("delegate_to_ratio_.50", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.50))), 0.50, false, false),
("delegate_to_ratio_.75", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.75))), 0.75, false, false),
("delegate_to_ratio_2.0", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(2.0))), 1.0, false, false),
("unsampled_parent_with_ratio_-1", Sampler::TraceIdRatioBased(-1.0), 0.0, true, false),
("unsampled_parent_with_ratio_.25", Sampler::TraceIdRatioBased(0.25), 0.25, true, false),
("unsampled_parent_with_ratio_.50", Sampler::TraceIdRatioBased(0.50), 0.5, true, false),
("unsampled_parent_with_ratio_.75", Sampler::TraceIdRatioBased(0.75), 0.75, true, false),
("unsampled_parent_with_ratio_2.0", Sampler::TraceIdRatioBased(2.0), 1.0, true, false),
("unsampled_parent_or_else_with_always_on", Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), 0.0, true, false),
("unsampled_parent_or_else_with_always_off", Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), 0.0, true, false),
("unsampled_parent_or_else_with_ratio_.25", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.25))), 0.0, true, false),
("sampled_parent_with_ratio_-1", Sampler::TraceIdRatioBased(-1.0), 0.0, true, true),
("sampled_parent_with_ratio_.25", Sampler::TraceIdRatioBased(0.25), 0.25, true, true),
("sampled_parent_with_ratio_2.0", Sampler::TraceIdRatioBased(2.0), 1.0, true, true),
("sampled_parent_or_else_with_always_on", Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), 1.0, true, true),
("sampled_parent_or_else_with_always_off", Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), 1.0, true, true),
("sampled_parent_or_else_with_ratio_.25", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.25))), 1.0, true, true),
("sampled_parent_span_with_never_sample", Sampler::AlwaysOff, 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::TRACE_FLAG_SAMPLED
} else {
0
};
Some(api::trace::SpanReference::new(
api::trace::TraceId::from_u128(1),
api::trace::SpanId::from_u64(1),
trace_flags,
false,
TraceState::default(),
))
} else {
None
};
let trace_id = api::trace::TraceId::from_u128(rng.gen());
if sampler
.should_sample(
parent_context.as_ref(),
trace_id,
name,
&api::trace::SpanKind::Internal,
&[],
&[],
)
.decision
== SamplingDecision::RecordAndSample
{
sampled += 1;
}
}
let mut tolerance = 0.0;
let got = sampled as f64 / total as f64;
if expectation > 0.0 && expectation < 1.0 {
let z = 4.75342;
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
);
}
}
}