Skip to main content

libdd_sampling/
datadog_sampler.rs

1// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use 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
11/// Type alias for sampling rules update callback
12/// Consolidated callback type used across crates for remote config sampling updates
13pub 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/// A composite sampler that applies rules in order of precedence
23#[derive(Clone, Debug)]
24pub struct DatadogSampler {
25    /// Sampling rules to apply, in order of precedence
26    rules: RulesSampler,
27
28    /// Service-based samplers provided by the Agent
29    service_samplers: ServicesSampler,
30
31    /// Rate limiter for limiting the number of spans per second
32    rate_limiter: RateLimiter,
33}
34
35impl DatadogSampler {
36    /// Creates a new DatadogSampler with the given rules
37    pub fn new(rules: Vec<SamplingRule>, rate_limit: i32) -> Self {
38        // Create rate limiter with default value of 100 if not provided
39        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    // Test-only helper that bypasses the agent-response parsing path.
49    #[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    /// Creates a callback for updating sampling rules from remote configuration.
68    ///
69    /// # Returns
70    ///
71    /// A boxed function that takes a slice of `SamplingRuleConfig` and updates the sampling rules.
72    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    /// Computes a key for service-based sampling
82    fn service_key(&self, span: &impl SpanProperties) -> String {
83        // `Cow<str>` implements `Display`, so no `into_owned()` allocation is needed here;
84        // `format!` will borrow directly from the span.
85        format!("service:{},env:{}", span.service(), span.env())
86    }
87
88    /// Finds the highest precedence rule that matches the span
89    fn find_matching_rule(&self, span: &impl SpanProperties) -> Option<SamplingRule> {
90        self.rules.find_matching_rule(|rule| rule.matches(span))
91    }
92
93    /// Returns the sampling mechanism used for the decision
94    fn get_sampling_mechanism(
95        &self,
96        rule: Option<&SamplingRule>,
97        used_agent_sampler: bool,
98    ) -> SamplingMechanism {
99        if let Some(rule) = rule {
100            // Provenance is set when rules come from remote configuration
101            // (see `on_rules_update`); locally configured rules use the default value.
102            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            // Should not happen in practice: agent rates default to covering all services.
111            mechanism::DEFAULT
112        }
113    }
114
115    /// Sample an incoming span based on the parent context and attributes.
116    ///
117    /// If a parent sampling decision is present it is inherited; otherwise the root-span
118    /// sampling pipeline is run via [`Self::sample_root`].
119    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    /// Sample the root span of a trace.
135    ///
136    /// Order of precedence:
137    /// 1. A matching local/remote sampling rule (with rate limiting on keep).
138    /// 2. Agent-provided per-service sampling rate.
139    /// 3. Default 100% keep.
140    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                // Rule kept the span, but the rate limiter dropped it.
160                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                // No agent rate for this service yet; keep with rate 1.0 until rates arrive.
173                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
190/// Formats a sampling rate with up to 6 decimal places, stripping trailing zeros.
191///
192/// Matches `strconv.FormatFloat(rate, 'f', 6, 64)` in Go, after trailing-zero
193/// stripping. Rates below 5e-7 round to `"0"`.
194///
195/// # Examples
196/// - `1.0` → `Some("1")`
197/// - `0.5` → `Some("0.5")`
198/// - `0.7654321` → `Some("0.765432")`
199/// - `0.100000` → `Some("0.1")`
200/// - `0.0000001` → `Some("0")` (rounds below the 6th decimal place)
201/// - `0.00000051` → `Some("0.000001")` (rounds up to one millionth)
202/// - `-0.1` → `None`
203/// - `1.1` → `None`
204fn 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    /// Returns the sampling mechanism used for this trace root
221    pub fn mechanism(&self) -> SamplingMechanism {
222        self.mechanism
223    }
224
225    /// Returns the sample rate used for this trace root
226    pub fn rate(&self) -> f64 {
227        self.rate
228    }
229
230    /// Returns the effective rate limit if rate limiting was applied
231    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    /// Returns Datadog-specific sampling tags to be added as attributes
252    ///
253    /// # Parameters
254    /// * `factory` - The attribute factory to use for creating attributes
255    ///
256    /// # Returns
257    /// An optional vector of attributes to add to the sampling result
258    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; // No root info, return empty attributes
264        };
265
266        let mut result: Vec<F::Attribute>;
267        // Add rate limiting tag if applicable
268        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        // Add the sampling decision trace tag with the mechanism
276        let mechanism = root_info.mechanism();
277        result.push(factory.create_string(SAMPLING_DECISION_MAKER_TAG_KEY, mechanism.to_cow()));
278
279        // Add the sample rate tag with the correct key based on the mechanism
280        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    // Test-only semantic convention constants
323    const HTTP_REQUEST_METHOD: &str = "http.request.method";
324    const SERVICE_NAME: &str = "service.name";
325
326    // HTTP status code attribute constants (for tests)
327    const HTTP_RESPONSE_STATUS_CODE: &str = "http.response.status_code";
328    const HTTP_STATUS_CODE: &str = "http.status_code";
329
330    // ============================================================================
331    // Test-only data structures
332    // ============================================================================
333
334    #[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            // Check for HTTP spans - label them all as client spans
427            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            // Default fallback
436            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    // ============================================================================
565    // Test helper functions
566    // ============================================================================
567
568    // Helper function to create a trace ID
569    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    // Helper function to create attributes for testing (with resource and env)
575    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    // Helper function to create attributes with service
583    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    // Helper function to create attributes with service plus arbitrary extra string tags.
596    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    // Helper function to create SamplingData for testing
610    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        // Create a rule without specifying any criteria
650        let rule = SamplingRule::new(
651            0.5, None, // No service
652            None, // No name
653            None, // No resource
654            None, // No tags
655            None, // Default provenance
656        );
657
658        // Verify fields are set to None or empty
659        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        // Verify no matchers were created
667        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        // Test that a rule with NO_RULE constants behaves the same as None
673        let rule_with_empty_strings = SamplingRule::new(
674            0.5,
675            Some(pattern::NO_RULE.to_string()), // Empty service string
676            Some(pattern::NO_RULE.to_string()), // Empty name string
677            Some(pattern::NO_RULE.to_string()), // Empty resource string
678            Some(HashMap::from([(
679                pattern::NO_RULE.to_string(),
680                pattern::NO_RULE.to_string(),
681            )])), // Empty tag
682            None,
683        );
684
685        // Verify that matchers aren't created for NO_RULE values
686        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        // Create a span with some attributes
692        let attributes = create_attributes("some-resource", "some-env");
693
694        // Both rules should match any span since they have no criteria
695        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        // Rule constrained on service, operation name, and a required tag value.
703        // `TestSpan::operation_name()` returns "http.client.request" when the span
704        // carries an `http.request.method` attribute (see `get_operation_name`).
705        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        // Matching span.
718        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        // Non-matching service.
728        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        // Missing required tag.
741        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        // Create two rules with different rates
757        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        // Rule with rate 1.0 should always sample
763        assert!(rule_always.sample(&trace_id));
764
765        // Rule with rate 0.0 should never sample
766        assert!(!rule_never.sample(&trace_id));
767    }
768
769    #[test]
770    fn test_datadog_sampler_creation() {
771        // Create a sampler with default config
772        let sampler = DatadogSampler::new(vec![], 100);
773        assert!(sampler.rules.is_empty());
774        assert!(sampler.service_samplers.is_empty());
775
776        // Create a sampler with rules
777        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        // Test with service and env
788        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        // Test with missing env
797        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        // Update with service rates
813        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        // Check number of samplers
820        assert_eq!(sampler.service_samplers.len(), 2);
821
822        // Verify keys exist
823        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        // Verify the sampling rates are correctly set
831        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        // Create rules with different priorities and service matchers
847        let rule1 = SamplingRule::new(
848            0.1,
849            Some("service1".to_string()),
850            None,
851            None,
852            None,
853            Some("customer".to_string()), // Highest priority
854        );
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()), // Middle priority
863        );
864
865        let rule3 = SamplingRule::new(
866            0.3,
867            Some("service*".to_string()), // Wildcard service
868            None,
869            None,
870            None,
871            Some("default".to_string()), // Lowest priority
872        );
873
874        let sampler = DatadogSampler::new(vec![rule1.clone(), rule2.clone(), rule3.clone()], 100);
875
876        // Test with a specific service that should match the first rule (rule1)
877        {
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        // Test with a specific service that should match the second rule (rule2)
895        {
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        // Test with a service that matches the wildcard rule (rule3)
913        {
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        // Test with a service that doesn't match any rule's service pattern
931        {
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        // Create rules with different provenances
951        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        // Test with customer rule
959        let mechanism1 = sampler.get_sampling_mechanism(Some(&rule_customer), false);
960        assert_eq!(mechanism1, mechanism::REMOTE_USER_TRACE_SAMPLING_RULE);
961
962        // Test with dynamic rule
963        let mechanism2 = sampler.get_sampling_mechanism(Some(&rule_dynamic), false);
964        assert_eq!(mechanism2, mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE);
965
966        // Test with default rule
967        let mechanism3 = sampler.get_sampling_mechanism(Some(&rule_default), false);
968        assert_eq!(mechanism3, mechanism::LOCAL_USER_TRACE_SAMPLING_RULE);
969
970        // Test with agent sampler
971        let mechanism4 = sampler.get_sampling_mechanism(None, true);
972        assert_eq!(mechanism4, mechanism::AGENT_RATE_BY_SERVICE);
973
974        // Test fallback case
975        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        // Test with RecordAndSample decision and LocalUserTraceSamplingRule mechanism
982        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        // Verify the number of attributes (decision_maker + priority + rule_rate + ksr)
999        assert_eq!(attrs.len(), 4);
1000
1001        // Check individual attributes
1002        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                    // For LocalUserTraceSamplingRule with KEEP, it should be USER_KEEP
1019                    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        // Test with rate limiting
1054        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        // With rate limiting, there should be one more attribute
1070        assert_eq!(attrs_with_limit.len(), 5);
1071
1072        // Check for rate limit attribute
1073        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        // Test with AgentRateByService mechanism to check for SAMPLING_AGENT_RATE_TAG_KEY
1089
1090        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        // Verify the number of attributes (should be 4: decision_maker + priority +
1107        // agent_rate + ksr)
1108        assert_eq!(agent_attrs.len(), 4);
1109
1110        // Check for agent rate tag and ksr tag
1111        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        // Also check that the SAMPLING_RULE_RATE_TAG_KEY is NOT present for agent mechanism
1142        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        // Exact values
1154        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        // Trailing zeros should be stripped
1160        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        // Truncation to 6 decimal places
1164        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        // Small values
1174        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        // Below the 6th decimal place rounds to "0"
1178        assert_eq!(format_sampling_rate(0.0000001), Some("0".to_string()));
1179        // Just above the round-up threshold lands on one millionth
1180        assert_eq!(
1181            format_sampling_rate(0.00000051),
1182            Some("0.000001".to_string())
1183        );
1184
1185        // Boundary values
1186        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        // Invalid rates
1190        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        // Create empty slices for attributes and links
1202        let empty_attrs: &[TestAttribute] = &[];
1203        let trace_id = create_trace_id();
1204
1205        // Test with sampled parent context
1206        let data_sampled = create_sampling_data(Some(true), &trace_id, "span", empty_attrs);
1207        let result_sampled = sampler.sample(&data_sampled);
1208
1209        // Should inherit the sampling decision from parent
1210        assert!(result_sampled.get_priority().is_keep());
1211        assert!(result_sampled
1212            .to_dd_sampling_tags(&TestAttributeFactory)
1213            .is_none());
1214
1215        // Test with non-sampled parent context
1216        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        // Should inherit the sampling decision from parent
1220        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        // Create a rule that always samples
1229        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        // Test with matching attributes
1243        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        // Should sample and add attributes
1248        assert!(result.get_priority().is_keep());
1249        assert!(result.to_dd_sampling_tags(&TestAttributeFactory).is_some());
1250
1251        // Test with non-matching attributes
1252        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        // Should still sample (default behavior when no rules match) and add attributes
1258        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        // Initialize sampler
1267        let sampler = DatadogSampler::new(vec![], 100);
1268
1269        // Add service rates for different service+env combinations
1270        let mut rates = HashMap::new();
1271        rates.insert("service:test-service,env:prod".to_string(), 1.0); // Always sample for test-service in prod
1272        rates.insert("service:other-service,env:prod".to_string(), 0.0); // Never sample for other-service in prod
1273
1274        sampler.update_service_rates(rates);
1275
1276        let trace_id = create_trace_id();
1277
1278        // Test with attributes that should lead to "service:test-service,env:prod" key
1279        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        // Expect RecordAndSample because service_key will be "service:test-service,env:prod" ->
1292        // rate 1.0
1293        assert!(
1294            result_sample.get_priority().is_keep(),
1295            "Span for test-service/prod should be sampled"
1296        );
1297
1298        // Test with attributes that should lead to "service:other-service,env:prod" key
1299        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        // Expect Drop because service_key will be "service:other-service,env:prod" -> rate 0.0
1312        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        // Helper to create attributes with a float value
1321        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        // Test case 1: Rule with exact value matching integer float
1333        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        // Should match integer float
1343        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        // Test case 2: Rule with wildcard pattern and non-integer float
1348        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        // Should match non-integer float with wildcard pattern
1358        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        // Test case 3: Rule with specific pattern and non-integer float
1363        // With our simplified logic, non-integer floats will never match non-wildcard patterns
1364        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        // Should NOT match the exact decimal value because non-integer floats only match wildcards
1377        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        // Test case 4: Pattern with partial wildcard '*' for suffix
1381        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        // Should NOT match decimal values as we don't do partial pattern matching for non-integer
1394        // floats
1395        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        // Test that the sampler correctly matches rules based on operation names
1402        // Operation name generation itself is tested in otel_mappings unit tests
1403
1404        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        // HTTP client request should match http_rule (operation name: http.client.request)
1418        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        // Span that doesn't match the rule should still be sampled (default behavior)
1423        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        // Create a sampler with initial rules
1431        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        // Verify initial state
1443        assert_eq!(sampler.rules.len(), 1);
1444
1445        // Get the callback
1446        let callback = sampler.on_rules_update();
1447
1448        // Create new rules directly as SamplingRuleConfig objects
1449        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        // Apply the update
1469        callback(&new_rules);
1470
1471        // Verify the rules were updated
1472        assert_eq!(sampler.rules.len(), 2);
1473
1474        // Test that the new rules work by finding a matching rule
1475        // Create attributes that will generate an operation name matching "http.*"
1476        // and service matching "web-*"
1477        let attrs = vec![
1478            TestAttribute::new(SERVICE_NAME, "web-frontend"),
1479            TestAttribute::new(HTTP_REQUEST_METHOD, "GET"), /* This will make operation name
1480                                                             * "http.client.request" */
1481        ];
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        // Test with empty rules array
1491        callback(&[]);
1492        assert_eq!(sampler.rules.len(), 0); // Should now have no rules
1493    }
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        // Valid JSON with rate_by_service
1501        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        // Invalid JSON — should not panic
1508        callback("not json");
1509
1510        // Missing rate_by_service — should not panic
1511        callback(r#"{"other_field":1}"#);
1512    }
1513
1514    #[test]
1515    fn test_rate_limiter_drop_branch() {
1516        // Rule with sample_rate=1.0 keeps everything, but rate_limit=0 then drops
1517        // every kept span via the rate limiter, exercising the drop branch.
1518        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}