1use crate::{
41 sdk::InstrumentationLibrary,
42 trace::{Link, SpanKind, TraceContextExt, TraceId, TraceState},
43 Context, KeyValue,
44};
45#[cfg(feature = "serialize")]
46use serde::{Deserialize, Serialize};
47
48pub trait ShouldSample: Send + Sync + std::fmt::Debug {
52 #[allow(clippy::too_many_arguments)]
54 fn should_sample(
55 &self,
56 parent_context: Option<&Context>,
57 trace_id: TraceId,
58 name: &str,
59 span_kind: &SpanKind,
60 attributes: &[KeyValue],
61 links: &[Link],
62 instrumentation_library: &InstrumentationLibrary,
63 ) -> SamplingResult;
64}
65
66#[derive(Clone, Debug, PartialEq)]
68pub struct SamplingResult {
69 pub decision: SamplingDecision,
71 pub attributes: Vec<KeyValue>,
73 pub trace_state: TraceState,
75}
76
77#[derive(Clone, Debug, PartialEq)]
79pub enum SamplingDecision {
80 Drop,
83 RecordOnly,
85 RecordAndSample,
87}
88
89#[cfg_attr(feature = "serialize", derive(Deserialize, Serialize))]
91#[derive(Clone, Debug, PartialEq)]
92pub enum Sampler {
93 AlwaysOn,
95 AlwaysOff,
97 ParentBased(Box<Sampler>),
99 TraceIdRatioBased(f64),
103}
104
105impl ShouldSample for Sampler {
106 fn should_sample(
107 &self,
108 parent_context: Option<&Context>,
109 trace_id: TraceId,
110 name: &str,
111 span_kind: &SpanKind,
112 attributes: &[KeyValue],
113 links: &[Link],
114 instrumentation_library: &InstrumentationLibrary,
115 ) -> SamplingResult {
116 let decision = match self {
117 Sampler::AlwaysOn => SamplingDecision::RecordAndSample,
119 Sampler::AlwaysOff => SamplingDecision::Drop,
121 Sampler::ParentBased(delegate_sampler) => {
123 parent_context.filter(|cx| cx.has_active_span()).map_or(
124 delegate_sampler
125 .should_sample(
126 parent_context,
127 trace_id,
128 name,
129 span_kind,
130 attributes,
131 links,
132 instrumentation_library,
133 )
134 .decision,
135 |ctx| {
136 let span = ctx.span();
137 let parent_span_context = span.span_context();
138 if parent_span_context.is_sampled() {
139 SamplingDecision::RecordAndSample
140 } else {
141 SamplingDecision::Drop
142 }
143 },
144 )
145 }
146 Sampler::TraceIdRatioBased(prob) => {
148 if *prob >= 1.0 {
149 SamplingDecision::RecordAndSample
150 } else {
151 let prob_upper_bound = (prob.max(0.0) * (1u64 << 63) as f64) as u64;
152 let rnd_from_trace_id = (trace_id.0 as u64) >> 1;
155
156 if rnd_from_trace_id < prob_upper_bound {
157 SamplingDecision::RecordAndSample
158 } else {
159 SamplingDecision::Drop
160 }
161 }
162 }
163 };
164
165 SamplingResult {
166 decision,
167 attributes: Vec::new(),
169 trace_state: match parent_context {
171 Some(ctx) => ctx.span().span_context().trace_state().clone(),
172 None => TraceState::default(),
173 },
174 }
175 }
176}
177
178#[cfg(all(test, feature = "testing", feature = "trace"))]
179mod tests {
180 use super::*;
181 use crate::sdk::trace::{Sampler, SamplingDecision, ShouldSample};
182 use crate::testing::trace::TestSpan;
183 use crate::trace::{SpanContext, SpanId, TraceFlags, TraceState};
184 use rand::Rng;
185
186 #[rustfmt::skip]
187 fn sampler_data() -> Vec<(&'static str, Sampler, f64, bool, bool)> {
188 vec![
189 ("never_sample", Sampler::AlwaysOff, 0.0, false, false),
191 ("always_sample", Sampler::AlwaysOn, 1.0, false, false),
192 ("ratio_-1", Sampler::TraceIdRatioBased(-1.0), 0.0, false, false),
193 ("ratio_.25", Sampler::TraceIdRatioBased(0.25), 0.25, false, false),
194 ("ratio_.50", Sampler::TraceIdRatioBased(0.50), 0.5, false, false),
195 ("ratio_.75", Sampler::TraceIdRatioBased(0.75), 0.75, false, false),
196 ("ratio_2.0", Sampler::TraceIdRatioBased(2.0), 1.0, false, false),
197
198 ("delegate_to_always_on", Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), 1.0, false, false),
200 ("delegate_to_always_off", Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), 0.0, false, false),
201 ("delegate_to_ratio_-1", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(-1.0))), 0.0, false, false),
202 ("delegate_to_ratio_.25", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.25))), 0.25, false, false),
203 ("delegate_to_ratio_.50", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.50))), 0.50, false, false),
204 ("delegate_to_ratio_.75", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.75))), 0.75, false, false),
205 ("delegate_to_ratio_2.0", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(2.0))), 1.0, false, false),
206
207 ("unsampled_parent_with_ratio_-1", Sampler::TraceIdRatioBased(-1.0), 0.0, true, false),
209 ("unsampled_parent_with_ratio_.25", Sampler::TraceIdRatioBased(0.25), 0.25, true, false),
210 ("unsampled_parent_with_ratio_.50", Sampler::TraceIdRatioBased(0.50), 0.5, true, false),
211 ("unsampled_parent_with_ratio_.75", Sampler::TraceIdRatioBased(0.75), 0.75, true, false),
212 ("unsampled_parent_with_ratio_2.0", Sampler::TraceIdRatioBased(2.0), 1.0, true, false),
213 ("unsampled_parent_or_else_with_always_on", Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), 0.0, true, false),
214 ("unsampled_parent_or_else_with_always_off", Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), 0.0, true, false),
215 ("unsampled_parent_or_else_with_ratio_.25", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.25))), 0.0, true, false),
216
217 ("sampled_parent_with_ratio_-1", Sampler::TraceIdRatioBased(-1.0), 0.0, true, true),
219 ("sampled_parent_with_ratio_.25", Sampler::TraceIdRatioBased(0.25), 0.25, true, true),
220 ("sampled_parent_with_ratio_2.0", Sampler::TraceIdRatioBased(2.0), 1.0, true, true),
221
222 ("sampled_parent_or_else_with_always_on", Sampler::ParentBased(Box::new(Sampler::AlwaysOn)), 1.0, true, true),
224 ("sampled_parent_or_else_with_always_off", Sampler::ParentBased(Box::new(Sampler::AlwaysOff)), 1.0, true, true),
225 ("sampled_parent_or_else_with_ratio_.25", Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(0.25))), 1.0, true, true),
226
227 ("sampled_parent_span_with_never_sample", Sampler::AlwaysOff, 0.0, true, true),
229 ]
230 }
231
232 #[test]
233 fn sampling() {
234 let total = 10_000;
235 let mut rng = rand::thread_rng();
236 for (name, sampler, expectation, parent, sample_parent) in sampler_data() {
237 let mut sampled = 0;
238 for _ in 0..total {
239 let parent_context = if parent {
240 let trace_flags = if sample_parent {
241 TraceFlags::SAMPLED
242 } else {
243 TraceFlags::default()
244 };
245 let span_context = SpanContext::new(
246 TraceId::from_u128(1),
247 SpanId::from_u64(1),
248 trace_flags,
249 false,
250 TraceState::default(),
251 );
252
253 Some(Context::current_with_span(TestSpan(span_context)))
254 } else {
255 None
256 };
257
258 let trace_id = TraceId::from(rng.gen::<[u8; 16]>());
259 if sampler
260 .should_sample(
261 parent_context.as_ref(),
262 trace_id,
263 name,
264 &SpanKind::Internal,
265 &[],
266 &[],
267 &InstrumentationLibrary::default(),
268 )
269 .decision
270 == SamplingDecision::RecordAndSample
271 {
272 sampled += 1;
273 }
274 }
275 let mut tolerance = 0.0;
276 let got = sampled as f64 / total as f64;
277
278 if expectation > 0.0 && expectation < 1.0 {
279 let z = 4.75342; tolerance = z * (got * (1.0 - got) / total as f64).sqrt();
282 }
283
284 let diff = (got - expectation).abs();
285 assert!(
286 diff <= tolerance,
287 "{} got {:?} (diff: {}), expected {} (w/tolerance: {})",
288 name,
289 got,
290 diff,
291 expectation,
292 tolerance
293 );
294 }
295 }
296
297 #[test]
298 fn filter_parent_sampler_for_active_spans() {
299 let sampler = Sampler::ParentBased(Box::new(Sampler::AlwaysOn));
300 let cx = Context::current_with_value("some_value");
301 let instrumentation_library = InstrumentationLibrary::default();
302 let result = sampler.should_sample(
303 Some(&cx),
304 TraceId::from_u128(1),
305 "should sample",
306 &SpanKind::Internal,
307 &[],
308 &[],
309 &instrumentation_library,
310 );
311
312 assert_eq!(result.decision, SamplingDecision::RecordAndSample);
313 }
314}