1use crate::dd_constants::{
5 RL_EFFECTIVE_RATE, SAMPLING_AGENT_RATE_TAG_KEY, SAMPLING_DECISION_MAKER_TAG_KEY,
6 SAMPLING_KNUTH_RATE_TAG_KEY, SAMPLING_PRIORITY_TAG_KEY, SAMPLING_RULE_RATE_TAG_KEY,
7};
8use crate::dd_sampling::{mechanism, priority, SamplingMechanism, SamplingPriority};
9use crate::sampling_rule_config::SamplingRuleConfig;
10
11pub type SamplingRulesCallback = Box<dyn for<'a> Fn(&'a [SamplingRuleConfig]) + Send + Sync>;
14
15use crate::types::{SamplingData, SpanProperties};
16
17use super::agent_service_sampler::{AgentRates, ServicesSampler};
18use super::rate_limiter::RateLimiter;
19use super::rules_sampler::RulesSampler;
20use super::sampling_rule::SamplingRule;
21
22#[derive(Clone, Debug)]
24pub struct DatadogSampler {
25 rules: RulesSampler,
27
28 service_samplers: ServicesSampler,
30
31 rate_limiter: RateLimiter,
33}
34
35impl DatadogSampler {
36 pub fn new(rules: Vec<SamplingRule>, rate_limit: i32) -> Self {
38 let limiter = RateLimiter::new(rate_limit, None);
40
41 DatadogSampler {
42 rules: RulesSampler::new(rules),
43 service_samplers: ServicesSampler::default(),
44 rate_limiter: limiter,
45 }
46 }
47
48 #[cfg(test)]
50 pub(crate) fn update_service_rates(&self, rates: impl IntoIterator<Item = (String, f64)>) {
51 self.service_samplers.update_rates(rates);
52 }
53
54 pub fn on_agent_response(&self) -> Box<dyn for<'a> Fn(&'a str) + Send + Sync> {
55 let service_samplers = self.service_samplers.clone();
56 Box::new(move |s: &str| {
57 let Ok(new_rates) = serde_json::de::from_str::<AgentRates>(s) else {
58 return;
59 };
60 let Some(new_rates) = new_rates.rate_by_service else {
61 return;
62 };
63 service_samplers.update_rates(new_rates.into_iter().map(|(k, v)| (k.to_string(), v)));
64 })
65 }
66
67 pub fn on_rules_update(&self) -> SamplingRulesCallback {
73 let rules_sampler = self.rules.clone();
74 Box::new(move |rule_configs: &[SamplingRuleConfig]| {
75 let new_rules = SamplingRule::from_configs(rule_configs.to_vec());
76
77 rules_sampler.update_rules(new_rules);
78 })
79 }
80
81 fn service_key(&self, span: &impl SpanProperties) -> String {
83 format!("service:{},env:{}", span.service(), span.env())
86 }
87
88 fn find_matching_rule(&self, span: &impl SpanProperties) -> Option<SamplingRule> {
90 self.rules.find_matching_rule(|rule| rule.matches(span))
91 }
92
93 fn get_sampling_mechanism(
95 &self,
96 rule: Option<&SamplingRule>,
97 used_agent_sampler: bool,
98 ) -> SamplingMechanism {
99 if let Some(rule) = rule {
100 match rule.provenance.as_str() {
103 "customer" => mechanism::REMOTE_USER_TRACE_SAMPLING_RULE,
104 "dynamic" => mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE,
105 _ => mechanism::LOCAL_USER_TRACE_SAMPLING_RULE,
106 }
107 } else if used_agent_sampler {
108 mechanism::AGENT_RATE_BY_SERVICE
109 } else {
110 mechanism::DEFAULT
112 }
113 }
114
115 pub fn sample(&self, data: &impl SamplingData) -> DdSamplingResult {
120 if let Some(is_parent_sampled) = data.is_parent_sampled() {
121 let priority = match is_parent_sampled {
122 false => priority::AUTO_REJECT,
123 true => priority::AUTO_KEEP,
124 };
125 return DdSamplingResult {
126 priority,
127 trace_root_info: None,
128 };
129 }
130
131 data.with_span_properties(self, |sampler, span| sampler.sample_root(data, span))
132 }
133
134 fn sample_root(
141 &self,
142 data: &impl SamplingData,
143 span: &impl SpanProperties,
144 ) -> DdSamplingResult {
145 let mut is_keep = true;
146 let mut used_agent_sampler = false;
147 let sample_rate;
148 let mut rl_effective_rate: Option<f64> = None;
149 let trace_id = data.trace_id();
150
151 let matching_rule = self.find_matching_rule(span);
152
153 if let Some(rule) = &matching_rule {
154 sample_rate = rule.sample_rate;
155
156 if !rule.sample(trace_id) {
157 is_keep = false;
158 } else if !self.rate_limiter.is_allowed() {
159 is_keep = false;
161 rl_effective_rate = Some(self.rate_limiter.effective_rate());
162 }
163 } else {
164 let service_key = self.service_key(span);
165 if let Some(sampler) = self.service_samplers.get(&service_key) {
166 used_agent_sampler = true;
167 sample_rate = sampler.sample_rate();
168 if !sampler.sample(trace_id) {
169 is_keep = false;
170 }
171 } else {
172 sample_rate = 1.0;
174 }
175 }
176
177 let mechanism = self.get_sampling_mechanism(matching_rule.as_ref(), used_agent_sampler);
178
179 DdSamplingResult {
180 priority: mechanism.to_priority(is_keep),
181 trace_root_info: Some(TraceRootSamplingInfo {
182 mechanism,
183 rate: sample_rate,
184 rl_effective_rate,
185 }),
186 }
187 }
188}
189
190fn format_sampling_rate(rate: f64) -> Option<String> {
202 if rate.is_nan() || !(0.0..=1.0).contains(&rate) {
203 return None;
204 }
205
206 if rate == 0.0 {
207 return Some("0".to_string());
208 }
209
210 let digits = 6_i32;
211 let magnitude = rate.abs().log10().floor() as i32;
212 let scale = 10f64.powi(digits - 1 - magnitude);
213 let rounded = (rate * scale).round() / scale;
214
215 let decimal_places = if magnitude >= digits - 1 {
217 0
218 } else {
219 (digits - 1 - magnitude) as usize
220 };
221
222 let s = format!("{:.prec$}", rounded, prec = decimal_places);
223 Some(if s.contains('.') {
225 s.trim_end_matches('0').trim_end_matches('.').to_string()
226 } else {
227 s
228 })
229}
230
231pub struct TraceRootSamplingInfo {
232 mechanism: SamplingMechanism,
233 rate: f64,
234 rl_effective_rate: Option<f64>,
235}
236
237impl TraceRootSamplingInfo {
238 pub fn mechanism(&self) -> SamplingMechanism {
240 self.mechanism
241 }
242
243 pub fn rate(&self) -> f64 {
245 self.rate
246 }
247
248 pub fn rl_effective_rate(&self) -> Option<f64> {
250 self.rl_effective_rate
251 }
252}
253
254pub struct DdSamplingResult {
255 priority: SamplingPriority,
256 trace_root_info: Option<TraceRootSamplingInfo>,
257}
258
259impl DdSamplingResult {
260 #[inline(always)]
261 pub fn get_priority(&self) -> SamplingPriority {
262 self.priority
263 }
264
265 pub fn get_trace_root_sampling_info(&self) -> &Option<TraceRootSamplingInfo> {
266 &self.trace_root_info
267 }
268
269 pub fn to_dd_sampling_tags<F>(&self, factory: &F) -> Option<Vec<F::Attribute>>
277 where
278 F: crate::types::AttributeFactory,
279 {
280 let Some(root_info) = &self.trace_root_info else {
281 return None; };
283
284 let mut result: Vec<F::Attribute>;
285 if let Some(limit) = root_info.rl_effective_rate() {
287 result = Vec::with_capacity(4);
288 result.push(factory.create_f64(RL_EFFECTIVE_RATE, limit));
289 } else {
290 result = Vec::with_capacity(3);
291 }
292
293 let mechanism = root_info.mechanism();
295 result.push(factory.create_string(SAMPLING_DECISION_MAKER_TAG_KEY, mechanism.to_cow()));
296
297 match mechanism {
299 mechanism::AGENT_RATE_BY_SERVICE => {
300 result.push(factory.create_f64(SAMPLING_AGENT_RATE_TAG_KEY, root_info.rate()));
301 if let Some(rate_str) = format_sampling_rate(root_info.rate()) {
302 result.push(factory.create_string(
303 SAMPLING_KNUTH_RATE_TAG_KEY,
304 std::borrow::Cow::Owned(rate_str),
305 ));
306 }
307 }
308 mechanism::REMOTE_USER_TRACE_SAMPLING_RULE
309 | mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE
310 | mechanism::LOCAL_USER_TRACE_SAMPLING_RULE => {
311 result.push(factory.create_f64(SAMPLING_RULE_RATE_TAG_KEY, root_info.rate()));
312 if let Some(rate_str) = format_sampling_rate(root_info.rate()) {
313 result.push(factory.create_string(
314 SAMPLING_KNUTH_RATE_TAG_KEY,
315 std::borrow::Cow::Owned(rate_str),
316 ));
317 }
318 }
319 _ => {}
320 }
321
322 let priority = self.priority;
323 result.push(factory.create_i64(SAMPLING_PRIORITY_TAG_KEY, priority.into_i8() as i64));
324
325 Some(result)
326 }
327}
328
329#[cfg(test)]
330mod tests {
331 use super::*;
332 use crate::constants::{
333 attr::{ENV_TAG, RESOURCE_TAG},
334 pattern,
335 };
336 use crate::types::{AttributeLike, TraceIdLike, ValueLike};
337 use std::borrow::Cow;
338 use std::collections::HashMap;
339
340 const HTTP_REQUEST_METHOD: &str = "http.request.method";
342 const SERVICE_NAME: &str = "service.name";
343
344 const HTTP_RESPONSE_STATUS_CODE: &str = "http.response.status_code";
346 const HTTP_STATUS_CODE: &str = "http.status_code";
347
348 #[derive(Clone, Debug, PartialEq, Eq)]
353 struct TestTraceId {
354 bytes: [u8; 16],
355 }
356
357 impl TestTraceId {
358 fn from_bytes(bytes: [u8; 16]) -> Self {
359 Self { bytes }
360 }
361 }
362
363 impl TraceIdLike for TestTraceId {
364 fn to_u128(&self) -> u128 {
365 u128::from_be_bytes(self.bytes)
366 }
367 }
368
369 #[derive(Clone, Debug, PartialEq)]
370 enum TestValue {
371 String(String),
372 I64(i64),
373 F64(f64),
374 }
375
376 impl ValueLike for TestValue {
377 fn as_float(&self) -> Option<f64> {
378 match self {
379 TestValue::I64(i) => Some(*i as f64),
380 TestValue::F64(f) => Some(*f),
381 _ => None,
382 }
383 }
384
385 fn as_str(&self) -> Option<Cow<'_, str>> {
386 match self {
387 TestValue::String(s) => Some(Cow::Borrowed(s.as_str())),
388 TestValue::I64(i) => Some(Cow::Owned(i.to_string())),
389 TestValue::F64(f) => Some(Cow::Owned(f.to_string())),
390 }
391 }
392 }
393
394 #[derive(Clone, Debug)]
395 struct TestAttribute {
396 key: String,
397 value: TestValue,
398 }
399
400 impl TestAttribute {
401 fn new(key: impl Into<String>, value: impl Into<TestValue>) -> Self {
402 Self {
403 key: key.into(),
404 value: value.into(),
405 }
406 }
407 }
408
409 impl AttributeLike for TestAttribute {
410 type Value = TestValue;
411
412 fn key(&self) -> &str {
413 &self.key
414 }
415
416 fn value(&self) -> &Self::Value {
417 &self.value
418 }
419 }
420
421 impl From<&str> for TestValue {
422 fn from(s: &str) -> Self {
423 TestValue::String(s.to_string())
424 }
425 }
426
427 impl From<String> for TestValue {
428 fn from(s: String) -> Self {
429 TestValue::String(s)
430 }
431 }
432
433 struct TestSpan<'a> {
434 name: &'a str,
435 attributes: &'a [TestAttribute],
436 }
437
438 impl<'a> TestSpan<'a> {
439 fn new(name: &'a str, attributes: &'a [TestAttribute]) -> Self {
440 Self { name, attributes }
441 }
442
443 fn get_operation_name(&self) -> Cow<'_, str> {
444 if self
446 .attributes
447 .iter()
448 .any(|attr| attr.key() == HTTP_REQUEST_METHOD)
449 {
450 return Cow::Borrowed("http.client.request");
451 }
452
453 Cow::Borrowed("internal")
455 }
456 }
457
458 impl SpanProperties for TestSpan<'_> {
459 type Attribute<'b>
460 = &'b TestAttribute
461 where
462 Self: 'b;
463
464 fn operation_name(&self) -> Cow<'_, str> {
465 self.get_operation_name()
466 }
467
468 fn service(&self) -> Cow<'_, str> {
469 self.attributes
470 .iter()
471 .find(|attr| attr.key() == SERVICE_NAME)
472 .and_then(|attr| attr.value().as_str())
473 .unwrap_or(Cow::Borrowed(""))
474 }
475
476 fn env(&self) -> Cow<'_, str> {
477 self.attributes
478 .iter()
479 .find(|attr| attr.key() == "datadog.env" || attr.key() == ENV_TAG)
480 .and_then(|attr| attr.value().as_str())
481 .unwrap_or(Cow::Borrowed(""))
482 }
483
484 fn resource(&self) -> Cow<'_, str> {
485 self.attributes
486 .iter()
487 .find(|attr| attr.key() == RESOURCE_TAG)
488 .and_then(|attr| attr.value().as_str())
489 .unwrap_or(Cow::Borrowed(self.name))
490 }
491
492 fn status_code(&self) -> Option<u32> {
493 self.attributes
494 .iter()
495 .find(|attr| {
496 attr.key() == HTTP_RESPONSE_STATUS_CODE || attr.key() == HTTP_STATUS_CODE
497 })
498 .and_then(|attr| match attr.value() {
499 TestValue::I64(i) => Some(*i as u32),
500 _ => None,
501 })
502 }
503
504 fn attributes(&self) -> impl Iterator<Item = &TestAttribute> + '_ {
505 self.attributes.iter()
506 }
507
508 fn get_alternate_key<'b>(&self, key: &'b str) -> Option<Cow<'b, str>> {
509 match key {
510 HTTP_RESPONSE_STATUS_CODE => Some(Cow::Borrowed(HTTP_STATUS_CODE)),
511 HTTP_REQUEST_METHOD => Some(Cow::Borrowed("http.method")),
512 _ => None,
513 }
514 }
515 }
516
517 struct TestSamplingData<'a> {
518 is_parent_sampled: Option<bool>,
519 trace_id: &'a TestTraceId,
520 name: &'a str,
521 attributes: &'a [TestAttribute],
522 }
523
524 impl<'a> TestSamplingData<'a> {
525 fn new(
526 is_parent_sampled: Option<bool>,
527 trace_id: &'a TestTraceId,
528 name: &'a str,
529 attributes: &'a [TestAttribute],
530 ) -> Self {
531 Self {
532 is_parent_sampled,
533 trace_id,
534 name,
535 attributes,
536 }
537 }
538 }
539
540 impl SamplingData for TestSamplingData<'_> {
541 type TraceId = TestTraceId;
542 type Properties<'b>
543 = TestSpan<'b>
544 where
545 Self: 'b;
546
547 fn is_parent_sampled(&self) -> Option<bool> {
548 self.is_parent_sampled
549 }
550
551 fn trace_id(&self) -> &Self::TraceId {
552 self.trace_id
553 }
554
555 fn with_span_properties<S, T, F>(&self, s: &S, f: F) -> T
556 where
557 F: for<'b> Fn(&S, &TestSpan<'b>) -> T,
558 {
559 let span = TestSpan::new(self.name, self.attributes);
560 f(s, &span)
561 }
562 }
563
564 struct TestAttributeFactory;
565
566 impl crate::types::AttributeFactory for TestAttributeFactory {
567 type Attribute = TestAttribute;
568
569 fn create_i64(&self, key: &'static str, value: i64) -> Self::Attribute {
570 TestAttribute::new(key, TestValue::I64(value))
571 }
572
573 fn create_f64(&self, key: &'static str, value: f64) -> Self::Attribute {
574 TestAttribute::new(key, TestValue::F64(value))
575 }
576
577 fn create_string(&self, key: &'static str, value: Cow<'static, str>) -> Self::Attribute {
578 TestAttribute::new(key, TestValue::String(value.into_owned()))
579 }
580 }
581
582 fn create_trace_id() -> TestTraceId {
588 let bytes = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
589 TestTraceId::from_bytes(bytes)
590 }
591
592 fn create_attributes(resource: &'static str, env: &'static str) -> Vec<TestAttribute> {
594 vec![
595 TestAttribute::new(RESOURCE_TAG, resource),
596 TestAttribute::new("datadog.env", env),
597 ]
598 }
599
600 fn create_attributes_with_service(
602 service: String,
603 resource: &'static str,
604 env: &'static str,
605 ) -> Vec<TestAttribute> {
606 vec![
607 TestAttribute::new(SERVICE_NAME, service),
608 TestAttribute::new(RESOURCE_TAG, resource),
609 TestAttribute::new("datadog.env", env),
610 ]
611 }
612
613 fn create_attributes_with_extra(
615 service: &'static str,
616 resource: &'static str,
617 env: &'static str,
618 extra: &[(&'static str, &'static str)],
619 ) -> Vec<TestAttribute> {
620 let mut attrs = create_attributes_with_service(service.to_string(), resource, env);
621 for (k, v) in extra {
622 attrs.push(TestAttribute::new(*k, *v));
623 }
624 attrs
625 }
626
627 fn create_sampling_data<'a>(
629 is_parent_sampled: Option<bool>,
630 trace_id: &'a TestTraceId,
631 name: &'a str,
632 attributes: &'a [TestAttribute],
633 ) -> TestSamplingData<'a> {
634 TestSamplingData::new(is_parent_sampled, trace_id, name, attributes)
635 }
636
637 #[test]
638 fn test_sampling_rule_creation() {
639 let rule = SamplingRule::new(
640 0.5,
641 Some("test-service".to_string()),
642 Some("test-name".to_string()),
643 Some("test-resource".to_string()),
644 Some(HashMap::from([(
645 "custom-tag".to_string(),
646 "tag-value".to_string(),
647 )])),
648 Some("customer".to_string()),
649 );
650
651 assert_eq!(rule.sample_rate, 0.5);
652 assert_eq!(rule.service_matcher.unwrap().pattern(), "test-service");
653 assert_eq!(rule.name_matcher.unwrap().pattern(), "test-name");
654 assert_eq!(
655 rule.resource_matcher.unwrap().pattern(),
656 "test-resource".to_string()
657 );
658 assert_eq!(
659 rule.tag_matchers.get("custom-tag").unwrap().pattern(),
660 "tag-value"
661 );
662 assert_eq!(rule.provenance, "customer");
663 }
664
665 #[test]
666 fn test_sampling_rule_with_no_rule() {
667 let rule = SamplingRule::new(
669 0.5, None, None, None, None, None, );
675
676 assert_eq!(rule.sample_rate, 0.5);
678 assert!(rule.service_matcher.is_none());
679 assert!(rule.name_matcher.is_none());
680 assert!(rule.resource_matcher.is_none());
681 assert!(rule.tag_matchers.is_empty());
682 assert_eq!(rule.provenance, "default");
683
684 assert!(rule.service_matcher.is_none());
686 assert!(rule.name_matcher.is_none());
687 assert!(rule.resource_matcher.is_none());
688 assert!(rule.tag_matchers.is_empty());
689
690 let rule_with_empty_strings = SamplingRule::new(
692 0.5,
693 Some(pattern::NO_RULE.to_string()), Some(pattern::NO_RULE.to_string()), Some(pattern::NO_RULE.to_string()), Some(HashMap::from([(
697 pattern::NO_RULE.to_string(),
698 pattern::NO_RULE.to_string(),
699 )])), None,
701 );
702
703 assert!(rule_with_empty_strings.service_matcher.is_none());
705 assert!(rule_with_empty_strings.name_matcher.is_none());
706 assert!(rule_with_empty_strings.resource_matcher.is_none());
707 assert!(rule_with_empty_strings.tag_matchers.is_empty());
708
709 let attributes = create_attributes("some-resource", "some-env");
711
712 let span = TestSpan::new("", &attributes);
714 assert!(rule.matches(&span));
715 assert!(rule_with_empty_strings.matches(&span));
716 }
717
718 #[test]
719 fn test_sampling_rule_matches() {
720 let rule = SamplingRule::new(
724 0.5,
725 Some("web-*".to_string()),
726 Some("http.client.*".to_string()),
727 None,
728 Some(HashMap::from([(
729 "custom_key".to_string(),
730 "custom_value".to_string(),
731 )])),
732 None,
733 );
734
735 let attrs = create_attributes_with_extra(
737 "web-foo",
738 "resource",
739 "production",
740 &[(HTTP_REQUEST_METHOD, "GET"), ("custom_key", "custom_value")],
741 );
742 let span = TestSpan::new("span-name", attrs.as_slice());
743 assert!(rule.matches(&span), "rule should match qualifying span");
744
745 let attrs_bad_service = create_attributes_with_extra(
747 "api-foo",
748 "resource",
749 "production",
750 &[(HTTP_REQUEST_METHOD, "GET"), ("custom_key", "custom_value")],
751 );
752 let span_bad_service = TestSpan::new("span-name", attrs_bad_service.as_slice());
753 assert!(
754 !rule.matches(&span_bad_service),
755 "rule should not match different service"
756 );
757
758 let attrs_no_tag = create_attributes_with_extra(
760 "web-foo",
761 "resource",
762 "production",
763 &[(HTTP_REQUEST_METHOD, "GET")],
764 );
765 let span_no_tag = TestSpan::new("span-name", attrs_no_tag.as_slice());
766 assert!(
767 !rule.matches(&span_no_tag),
768 "rule should not match without required tag"
769 );
770 }
771
772 #[test]
773 fn test_sample_method() {
774 let rule_always = SamplingRule::new(1.0, None, None, None, None, None);
776 let rule_never = SamplingRule::new(0.0, None, None, None, None, None);
777
778 let trace_id = create_trace_id();
779
780 assert!(rule_always.sample(&trace_id));
782
783 assert!(!rule_never.sample(&trace_id));
785 }
786
787 #[test]
788 fn test_datadog_sampler_creation() {
789 let sampler = DatadogSampler::new(vec![], 100);
791 assert!(sampler.rules.is_empty());
792 assert!(sampler.service_samplers.is_empty());
793
794 let rule = SamplingRule::new(0.5, None, None, None, None, None);
796 let sampler_with_rules = DatadogSampler::new(vec![rule], 200);
797 assert_eq!(sampler_with_rules.rules.len(), 1);
798 }
799
800 #[test]
801 fn test_service_key_generation() {
802 let test_service_name = "test-service".to_string();
803 let sampler = DatadogSampler::new(vec![], 100);
804
805 let attrs =
807 create_attributes_with_service(test_service_name.clone(), "resource", "production");
808 let span = TestSpan::new("test-span", attrs.as_slice());
809 assert_eq!(
810 sampler.service_key(&span),
811 format!("service:{test_service_name},env:production")
812 );
813
814 let attrs_no_env = vec![
816 TestAttribute::new(SERVICE_NAME, test_service_name.clone()),
817 TestAttribute::new(RESOURCE_TAG, "resource"),
818 ];
819 let span = TestSpan::new("test-span", attrs_no_env.as_slice());
820 assert_eq!(
821 sampler.service_key(&span),
822 format!("service:{test_service_name},env:")
823 );
824 }
825
826 #[test]
827 fn test_update_service_rates() {
828 let sampler = DatadogSampler::new(vec![], 100);
829
830 let mut rates = HashMap::new();
832 rates.insert("service:web,env:prod".to_string(), 0.5);
833 rates.insert("service:api,env:prod".to_string(), 0.75);
834
835 sampler.service_samplers.update_rates(rates);
836
837 assert_eq!(sampler.service_samplers.len(), 2);
839
840 assert!(sampler
842 .service_samplers
843 .contains_key("service:web,env:prod"));
844 assert!(sampler
845 .service_samplers
846 .contains_key("service:api,env:prod"));
847
848 if let Some(web_sampler) = sampler.service_samplers.get("service:web,env:prod") {
850 assert_eq!(web_sampler.sample_rate(), 0.5);
851 } else {
852 panic!("Web service sampler not found");
853 }
854
855 if let Some(api_sampler) = sampler.service_samplers.get("service:api,env:prod") {
856 assert_eq!(api_sampler.sample_rate(), 0.75);
857 } else {
858 panic!("API service sampler not found");
859 }
860 }
861
862 #[test]
863 fn test_find_matching_rule() {
864 let rule1 = SamplingRule::new(
866 0.1,
867 Some("service1".to_string()),
868 None,
869 None,
870 None,
871 Some("customer".to_string()), );
873
874 let rule2 = SamplingRule::new(
875 0.2,
876 Some("service2".to_string()),
877 None,
878 None,
879 None,
880 Some("dynamic".to_string()), );
882
883 let rule3 = SamplingRule::new(
884 0.3,
885 Some("service*".to_string()), None,
887 None,
888 None,
889 Some("default".to_string()), );
891
892 let sampler = DatadogSampler::new(vec![rule1.clone(), rule2.clone(), rule3.clone()], 100);
893
894 {
896 let attrs1 = create_attributes_with_service(
897 "service1".to_string(),
898 "resource_val_for_attr1",
899 "prod",
900 );
901 let span = TestSpan::new("test-span", attrs1.as_slice());
902 let matching_rule_for_attrs1 = sampler.find_matching_rule(&span);
903 assert!(
904 matching_rule_for_attrs1.is_some(),
905 "Expected rule1 to match for service1"
906 );
907 let rule = matching_rule_for_attrs1.unwrap();
908 assert_eq!(rule.sample_rate, 0.1, "Expected rule1 sample rate");
909 assert_eq!(rule.provenance, "customer", "Expected rule1 provenance");
910 }
911
912 {
914 let attrs2 = create_attributes_with_service(
915 "service2".to_string(),
916 "resource_val_for_attr2",
917 "prod",
918 );
919 let span = TestSpan::new("test-span", attrs2.as_slice());
920 let matching_rule_for_attrs2 = sampler.find_matching_rule(&span);
921 assert!(
922 matching_rule_for_attrs2.is_some(),
923 "Expected rule2 to match for service2"
924 );
925 let rule = matching_rule_for_attrs2.unwrap();
926 assert_eq!(rule.sample_rate, 0.2, "Expected rule2 sample rate");
927 assert_eq!(rule.provenance, "dynamic", "Expected rule2 provenance");
928 }
929
930 {
932 let attrs3 = create_attributes_with_service(
933 "service3".to_string(),
934 "resource_val_for_attr3",
935 "prod",
936 );
937 let span = TestSpan::new("test-span", attrs3.as_slice());
938 let matching_rule_for_attrs3 = sampler.find_matching_rule(&span);
939 assert!(
940 matching_rule_for_attrs3.is_some(),
941 "Expected rule3 to match for service3"
942 );
943 let rule = matching_rule_for_attrs3.unwrap();
944 assert_eq!(rule.sample_rate, 0.3, "Expected rule3 sample rate");
945 assert_eq!(rule.provenance, "default", "Expected rule3 provenance");
946 }
947
948 {
950 let attrs4 = create_attributes_with_service(
951 "other_sampler_service".to_string(),
952 "resource_val_for_attr4",
953 "prod",
954 );
955 let span = TestSpan::new("test-span", attrs4.as_slice());
956 let matching_rule_for_attrs4 = sampler.find_matching_rule(&span);
957 assert!(
958 matching_rule_for_attrs4.is_none(),
959 "Expected no rule to match for service 'other_sampler_service'"
960 );
961 }
962 }
963
964 #[test]
965 fn test_get_sampling_mechanism() {
966 let sampler = DatadogSampler::new(vec![], 100);
967
968 let rule_customer =
970 SamplingRule::new(0.1, None, None, None, None, Some("customer".to_string()));
971 let rule_dynamic =
972 SamplingRule::new(0.2, None, None, None, None, Some("dynamic".to_string()));
973 let rule_default =
974 SamplingRule::new(0.3, None, None, None, None, Some("default".to_string()));
975
976 let mechanism1 = sampler.get_sampling_mechanism(Some(&rule_customer), false);
978 assert_eq!(mechanism1, mechanism::REMOTE_USER_TRACE_SAMPLING_RULE);
979
980 let mechanism2 = sampler.get_sampling_mechanism(Some(&rule_dynamic), false);
982 assert_eq!(mechanism2, mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE);
983
984 let mechanism3 = sampler.get_sampling_mechanism(Some(&rule_default), false);
986 assert_eq!(mechanism3, mechanism::LOCAL_USER_TRACE_SAMPLING_RULE);
987
988 let mechanism4 = sampler.get_sampling_mechanism(None, true);
990 assert_eq!(mechanism4, mechanism::AGENT_RATE_BY_SERVICE);
991
992 let mechanism5 = sampler.get_sampling_mechanism(None, false);
994 assert_eq!(mechanism5, mechanism::DEFAULT);
995 }
996
997 #[test]
998 fn test_add_dd_sampling_tags() {
999 let sample_rate = 0.5;
1001 let is_sampled = true;
1002 let mechanism = mechanism::LOCAL_USER_TRACE_SAMPLING_RULE;
1003 let sampling_result = DdSamplingResult {
1004 priority: mechanism.to_priority(is_sampled),
1005 trace_root_info: Some(TraceRootSamplingInfo {
1006 mechanism,
1007 rate: 0.5,
1008 rl_effective_rate: None,
1009 }),
1010 };
1011
1012 let attrs = sampling_result
1013 .to_dd_sampling_tags(&TestAttributeFactory)
1014 .unwrap_or_default();
1015
1016 assert_eq!(attrs.len(), 4);
1018
1019 let mut found_decision_maker = false;
1021 let mut found_priority = false;
1022 let mut found_rule_rate = false;
1023 let mut found_ksr = false;
1024
1025 for attr in &attrs {
1026 match attr.key() {
1027 SAMPLING_DECISION_MAKER_TAG_KEY => {
1028 let value_str = match attr.value() {
1029 TestValue::String(s) => s.to_string(),
1030 _ => panic!("Expected string value for decision maker tag"),
1031 };
1032 assert_eq!(value_str, mechanism.to_cow());
1033 found_decision_maker = true;
1034 }
1035 SAMPLING_PRIORITY_TAG_KEY => {
1036 let expected_priority = mechanism.to_priority(true).into_i8() as i64;
1038
1039 let value_int = match attr.value() {
1040 TestValue::I64(i) => *i,
1041 _ => panic!("Expected integer value for priority tag"),
1042 };
1043 assert_eq!(value_int, expected_priority);
1044 found_priority = true;
1045 }
1046 SAMPLING_RULE_RATE_TAG_KEY => {
1047 let value_float = match attr.value() {
1048 TestValue::F64(f) => *f,
1049 _ => panic!("Expected float value for rule rate tag"),
1050 };
1051 assert_eq!(value_float, sample_rate);
1052 found_rule_rate = true;
1053 }
1054 SAMPLING_KNUTH_RATE_TAG_KEY => {
1055 let value_str = match attr.value() {
1056 TestValue::String(s) => s.to_string(),
1057 _ => panic!("Expected string value for ksr tag"),
1058 };
1059 assert_eq!(value_str, "0.5");
1060 found_ksr = true;
1061 }
1062 _ => {}
1063 }
1064 }
1065
1066 assert!(found_decision_maker, "Missing decision maker tag");
1067 assert!(found_priority, "Missing priority tag");
1068 assert!(found_rule_rate, "Missing rule rate tag");
1069 assert!(found_ksr, "Missing knuth sampling rate tag");
1070
1071 let rate_limit = 0.5;
1073 let is_sampled = false;
1074 let mechanism = mechanism::LOCAL_USER_TRACE_SAMPLING_RULE;
1075 let sampling_result = DdSamplingResult {
1076 priority: mechanism.to_priority(is_sampled),
1077 trace_root_info: Some(TraceRootSamplingInfo {
1078 mechanism,
1079 rate: 0.5,
1080 rl_effective_rate: Some(rate_limit),
1081 }),
1082 };
1083 let attrs_with_limit = sampling_result
1084 .to_dd_sampling_tags(&TestAttributeFactory)
1085 .unwrap_or_default();
1086
1087 assert_eq!(attrs_with_limit.len(), 5);
1089
1090 let mut found_limit = false;
1092 for attr in &attrs_with_limit {
1093 if attr.key() == RL_EFFECTIVE_RATE {
1094 let value_float = match attr.value() {
1095 TestValue::F64(f) => *f,
1096 _ => panic!("Expected float value for rate limit tag"),
1097 };
1098 assert_eq!(value_float, rate_limit);
1099 found_limit = true;
1100 break;
1101 }
1102 }
1103
1104 assert!(found_limit, "Missing rate limit tag");
1105
1106 let agent_rate = 0.75;
1109 let is_sampled = false;
1110 let mechanism = mechanism::AGENT_RATE_BY_SERVICE;
1111 let sampling_result = DdSamplingResult {
1112 priority: mechanism.to_priority(is_sampled),
1113 trace_root_info: Some(TraceRootSamplingInfo {
1114 mechanism,
1115 rate: agent_rate,
1116 rl_effective_rate: None,
1117 }),
1118 };
1119
1120 let agent_attrs = sampling_result
1121 .to_dd_sampling_tags(&TestAttributeFactory)
1122 .unwrap_or_default();
1123
1124 assert_eq!(agent_attrs.len(), 4);
1127
1128 let mut found_agent_rate = false;
1130 let mut found_ksr = false;
1131 for attr in &agent_attrs {
1132 match attr.key() {
1133 SAMPLING_AGENT_RATE_TAG_KEY => {
1134 let value_float = match attr.value() {
1135 TestValue::F64(f) => *f,
1136 _ => panic!("Expected float value for agent rate tag"),
1137 };
1138 assert_eq!(value_float, agent_rate);
1139 found_agent_rate = true;
1140 }
1141 SAMPLING_KNUTH_RATE_TAG_KEY => {
1142 let value_str = match attr.value() {
1143 TestValue::String(s) => s.to_string(),
1144 _ => panic!("Expected string value for ksr tag"),
1145 };
1146 assert_eq!(value_str, "0.75");
1147 found_ksr = true;
1148 }
1149 _ => {}
1150 }
1151 }
1152
1153 assert!(found_agent_rate, "Missing agent rate tag");
1154 assert!(
1155 found_ksr,
1156 "Missing knuth sampling rate tag for agent mechanism"
1157 );
1158
1159 for attr in &agent_attrs {
1161 assert_ne!(
1162 attr.key(),
1163 SAMPLING_RULE_RATE_TAG_KEY,
1164 "Rule rate tag should not be present for agent mechanism"
1165 );
1166 }
1167 }
1168
1169 #[test]
1170 fn test_format_sampling_rate() {
1171 assert_eq!(format_sampling_rate(1.0), Some("1".to_string()));
1173 assert_eq!(format_sampling_rate(0.5), Some("0.5".to_string()));
1174 assert_eq!(format_sampling_rate(0.1), Some("0.1".to_string()));
1175 assert_eq!(format_sampling_rate(0.0), Some("0".to_string()));
1176
1177 assert_eq!(format_sampling_rate(0.100000), Some("0.1".to_string()));
1179 assert_eq!(format_sampling_rate(0.500000), Some("0.5".to_string()));
1180
1181 assert_eq!(
1183 format_sampling_rate(0.7654321),
1184 Some("0.765432".to_string())
1185 );
1186 assert_eq!(
1187 format_sampling_rate(0.123456789),
1188 Some("0.123457".to_string())
1189 );
1190
1191 assert_eq!(format_sampling_rate(0.001), Some("0.001".to_string()));
1193
1194 assert_eq!(format_sampling_rate(0.75), Some("0.75".to_string()));
1196 assert_eq!(format_sampling_rate(0.999999), Some("0.999999".to_string()));
1197
1198 assert_eq!(format_sampling_rate(-0.1), None);
1200 assert_eq!(format_sampling_rate(1.1), None);
1201 assert_eq!(format_sampling_rate(f64::NAN), None);
1202 assert_eq!(format_sampling_rate(f64::INFINITY), None);
1203 assert_eq!(format_sampling_rate(f64::NEG_INFINITY), None);
1204 }
1205
1206 #[test]
1207 fn test_should_sample_parent_context() {
1208 let sampler = DatadogSampler::new(vec![], 100);
1209
1210 let empty_attrs: &[TestAttribute] = &[];
1212 let trace_id = create_trace_id();
1213
1214 let data_sampled = create_sampling_data(Some(true), &trace_id, "span", empty_attrs);
1216 let result_sampled = sampler.sample(&data_sampled);
1217
1218 assert!(result_sampled.get_priority().is_keep());
1220 assert!(result_sampled
1221 .to_dd_sampling_tags(&TestAttributeFactory)
1222 .is_none());
1223
1224 let data_not_sampled = create_sampling_data(Some(false), &trace_id, "span", empty_attrs);
1226 let result_not_sampled = sampler.sample(&data_not_sampled);
1227
1228 assert!(!result_not_sampled.get_priority().is_keep());
1230 assert!(result_not_sampled
1231 .to_dd_sampling_tags(&TestAttributeFactory)
1232 .is_none());
1233 }
1234
1235 #[test]
1236 fn test_should_sample_with_rule() {
1237 let rule = SamplingRule::new(
1239 1.0,
1240 Some("test-service".to_string()),
1241 None,
1242 None,
1243 None,
1244 None,
1245 );
1246
1247 let sampler = DatadogSampler::new(vec![rule], 100);
1248
1249 let trace_id = create_trace_id();
1250
1251 let attrs = create_attributes("resource", "prod");
1253 let data = create_sampling_data(None, &trace_id, "span", attrs.as_slice());
1254 let result = sampler.sample(&data);
1255
1256 assert!(result.get_priority().is_keep());
1258 assert!(result.to_dd_sampling_tags(&TestAttributeFactory).is_some());
1259
1260 let attrs_no_match = create_attributes("other-resource", "prod");
1262 let data_no_match =
1263 create_sampling_data(None, &trace_id, "span", attrs_no_match.as_slice());
1264 let result_no_match = sampler.sample(&data_no_match);
1265
1266 assert!(result_no_match.get_priority().is_keep());
1268 assert!(result_no_match
1269 .to_dd_sampling_tags(&TestAttributeFactory)
1270 .is_some());
1271 }
1272
1273 #[test]
1274 fn test_should_sample_with_service_rates() {
1275 let sampler = DatadogSampler::new(vec![], 100);
1277
1278 let mut rates = HashMap::new();
1280 rates.insert("service:test-service,env:prod".to_string(), 1.0); rates.insert("service:other-service,env:prod".to_string(), 0.0); sampler.update_service_rates(rates);
1284
1285 let trace_id = create_trace_id();
1286
1287 let attrs_sample = create_attributes_with_service(
1289 "test-service".to_string(),
1290 "any_resource_name_matching_env",
1291 "prod",
1292 );
1293 let data_sample = create_sampling_data(
1294 None,
1295 &trace_id,
1296 "span_for_test_service",
1297 attrs_sample.as_slice(),
1298 );
1299 let result_sample = sampler.sample(&data_sample);
1300 assert!(
1303 result_sample.get_priority().is_keep(),
1304 "Span for test-service/prod should be sampled"
1305 );
1306
1307 let attrs_no_sample = create_attributes_with_service(
1309 "other-service".to_string(),
1310 "any_resource_name_matching_env",
1311 "prod",
1312 );
1313 let data_no_sample = create_sampling_data(
1314 None,
1315 &trace_id,
1316 "span_for_other_service",
1317 attrs_no_sample.as_slice(),
1318 );
1319 let result_no_sample = sampler.sample(&data_no_sample);
1320 assert!(
1322 !result_no_sample.get_priority().is_keep(),
1323 "Span for other-service/prod should be dropped"
1324 );
1325 }
1326
1327 #[test]
1328 fn test_sampling_rule_matches_float_attributes() {
1329 fn create_attributes_with_float(
1331 tag_key: &'static str,
1332 float_value: f64,
1333 ) -> Vec<TestAttribute> {
1334 vec![
1335 TestAttribute::new(RESOURCE_TAG, "resource"),
1336 TestAttribute::new(ENV_TAG, "prod"),
1337 TestAttribute::new(tag_key, TestValue::F64(float_value)),
1338 ]
1339 }
1340
1341 let rule_integer = SamplingRule::new(
1343 0.5,
1344 None,
1345 None,
1346 None,
1347 Some(HashMap::from([("float_tag".to_string(), "42".to_string())])),
1348 None,
1349 );
1350
1351 let integer_float_attrs = create_attributes_with_float("float_tag", 42.0);
1353 let span = TestSpan::new("test-span", integer_float_attrs.as_slice());
1354 assert!(rule_integer.matches(&span));
1355
1356 let rule_wildcard = SamplingRule::new(
1358 0.5,
1359 None,
1360 None,
1361 None,
1362 Some(HashMap::from([("float_tag".to_string(), "*".to_string())])),
1363 None,
1364 );
1365
1366 let decimal_float_attrs = create_attributes_with_float("float_tag", 42.5);
1368 let span = TestSpan::new("test-span", decimal_float_attrs.as_slice());
1369 assert!(rule_wildcard.matches(&span));
1370
1371 let rule_specific = SamplingRule::new(
1374 0.5,
1375 None,
1376 None,
1377 None,
1378 Some(HashMap::from([(
1379 "float_tag".to_string(),
1380 "42.5".to_string(),
1381 )])),
1382 None,
1383 );
1384
1385 let decimal_float_attrs = create_attributes_with_float("float_tag", 42.5);
1387 let span = TestSpan::new("test-span", decimal_float_attrs.as_slice());
1388 assert!(!rule_specific.matches(&span));
1389 let rule_prefix = SamplingRule::new(
1391 0.5,
1392 None,
1393 None,
1394 None,
1395 Some(HashMap::from([(
1396 "float_tag".to_string(),
1397 "42.*".to_string(),
1398 )])),
1399 None,
1400 );
1401
1402 let span = TestSpan::new("test-span", decimal_float_attrs.as_slice());
1405 assert!(!rule_prefix.matches(&span));
1406 }
1407
1408 #[test]
1409 fn test_operation_name() {
1410 let http_rule = SamplingRule::new(
1414 1.0,
1415 None,
1416 Some("http.*.request".to_string()),
1417 None,
1418 None,
1419 Some("default".to_string()),
1420 );
1421
1422 let sampler = DatadogSampler::new(vec![http_rule], 100);
1423
1424 let trace_id = create_trace_id();
1425
1426 let http_client_attrs = vec![TestAttribute::new(HTTP_REQUEST_METHOD, "GET")];
1428 let data = create_sampling_data(None, &trace_id, "test-span", &http_client_attrs);
1429 assert!(sampler.sample(&data).get_priority().is_keep());
1430
1431 let internal_attrs = vec![TestAttribute::new("custom.tag", "value")];
1433 let data = create_sampling_data(None, &trace_id, "test-span", &internal_attrs);
1434 assert!(sampler.sample(&data).get_priority().is_keep());
1435 }
1436
1437 #[test]
1438 fn test_on_rules_update_callback() {
1439 let initial_rule = SamplingRule::new(
1441 0.1,
1442 Some("initial-service".to_string()),
1443 None,
1444 None,
1445 None,
1446 Some("default".to_string()),
1447 );
1448
1449 let sampler = DatadogSampler::new(vec![initial_rule], 100);
1450
1451 assert_eq!(sampler.rules.len(), 1);
1453
1454 let callback = sampler.on_rules_update();
1456
1457 let new_rules = vec![
1459 SamplingRuleConfig {
1460 sample_rate: 0.5,
1461 service: Some("web-*".to_string()),
1462 name: Some("http.*".to_string()),
1463 resource: None,
1464 tags: std::collections::HashMap::new(),
1465 provenance: "customer".to_string(),
1466 },
1467 SamplingRuleConfig {
1468 sample_rate: 0.2,
1469 service: Some("api-*".to_string()),
1470 name: None,
1471 resource: Some("/api/*".to_string()),
1472 tags: [("env".to_string(), "prod".to_string())].into(),
1473 provenance: "dynamic".to_string(),
1474 },
1475 ];
1476
1477 callback(&new_rules);
1479
1480 assert_eq!(sampler.rules.len(), 2);
1482
1483 let attrs = vec![
1487 TestAttribute::new(SERVICE_NAME, "web-frontend"),
1488 TestAttribute::new(HTTP_REQUEST_METHOD, "GET"), ];
1491 let span = TestSpan::new("test-span", attrs.as_slice());
1492
1493 let matching_rule = sampler.find_matching_rule(&span);
1494 assert!(matching_rule.is_some(), "Expected to find a matching rule for service 'web-frontend' and name 'http.client.request'");
1495 let rule = matching_rule.unwrap();
1496 assert_eq!(rule.sample_rate, 0.5);
1497 assert_eq!(rule.provenance, "customer");
1498
1499 callback(&[]);
1501 assert_eq!(sampler.rules.len(), 0); }
1503
1504 #[test]
1505 fn test_on_agent_response_updates_service_rates() {
1506 let sampler = DatadogSampler::new(vec![], 100);
1507 let callback = sampler.on_agent_response();
1508
1509 let json = r#"{"rate_by_service":{"service:web,env:prod":0.5}}"#;
1511 callback(json);
1512 assert!(sampler
1513 .service_samplers
1514 .contains_key("service:web,env:prod"));
1515
1516 callback("not json");
1518
1519 callback(r#"{"other_field":1}"#);
1521 }
1522
1523 #[test]
1524 fn test_rate_limiter_drop_branch() {
1525 let always_keep = SamplingRule::new(1.0, None, None, None, None, None);
1528 let sampler = DatadogSampler::new(vec![always_keep], 0);
1529 let trace_id = TestTraceId::from_bytes([0u8; 16]);
1530 let attributes = create_attributes("res", "prod");
1531 let data = create_sampling_data(None, &trace_id, "op", &attributes);
1532 let decision = sampler.sample(&data);
1533 assert_eq!(
1534 decision.priority,
1535 priority::USER_REJECT,
1536 "rule kept span, rate_limit=0 should then drop it"
1537 );
1538 }
1539
1540 #[test]
1541 fn test_get_trace_root_sampling_info() {
1542 let sampler = DatadogSampler::new(vec![], 100);
1543 let trace_id = TestTraceId::from_bytes([0u8; 16]);
1544 let attributes = create_attributes("res", "prod");
1545 let data = create_sampling_data(None, &trace_id, "op", &attributes);
1546 let decision = sampler.sample(&data);
1547 let _info = decision.get_trace_root_sampling_info();
1548 }
1549}