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