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> {
205 if rate.is_nan() || !(0.0..=1.0).contains(&rate) {
206 return None;
207 }
208
209 let s = format!("{rate:.6}");
210 Some(s.trim_end_matches('0').trim_end_matches('.').to_string())
211}
212
213pub struct TraceRootSamplingInfo {
214 mechanism: SamplingMechanism,
215 rate: f64,
216 rl_effective_rate: Option<f64>,
217}
218
219impl TraceRootSamplingInfo {
220 pub fn mechanism(&self) -> SamplingMechanism {
222 self.mechanism
223 }
224
225 pub fn rate(&self) -> f64 {
227 self.rate
228 }
229
230 pub fn rl_effective_rate(&self) -> Option<f64> {
232 self.rl_effective_rate
233 }
234}
235
236pub struct DdSamplingResult {
237 priority: SamplingPriority,
238 trace_root_info: Option<TraceRootSamplingInfo>,
239}
240
241impl DdSamplingResult {
242 #[inline(always)]
243 pub fn get_priority(&self) -> SamplingPriority {
244 self.priority
245 }
246
247 pub fn get_trace_root_sampling_info(&self) -> &Option<TraceRootSamplingInfo> {
248 &self.trace_root_info
249 }
250
251 pub fn to_dd_sampling_tags<F>(&self, factory: &F) -> Option<Vec<F::Attribute>>
259 where
260 F: crate::types::AttributeFactory,
261 {
262 let Some(root_info) = &self.trace_root_info else {
263 return None; };
265
266 let mut result: Vec<F::Attribute>;
267 if let Some(limit) = root_info.rl_effective_rate() {
269 result = Vec::with_capacity(4);
270 result.push(factory.create_f64(RL_EFFECTIVE_RATE, limit));
271 } else {
272 result = Vec::with_capacity(3);
273 }
274
275 let mechanism = root_info.mechanism();
277 result.push(factory.create_string(SAMPLING_DECISION_MAKER_TAG_KEY, mechanism.to_cow()));
278
279 match mechanism {
281 mechanism::AGENT_RATE_BY_SERVICE => {
282 result.push(factory.create_f64(SAMPLING_AGENT_RATE_TAG_KEY, root_info.rate()));
283 if let Some(rate_str) = format_sampling_rate(root_info.rate()) {
284 result.push(factory.create_string(
285 SAMPLING_KNUTH_RATE_TAG_KEY,
286 std::borrow::Cow::Owned(rate_str),
287 ));
288 }
289 }
290 mechanism::REMOTE_USER_TRACE_SAMPLING_RULE
291 | mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE
292 | mechanism::LOCAL_USER_TRACE_SAMPLING_RULE => {
293 result.push(factory.create_f64(SAMPLING_RULE_RATE_TAG_KEY, root_info.rate()));
294 if let Some(rate_str) = format_sampling_rate(root_info.rate()) {
295 result.push(factory.create_string(
296 SAMPLING_KNUTH_RATE_TAG_KEY,
297 std::borrow::Cow::Owned(rate_str),
298 ));
299 }
300 }
301 _ => {}
302 }
303
304 let priority = self.priority;
305 result.push(factory.create_i64(SAMPLING_PRIORITY_TAG_KEY, priority.into_i8() as i64));
306
307 Some(result)
308 }
309}
310
311#[cfg(test)]
312mod tests {
313 use super::*;
314 use crate::constants::{
315 attr::{ENV_TAG, RESOURCE_TAG},
316 pattern,
317 };
318 use crate::types::{AttributeLike, TraceIdLike, ValueLike};
319 use std::borrow::Cow;
320 use std::collections::HashMap;
321
322 const HTTP_REQUEST_METHOD: &str = "http.request.method";
324 const SERVICE_NAME: &str = "service.name";
325
326 const HTTP_RESPONSE_STATUS_CODE: &str = "http.response.status_code";
328 const HTTP_STATUS_CODE: &str = "http.status_code";
329
330 #[derive(Clone, Debug, PartialEq, Eq)]
335 struct TestTraceId {
336 bytes: [u8; 16],
337 }
338
339 impl TestTraceId {
340 fn from_bytes(bytes: [u8; 16]) -> Self {
341 Self { bytes }
342 }
343 }
344
345 impl TraceIdLike for TestTraceId {
346 fn to_u128(&self) -> u128 {
347 u128::from_be_bytes(self.bytes)
348 }
349 }
350
351 #[derive(Clone, Debug, PartialEq)]
352 enum TestValue {
353 String(String),
354 I64(i64),
355 F64(f64),
356 }
357
358 impl ValueLike for TestValue {
359 fn as_float(&self) -> Option<f64> {
360 match self {
361 TestValue::I64(i) => Some(*i as f64),
362 TestValue::F64(f) => Some(*f),
363 _ => None,
364 }
365 }
366
367 fn as_str(&self) -> Option<Cow<'_, str>> {
368 match self {
369 TestValue::String(s) => Some(Cow::Borrowed(s.as_str())),
370 TestValue::I64(i) => Some(Cow::Owned(i.to_string())),
371 TestValue::F64(f) => Some(Cow::Owned(f.to_string())),
372 }
373 }
374 }
375
376 #[derive(Clone, Debug)]
377 struct TestAttribute {
378 key: String,
379 value: TestValue,
380 }
381
382 impl TestAttribute {
383 fn new(key: impl Into<String>, value: impl Into<TestValue>) -> Self {
384 Self {
385 key: key.into(),
386 value: value.into(),
387 }
388 }
389 }
390
391 impl AttributeLike for TestAttribute {
392 type Value = TestValue;
393
394 fn key(&self) -> &str {
395 &self.key
396 }
397
398 fn value(&self) -> &Self::Value {
399 &self.value
400 }
401 }
402
403 impl From<&str> for TestValue {
404 fn from(s: &str) -> Self {
405 TestValue::String(s.to_string())
406 }
407 }
408
409 impl From<String> for TestValue {
410 fn from(s: String) -> Self {
411 TestValue::String(s)
412 }
413 }
414
415 struct TestSpan<'a> {
416 name: &'a str,
417 attributes: &'a [TestAttribute],
418 }
419
420 impl<'a> TestSpan<'a> {
421 fn new(name: &'a str, attributes: &'a [TestAttribute]) -> Self {
422 Self { name, attributes }
423 }
424
425 fn get_operation_name(&self) -> Cow<'_, str> {
426 if self
428 .attributes
429 .iter()
430 .any(|attr| attr.key() == HTTP_REQUEST_METHOD)
431 {
432 return Cow::Borrowed("http.client.request");
433 }
434
435 Cow::Borrowed("internal")
437 }
438 }
439
440 impl SpanProperties for TestSpan<'_> {
441 type Attribute<'b>
442 = &'b TestAttribute
443 where
444 Self: 'b;
445
446 fn operation_name(&self) -> Cow<'_, str> {
447 self.get_operation_name()
448 }
449
450 fn service(&self) -> Cow<'_, str> {
451 self.attributes
452 .iter()
453 .find(|attr| attr.key() == SERVICE_NAME)
454 .and_then(|attr| attr.value().as_str())
455 .unwrap_or(Cow::Borrowed(""))
456 }
457
458 fn env(&self) -> Cow<'_, str> {
459 self.attributes
460 .iter()
461 .find(|attr| attr.key() == "datadog.env" || attr.key() == ENV_TAG)
462 .and_then(|attr| attr.value().as_str())
463 .unwrap_or(Cow::Borrowed(""))
464 }
465
466 fn resource(&self) -> Cow<'_, str> {
467 self.attributes
468 .iter()
469 .find(|attr| attr.key() == RESOURCE_TAG)
470 .and_then(|attr| attr.value().as_str())
471 .unwrap_or(Cow::Borrowed(self.name))
472 }
473
474 fn status_code(&self) -> Option<u32> {
475 self.attributes
476 .iter()
477 .find(|attr| {
478 attr.key() == HTTP_RESPONSE_STATUS_CODE || attr.key() == HTTP_STATUS_CODE
479 })
480 .and_then(|attr| match attr.value() {
481 TestValue::I64(i) => Some(*i as u32),
482 _ => None,
483 })
484 }
485
486 fn attributes(&self) -> impl Iterator<Item = &TestAttribute> + '_ {
487 self.attributes.iter()
488 }
489
490 fn get_alternate_key<'b>(&self, key: &'b str) -> Option<Cow<'b, str>> {
491 match key {
492 HTTP_RESPONSE_STATUS_CODE => Some(Cow::Borrowed(HTTP_STATUS_CODE)),
493 HTTP_REQUEST_METHOD => Some(Cow::Borrowed("http.method")),
494 _ => None,
495 }
496 }
497 }
498
499 struct TestSamplingData<'a> {
500 is_parent_sampled: Option<bool>,
501 trace_id: &'a TestTraceId,
502 name: &'a str,
503 attributes: &'a [TestAttribute],
504 }
505
506 impl<'a> TestSamplingData<'a> {
507 fn new(
508 is_parent_sampled: Option<bool>,
509 trace_id: &'a TestTraceId,
510 name: &'a str,
511 attributes: &'a [TestAttribute],
512 ) -> Self {
513 Self {
514 is_parent_sampled,
515 trace_id,
516 name,
517 attributes,
518 }
519 }
520 }
521
522 impl SamplingData for TestSamplingData<'_> {
523 type TraceId = TestTraceId;
524 type Properties<'b>
525 = TestSpan<'b>
526 where
527 Self: 'b;
528
529 fn is_parent_sampled(&self) -> Option<bool> {
530 self.is_parent_sampled
531 }
532
533 fn trace_id(&self) -> &Self::TraceId {
534 self.trace_id
535 }
536
537 fn with_span_properties<S, T, F>(&self, s: &S, f: F) -> T
538 where
539 F: for<'b> Fn(&S, &TestSpan<'b>) -> T,
540 {
541 let span = TestSpan::new(self.name, self.attributes);
542 f(s, &span)
543 }
544 }
545
546 struct TestAttributeFactory;
547
548 impl crate::types::AttributeFactory for TestAttributeFactory {
549 type Attribute = TestAttribute;
550
551 fn create_i64(&self, key: &'static str, value: i64) -> Self::Attribute {
552 TestAttribute::new(key, TestValue::I64(value))
553 }
554
555 fn create_f64(&self, key: &'static str, value: f64) -> Self::Attribute {
556 TestAttribute::new(key, TestValue::F64(value))
557 }
558
559 fn create_string(&self, key: &'static str, value: Cow<'static, str>) -> Self::Attribute {
560 TestAttribute::new(key, TestValue::String(value.into_owned()))
561 }
562 }
563
564 fn create_trace_id() -> TestTraceId {
570 let bytes = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
571 TestTraceId::from_bytes(bytes)
572 }
573
574 fn create_attributes(resource: &'static str, env: &'static str) -> Vec<TestAttribute> {
576 vec![
577 TestAttribute::new(RESOURCE_TAG, resource),
578 TestAttribute::new("datadog.env", env),
579 ]
580 }
581
582 fn create_attributes_with_service(
584 service: String,
585 resource: &'static str,
586 env: &'static str,
587 ) -> Vec<TestAttribute> {
588 vec![
589 TestAttribute::new(SERVICE_NAME, service),
590 TestAttribute::new(RESOURCE_TAG, resource),
591 TestAttribute::new("datadog.env", env),
592 ]
593 }
594
595 fn create_attributes_with_extra(
597 service: &'static str,
598 resource: &'static str,
599 env: &'static str,
600 extra: &[(&'static str, &'static str)],
601 ) -> Vec<TestAttribute> {
602 let mut attrs = create_attributes_with_service(service.to_string(), resource, env);
603 for (k, v) in extra {
604 attrs.push(TestAttribute::new(*k, *v));
605 }
606 attrs
607 }
608
609 fn create_sampling_data<'a>(
611 is_parent_sampled: Option<bool>,
612 trace_id: &'a TestTraceId,
613 name: &'a str,
614 attributes: &'a [TestAttribute],
615 ) -> TestSamplingData<'a> {
616 TestSamplingData::new(is_parent_sampled, trace_id, name, attributes)
617 }
618
619 #[test]
620 fn test_sampling_rule_creation() {
621 let rule = SamplingRule::new(
622 0.5,
623 Some("test-service".to_string()),
624 Some("test-name".to_string()),
625 Some("test-resource".to_string()),
626 Some(HashMap::from([(
627 "custom-tag".to_string(),
628 "tag-value".to_string(),
629 )])),
630 Some("customer".to_string()),
631 );
632
633 assert_eq!(rule.sample_rate, 0.5);
634 assert_eq!(rule.service_matcher.unwrap().pattern(), "test-service");
635 assert_eq!(rule.name_matcher.unwrap().pattern(), "test-name");
636 assert_eq!(
637 rule.resource_matcher.unwrap().pattern(),
638 "test-resource".to_string()
639 );
640 assert_eq!(
641 rule.tag_matchers.get("custom-tag").unwrap().pattern(),
642 "tag-value"
643 );
644 assert_eq!(rule.provenance, "customer");
645 }
646
647 #[test]
648 fn test_sampling_rule_with_no_rule() {
649 let rule = SamplingRule::new(
651 0.5, None, None, None, None, None, );
657
658 assert_eq!(rule.sample_rate, 0.5);
660 assert!(rule.service_matcher.is_none());
661 assert!(rule.name_matcher.is_none());
662 assert!(rule.resource_matcher.is_none());
663 assert!(rule.tag_matchers.is_empty());
664 assert_eq!(rule.provenance, "default");
665
666 assert!(rule.service_matcher.is_none());
668 assert!(rule.name_matcher.is_none());
669 assert!(rule.resource_matcher.is_none());
670 assert!(rule.tag_matchers.is_empty());
671
672 let rule_with_empty_strings = SamplingRule::new(
674 0.5,
675 Some(pattern::NO_RULE.to_string()), Some(pattern::NO_RULE.to_string()), Some(pattern::NO_RULE.to_string()), Some(HashMap::from([(
679 pattern::NO_RULE.to_string(),
680 pattern::NO_RULE.to_string(),
681 )])), None,
683 );
684
685 assert!(rule_with_empty_strings.service_matcher.is_none());
687 assert!(rule_with_empty_strings.name_matcher.is_none());
688 assert!(rule_with_empty_strings.resource_matcher.is_none());
689 assert!(rule_with_empty_strings.tag_matchers.is_empty());
690
691 let attributes = create_attributes("some-resource", "some-env");
693
694 let span = TestSpan::new("", &attributes);
696 assert!(rule.matches(&span));
697 assert!(rule_with_empty_strings.matches(&span));
698 }
699
700 #[test]
701 fn test_sampling_rule_matches() {
702 let rule = SamplingRule::new(
706 0.5,
707 Some("web-*".to_string()),
708 Some("http.client.*".to_string()),
709 None,
710 Some(HashMap::from([(
711 "custom_key".to_string(),
712 "custom_value".to_string(),
713 )])),
714 None,
715 );
716
717 let attrs = create_attributes_with_extra(
719 "web-foo",
720 "resource",
721 "production",
722 &[(HTTP_REQUEST_METHOD, "GET"), ("custom_key", "custom_value")],
723 );
724 let span = TestSpan::new("span-name", attrs.as_slice());
725 assert!(rule.matches(&span), "rule should match qualifying span");
726
727 let attrs_bad_service = create_attributes_with_extra(
729 "api-foo",
730 "resource",
731 "production",
732 &[(HTTP_REQUEST_METHOD, "GET"), ("custom_key", "custom_value")],
733 );
734 let span_bad_service = TestSpan::new("span-name", attrs_bad_service.as_slice());
735 assert!(
736 !rule.matches(&span_bad_service),
737 "rule should not match different service"
738 );
739
740 let attrs_no_tag = create_attributes_with_extra(
742 "web-foo",
743 "resource",
744 "production",
745 &[(HTTP_REQUEST_METHOD, "GET")],
746 );
747 let span_no_tag = TestSpan::new("span-name", attrs_no_tag.as_slice());
748 assert!(
749 !rule.matches(&span_no_tag),
750 "rule should not match without required tag"
751 );
752 }
753
754 #[test]
755 fn test_sample_method() {
756 let rule_always = SamplingRule::new(1.0, None, None, None, None, None);
758 let rule_never = SamplingRule::new(0.0, None, None, None, None, None);
759
760 let trace_id = create_trace_id();
761
762 assert!(rule_always.sample(&trace_id));
764
765 assert!(!rule_never.sample(&trace_id));
767 }
768
769 #[test]
770 fn test_datadog_sampler_creation() {
771 let sampler = DatadogSampler::new(vec![], 100);
773 assert!(sampler.rules.is_empty());
774 assert!(sampler.service_samplers.is_empty());
775
776 let rule = SamplingRule::new(0.5, None, None, None, None, None);
778 let sampler_with_rules = DatadogSampler::new(vec![rule], 200);
779 assert_eq!(sampler_with_rules.rules.len(), 1);
780 }
781
782 #[test]
783 fn test_service_key_generation() {
784 let test_service_name = "test-service".to_string();
785 let sampler = DatadogSampler::new(vec![], 100);
786
787 let attrs =
789 create_attributes_with_service(test_service_name.clone(), "resource", "production");
790 let span = TestSpan::new("test-span", attrs.as_slice());
791 assert_eq!(
792 sampler.service_key(&span),
793 format!("service:{test_service_name},env:production")
794 );
795
796 let attrs_no_env = vec![
798 TestAttribute::new(SERVICE_NAME, test_service_name.clone()),
799 TestAttribute::new(RESOURCE_TAG, "resource"),
800 ];
801 let span = TestSpan::new("test-span", attrs_no_env.as_slice());
802 assert_eq!(
803 sampler.service_key(&span),
804 format!("service:{test_service_name},env:")
805 );
806 }
807
808 #[test]
809 fn test_update_service_rates() {
810 let sampler = DatadogSampler::new(vec![], 100);
811
812 let mut rates = HashMap::new();
814 rates.insert("service:web,env:prod".to_string(), 0.5);
815 rates.insert("service:api,env:prod".to_string(), 0.75);
816
817 sampler.service_samplers.update_rates(rates);
818
819 assert_eq!(sampler.service_samplers.len(), 2);
821
822 assert!(sampler
824 .service_samplers
825 .contains_key("service:web,env:prod"));
826 assert!(sampler
827 .service_samplers
828 .contains_key("service:api,env:prod"));
829
830 if let Some(web_sampler) = sampler.service_samplers.get("service:web,env:prod") {
832 assert_eq!(web_sampler.sample_rate(), 0.5);
833 } else {
834 panic!("Web service sampler not found");
835 }
836
837 if let Some(api_sampler) = sampler.service_samplers.get("service:api,env:prod") {
838 assert_eq!(api_sampler.sample_rate(), 0.75);
839 } else {
840 panic!("API service sampler not found");
841 }
842 }
843
844 #[test]
845 fn test_find_matching_rule() {
846 let rule1 = SamplingRule::new(
848 0.1,
849 Some("service1".to_string()),
850 None,
851 None,
852 None,
853 Some("customer".to_string()), );
855
856 let rule2 = SamplingRule::new(
857 0.2,
858 Some("service2".to_string()),
859 None,
860 None,
861 None,
862 Some("dynamic".to_string()), );
864
865 let rule3 = SamplingRule::new(
866 0.3,
867 Some("service*".to_string()), None,
869 None,
870 None,
871 Some("default".to_string()), );
873
874 let sampler = DatadogSampler::new(vec![rule1.clone(), rule2.clone(), rule3.clone()], 100);
875
876 {
878 let attrs1 = create_attributes_with_service(
879 "service1".to_string(),
880 "resource_val_for_attr1",
881 "prod",
882 );
883 let span = TestSpan::new("test-span", attrs1.as_slice());
884 let matching_rule_for_attrs1 = sampler.find_matching_rule(&span);
885 assert!(
886 matching_rule_for_attrs1.is_some(),
887 "Expected rule1 to match for service1"
888 );
889 let rule = matching_rule_for_attrs1.unwrap();
890 assert_eq!(rule.sample_rate, 0.1, "Expected rule1 sample rate");
891 assert_eq!(rule.provenance, "customer", "Expected rule1 provenance");
892 }
893
894 {
896 let attrs2 = create_attributes_with_service(
897 "service2".to_string(),
898 "resource_val_for_attr2",
899 "prod",
900 );
901 let span = TestSpan::new("test-span", attrs2.as_slice());
902 let matching_rule_for_attrs2 = sampler.find_matching_rule(&span);
903 assert!(
904 matching_rule_for_attrs2.is_some(),
905 "Expected rule2 to match for service2"
906 );
907 let rule = matching_rule_for_attrs2.unwrap();
908 assert_eq!(rule.sample_rate, 0.2, "Expected rule2 sample rate");
909 assert_eq!(rule.provenance, "dynamic", "Expected rule2 provenance");
910 }
911
912 {
914 let attrs3 = create_attributes_with_service(
915 "service3".to_string(),
916 "resource_val_for_attr3",
917 "prod",
918 );
919 let span = TestSpan::new("test-span", attrs3.as_slice());
920 let matching_rule_for_attrs3 = sampler.find_matching_rule(&span);
921 assert!(
922 matching_rule_for_attrs3.is_some(),
923 "Expected rule3 to match for service3"
924 );
925 let rule = matching_rule_for_attrs3.unwrap();
926 assert_eq!(rule.sample_rate, 0.3, "Expected rule3 sample rate");
927 assert_eq!(rule.provenance, "default", "Expected rule3 provenance");
928 }
929
930 {
932 let attrs4 = create_attributes_with_service(
933 "other_sampler_service".to_string(),
934 "resource_val_for_attr4",
935 "prod",
936 );
937 let span = TestSpan::new("test-span", attrs4.as_slice());
938 let matching_rule_for_attrs4 = sampler.find_matching_rule(&span);
939 assert!(
940 matching_rule_for_attrs4.is_none(),
941 "Expected no rule to match for service 'other_sampler_service'"
942 );
943 }
944 }
945
946 #[test]
947 fn test_get_sampling_mechanism() {
948 let sampler = DatadogSampler::new(vec![], 100);
949
950 let rule_customer =
952 SamplingRule::new(0.1, None, None, None, None, Some("customer".to_string()));
953 let rule_dynamic =
954 SamplingRule::new(0.2, None, None, None, None, Some("dynamic".to_string()));
955 let rule_default =
956 SamplingRule::new(0.3, None, None, None, None, Some("default".to_string()));
957
958 let mechanism1 = sampler.get_sampling_mechanism(Some(&rule_customer), false);
960 assert_eq!(mechanism1, mechanism::REMOTE_USER_TRACE_SAMPLING_RULE);
961
962 let mechanism2 = sampler.get_sampling_mechanism(Some(&rule_dynamic), false);
964 assert_eq!(mechanism2, mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE);
965
966 let mechanism3 = sampler.get_sampling_mechanism(Some(&rule_default), false);
968 assert_eq!(mechanism3, mechanism::LOCAL_USER_TRACE_SAMPLING_RULE);
969
970 let mechanism4 = sampler.get_sampling_mechanism(None, true);
972 assert_eq!(mechanism4, mechanism::AGENT_RATE_BY_SERVICE);
973
974 let mechanism5 = sampler.get_sampling_mechanism(None, false);
976 assert_eq!(mechanism5, mechanism::DEFAULT);
977 }
978
979 #[test]
980 fn test_add_dd_sampling_tags() {
981 let sample_rate = 0.5;
983 let is_sampled = true;
984 let mechanism = mechanism::LOCAL_USER_TRACE_SAMPLING_RULE;
985 let sampling_result = DdSamplingResult {
986 priority: mechanism.to_priority(is_sampled),
987 trace_root_info: Some(TraceRootSamplingInfo {
988 mechanism,
989 rate: 0.5,
990 rl_effective_rate: None,
991 }),
992 };
993
994 let attrs = sampling_result
995 .to_dd_sampling_tags(&TestAttributeFactory)
996 .unwrap_or_default();
997
998 assert_eq!(attrs.len(), 4);
1000
1001 let mut found_decision_maker = false;
1003 let mut found_priority = false;
1004 let mut found_rule_rate = false;
1005 let mut found_ksr = false;
1006
1007 for attr in &attrs {
1008 match attr.key() {
1009 SAMPLING_DECISION_MAKER_TAG_KEY => {
1010 let value_str = match attr.value() {
1011 TestValue::String(s) => s.to_string(),
1012 _ => panic!("Expected string value for decision maker tag"),
1013 };
1014 assert_eq!(value_str, mechanism.to_cow());
1015 found_decision_maker = true;
1016 }
1017 SAMPLING_PRIORITY_TAG_KEY => {
1018 let expected_priority = mechanism.to_priority(true).into_i8() as i64;
1020
1021 let value_int = match attr.value() {
1022 TestValue::I64(i) => *i,
1023 _ => panic!("Expected integer value for priority tag"),
1024 };
1025 assert_eq!(value_int, expected_priority);
1026 found_priority = true;
1027 }
1028 SAMPLING_RULE_RATE_TAG_KEY => {
1029 let value_float = match attr.value() {
1030 TestValue::F64(f) => *f,
1031 _ => panic!("Expected float value for rule rate tag"),
1032 };
1033 assert_eq!(value_float, sample_rate);
1034 found_rule_rate = true;
1035 }
1036 SAMPLING_KNUTH_RATE_TAG_KEY => {
1037 let value_str = match attr.value() {
1038 TestValue::String(s) => s.to_string(),
1039 _ => panic!("Expected string value for ksr tag"),
1040 };
1041 assert_eq!(value_str, "0.5");
1042 found_ksr = true;
1043 }
1044 _ => {}
1045 }
1046 }
1047
1048 assert!(found_decision_maker, "Missing decision maker tag");
1049 assert!(found_priority, "Missing priority tag");
1050 assert!(found_rule_rate, "Missing rule rate tag");
1051 assert!(found_ksr, "Missing knuth sampling rate tag");
1052
1053 let rate_limit = 0.5;
1055 let is_sampled = false;
1056 let mechanism = mechanism::LOCAL_USER_TRACE_SAMPLING_RULE;
1057 let sampling_result = DdSamplingResult {
1058 priority: mechanism.to_priority(is_sampled),
1059 trace_root_info: Some(TraceRootSamplingInfo {
1060 mechanism,
1061 rate: 0.5,
1062 rl_effective_rate: Some(rate_limit),
1063 }),
1064 };
1065 let attrs_with_limit = sampling_result
1066 .to_dd_sampling_tags(&TestAttributeFactory)
1067 .unwrap_or_default();
1068
1069 assert_eq!(attrs_with_limit.len(), 5);
1071
1072 let mut found_limit = false;
1074 for attr in &attrs_with_limit {
1075 if attr.key() == RL_EFFECTIVE_RATE {
1076 let value_float = match attr.value() {
1077 TestValue::F64(f) => *f,
1078 _ => panic!("Expected float value for rate limit tag"),
1079 };
1080 assert_eq!(value_float, rate_limit);
1081 found_limit = true;
1082 break;
1083 }
1084 }
1085
1086 assert!(found_limit, "Missing rate limit tag");
1087
1088 let agent_rate = 0.75;
1091 let is_sampled = false;
1092 let mechanism = mechanism::AGENT_RATE_BY_SERVICE;
1093 let sampling_result = DdSamplingResult {
1094 priority: mechanism.to_priority(is_sampled),
1095 trace_root_info: Some(TraceRootSamplingInfo {
1096 mechanism,
1097 rate: agent_rate,
1098 rl_effective_rate: None,
1099 }),
1100 };
1101
1102 let agent_attrs = sampling_result
1103 .to_dd_sampling_tags(&TestAttributeFactory)
1104 .unwrap_or_default();
1105
1106 assert_eq!(agent_attrs.len(), 4);
1109
1110 let mut found_agent_rate = false;
1112 let mut found_ksr = false;
1113 for attr in &agent_attrs {
1114 match attr.key() {
1115 SAMPLING_AGENT_RATE_TAG_KEY => {
1116 let value_float = match attr.value() {
1117 TestValue::F64(f) => *f,
1118 _ => panic!("Expected float value for agent rate tag"),
1119 };
1120 assert_eq!(value_float, agent_rate);
1121 found_agent_rate = true;
1122 }
1123 SAMPLING_KNUTH_RATE_TAG_KEY => {
1124 let value_str = match attr.value() {
1125 TestValue::String(s) => s.to_string(),
1126 _ => panic!("Expected string value for ksr tag"),
1127 };
1128 assert_eq!(value_str, "0.75");
1129 found_ksr = true;
1130 }
1131 _ => {}
1132 }
1133 }
1134
1135 assert!(found_agent_rate, "Missing agent rate tag");
1136 assert!(
1137 found_ksr,
1138 "Missing knuth sampling rate tag for agent mechanism"
1139 );
1140
1141 for attr in &agent_attrs {
1143 assert_ne!(
1144 attr.key(),
1145 SAMPLING_RULE_RATE_TAG_KEY,
1146 "Rule rate tag should not be present for agent mechanism"
1147 );
1148 }
1149 }
1150
1151 #[test]
1152 fn test_format_sampling_rate() {
1153 assert_eq!(format_sampling_rate(1.0), Some("1".to_string()));
1155 assert_eq!(format_sampling_rate(0.5), Some("0.5".to_string()));
1156 assert_eq!(format_sampling_rate(0.1), Some("0.1".to_string()));
1157 assert_eq!(format_sampling_rate(0.0), Some("0".to_string()));
1158
1159 assert_eq!(format_sampling_rate(0.100000), Some("0.1".to_string()));
1161 assert_eq!(format_sampling_rate(0.500000), Some("0.5".to_string()));
1162
1163 assert_eq!(
1165 format_sampling_rate(0.7654321),
1166 Some("0.765432".to_string())
1167 );
1168 assert_eq!(
1169 format_sampling_rate(0.123456789),
1170 Some("0.123457".to_string())
1171 );
1172
1173 assert_eq!(format_sampling_rate(0.001), Some("0.001".to_string()));
1175 assert_eq!(format_sampling_rate(0.000001), Some("0.000001".to_string()));
1176
1177 assert_eq!(format_sampling_rate(0.0000001), Some("0".to_string()));
1179 assert_eq!(
1181 format_sampling_rate(0.00000051),
1182 Some("0.000001".to_string())
1183 );
1184
1185 assert_eq!(format_sampling_rate(0.75), Some("0.75".to_string()));
1187 assert_eq!(format_sampling_rate(0.999999), Some("0.999999".to_string()));
1188
1189 assert_eq!(format_sampling_rate(-0.1), None);
1191 assert_eq!(format_sampling_rate(1.1), None);
1192 assert_eq!(format_sampling_rate(f64::NAN), None);
1193 assert_eq!(format_sampling_rate(f64::INFINITY), None);
1194 assert_eq!(format_sampling_rate(f64::NEG_INFINITY), None);
1195 }
1196
1197 #[test]
1198 fn test_should_sample_parent_context() {
1199 let sampler = DatadogSampler::new(vec![], 100);
1200
1201 let empty_attrs: &[TestAttribute] = &[];
1203 let trace_id = create_trace_id();
1204
1205 let data_sampled = create_sampling_data(Some(true), &trace_id, "span", empty_attrs);
1207 let result_sampled = sampler.sample(&data_sampled);
1208
1209 assert!(result_sampled.get_priority().is_keep());
1211 assert!(result_sampled
1212 .to_dd_sampling_tags(&TestAttributeFactory)
1213 .is_none());
1214
1215 let data_not_sampled = create_sampling_data(Some(false), &trace_id, "span", empty_attrs);
1217 let result_not_sampled = sampler.sample(&data_not_sampled);
1218
1219 assert!(!result_not_sampled.get_priority().is_keep());
1221 assert!(result_not_sampled
1222 .to_dd_sampling_tags(&TestAttributeFactory)
1223 .is_none());
1224 }
1225
1226 #[test]
1227 fn test_should_sample_with_rule() {
1228 let rule = SamplingRule::new(
1230 1.0,
1231 Some("test-service".to_string()),
1232 None,
1233 None,
1234 None,
1235 None,
1236 );
1237
1238 let sampler = DatadogSampler::new(vec![rule], 100);
1239
1240 let trace_id = create_trace_id();
1241
1242 let attrs = create_attributes("resource", "prod");
1244 let data = create_sampling_data(None, &trace_id, "span", attrs.as_slice());
1245 let result = sampler.sample(&data);
1246
1247 assert!(result.get_priority().is_keep());
1249 assert!(result.to_dd_sampling_tags(&TestAttributeFactory).is_some());
1250
1251 let attrs_no_match = create_attributes("other-resource", "prod");
1253 let data_no_match =
1254 create_sampling_data(None, &trace_id, "span", attrs_no_match.as_slice());
1255 let result_no_match = sampler.sample(&data_no_match);
1256
1257 assert!(result_no_match.get_priority().is_keep());
1259 assert!(result_no_match
1260 .to_dd_sampling_tags(&TestAttributeFactory)
1261 .is_some());
1262 }
1263
1264 #[test]
1265 fn test_should_sample_with_service_rates() {
1266 let sampler = DatadogSampler::new(vec![], 100);
1268
1269 let mut rates = HashMap::new();
1271 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);
1275
1276 let trace_id = create_trace_id();
1277
1278 let attrs_sample = create_attributes_with_service(
1280 "test-service".to_string(),
1281 "any_resource_name_matching_env",
1282 "prod",
1283 );
1284 let data_sample = create_sampling_data(
1285 None,
1286 &trace_id,
1287 "span_for_test_service",
1288 attrs_sample.as_slice(),
1289 );
1290 let result_sample = sampler.sample(&data_sample);
1291 assert!(
1294 result_sample.get_priority().is_keep(),
1295 "Span for test-service/prod should be sampled"
1296 );
1297
1298 let attrs_no_sample = create_attributes_with_service(
1300 "other-service".to_string(),
1301 "any_resource_name_matching_env",
1302 "prod",
1303 );
1304 let data_no_sample = create_sampling_data(
1305 None,
1306 &trace_id,
1307 "span_for_other_service",
1308 attrs_no_sample.as_slice(),
1309 );
1310 let result_no_sample = sampler.sample(&data_no_sample);
1311 assert!(
1313 !result_no_sample.get_priority().is_keep(),
1314 "Span for other-service/prod should be dropped"
1315 );
1316 }
1317
1318 #[test]
1319 fn test_sampling_rule_matches_float_attributes() {
1320 fn create_attributes_with_float(
1322 tag_key: &'static str,
1323 float_value: f64,
1324 ) -> Vec<TestAttribute> {
1325 vec![
1326 TestAttribute::new(RESOURCE_TAG, "resource"),
1327 TestAttribute::new(ENV_TAG, "prod"),
1328 TestAttribute::new(tag_key, TestValue::F64(float_value)),
1329 ]
1330 }
1331
1332 let rule_integer = SamplingRule::new(
1334 0.5,
1335 None,
1336 None,
1337 None,
1338 Some(HashMap::from([("float_tag".to_string(), "42".to_string())])),
1339 None,
1340 );
1341
1342 let integer_float_attrs = create_attributes_with_float("float_tag", 42.0);
1344 let span = TestSpan::new("test-span", integer_float_attrs.as_slice());
1345 assert!(rule_integer.matches(&span));
1346
1347 let rule_wildcard = SamplingRule::new(
1349 0.5,
1350 None,
1351 None,
1352 None,
1353 Some(HashMap::from([("float_tag".to_string(), "*".to_string())])),
1354 None,
1355 );
1356
1357 let decimal_float_attrs = create_attributes_with_float("float_tag", 42.5);
1359 let span = TestSpan::new("test-span", decimal_float_attrs.as_slice());
1360 assert!(rule_wildcard.matches(&span));
1361
1362 let rule_specific = SamplingRule::new(
1365 0.5,
1366 None,
1367 None,
1368 None,
1369 Some(HashMap::from([(
1370 "float_tag".to_string(),
1371 "42.5".to_string(),
1372 )])),
1373 None,
1374 );
1375
1376 let decimal_float_attrs = create_attributes_with_float("float_tag", 42.5);
1378 let span = TestSpan::new("test-span", decimal_float_attrs.as_slice());
1379 assert!(!rule_specific.matches(&span));
1380 let rule_prefix = SamplingRule::new(
1382 0.5,
1383 None,
1384 None,
1385 None,
1386 Some(HashMap::from([(
1387 "float_tag".to_string(),
1388 "42.*".to_string(),
1389 )])),
1390 None,
1391 );
1392
1393 let span = TestSpan::new("test-span", decimal_float_attrs.as_slice());
1396 assert!(!rule_prefix.matches(&span));
1397 }
1398
1399 #[test]
1400 fn test_operation_name() {
1401 let http_rule = SamplingRule::new(
1405 1.0,
1406 None,
1407 Some("http.*.request".to_string()),
1408 None,
1409 None,
1410 Some("default".to_string()),
1411 );
1412
1413 let sampler = DatadogSampler::new(vec![http_rule], 100);
1414
1415 let trace_id = create_trace_id();
1416
1417 let http_client_attrs = vec![TestAttribute::new(HTTP_REQUEST_METHOD, "GET")];
1419 let data = create_sampling_data(None, &trace_id, "test-span", &http_client_attrs);
1420 assert!(sampler.sample(&data).get_priority().is_keep());
1421
1422 let internal_attrs = vec![TestAttribute::new("custom.tag", "value")];
1424 let data = create_sampling_data(None, &trace_id, "test-span", &internal_attrs);
1425 assert!(sampler.sample(&data).get_priority().is_keep());
1426 }
1427
1428 #[test]
1429 fn test_on_rules_update_callback() {
1430 let initial_rule = SamplingRule::new(
1432 0.1,
1433 Some("initial-service".to_string()),
1434 None,
1435 None,
1436 None,
1437 Some("default".to_string()),
1438 );
1439
1440 let sampler = DatadogSampler::new(vec![initial_rule], 100);
1441
1442 assert_eq!(sampler.rules.len(), 1);
1444
1445 let callback = sampler.on_rules_update();
1447
1448 let new_rules = vec![
1450 SamplingRuleConfig {
1451 sample_rate: 0.5,
1452 service: Some("web-*".to_string()),
1453 name: Some("http.*".to_string()),
1454 resource: None,
1455 tags: std::collections::HashMap::new(),
1456 provenance: "customer".to_string(),
1457 },
1458 SamplingRuleConfig {
1459 sample_rate: 0.2,
1460 service: Some("api-*".to_string()),
1461 name: None,
1462 resource: Some("/api/*".to_string()),
1463 tags: [("env".to_string(), "prod".to_string())].into(),
1464 provenance: "dynamic".to_string(),
1465 },
1466 ];
1467
1468 callback(&new_rules);
1470
1471 assert_eq!(sampler.rules.len(), 2);
1473
1474 let attrs = vec![
1478 TestAttribute::new(SERVICE_NAME, "web-frontend"),
1479 TestAttribute::new(HTTP_REQUEST_METHOD, "GET"), ];
1482 let span = TestSpan::new("test-span", attrs.as_slice());
1483
1484 let matching_rule = sampler.find_matching_rule(&span);
1485 assert!(matching_rule.is_some(), "Expected to find a matching rule for service 'web-frontend' and name 'http.client.request'");
1486 let rule = matching_rule.unwrap();
1487 assert_eq!(rule.sample_rate, 0.5);
1488 assert_eq!(rule.provenance, "customer");
1489
1490 callback(&[]);
1492 assert_eq!(sampler.rules.len(), 0); }
1494
1495 #[test]
1496 fn test_on_agent_response_updates_service_rates() {
1497 let sampler = DatadogSampler::new(vec![], 100);
1498 let callback = sampler.on_agent_response();
1499
1500 let json = r#"{"rate_by_service":{"service:web,env:prod":0.5}}"#;
1502 callback(json);
1503 assert!(sampler
1504 .service_samplers
1505 .contains_key("service:web,env:prod"));
1506
1507 callback("not json");
1509
1510 callback(r#"{"other_field":1}"#);
1512 }
1513
1514 #[test]
1515 fn test_rate_limiter_drop_branch() {
1516 let always_keep = SamplingRule::new(1.0, None, None, None, None, None);
1519 let sampler = DatadogSampler::new(vec![always_keep], 0);
1520 let trace_id = TestTraceId::from_bytes([0u8; 16]);
1521 let attributes = create_attributes("res", "prod");
1522 let data = create_sampling_data(None, &trace_id, "op", &attributes);
1523 let decision = sampler.sample(&data);
1524 assert_eq!(
1525 decision.priority,
1526 priority::USER_REJECT,
1527 "rule kept span, rate_limit=0 should then drop it"
1528 );
1529 }
1530
1531 #[test]
1532 fn test_get_trace_root_sampling_info() {
1533 let sampler = DatadogSampler::new(vec![], 100);
1534 let trace_id = TestTraceId::from_bytes([0u8; 16]);
1535 let attributes = create_attributes("res", "prod");
1536 let data = create_sampling_data(None, &trace_id, "op", &attributes);
1537 let decision = sampler.sample(&data);
1538 let _info = decision.get_trace_root_sampling_info();
1539 }
1540}