chore(sampling): pr review feedback

Author: bantonsson

.codecov.yml

@@ -55,6 +55,10 @@ component_management:
       name: libdd-profiling-ffi  # this is a display name, and can be changed freely
       paths:
         - libdd-profiling-ffi
+    - component_id: sampling  # this is an identifier that should not be changed
+      name: libdd-sampling  # this is a display name, and can be changed freely
+      paths:
+        - libdd-sampling
     - component_id: sidecar  # this is an identifier that should not be changed
       name: datadog-sidecar  # this is a display name, and can be changed freely
       paths:

libdd-sampling/Cargo.toml

@@ -26,6 +26,7 @@ required-features = ["v04_span"]
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 lru = "0.16.3"
+libdd-common = { path = "../libdd-common", version = "4.0.0" }
 libdd-trace-utils = { path = "../libdd-trace-utils", version = "3.0.1", optional = true }
 
 [features]
@@ -34,4 +35,3 @@ v04_span = ["dep:libdd-trace-utils"]
 [dev-dependencies]
 criterion = "0.5"
 libdd-common = { path = "../libdd-common", version = "4.0.0", features = ["bench-utils"] }
-libdd-trace-utils = { path = "../libdd-trace-utils", version = "3.0.1" }

libdd-sampling/src/agent_service_sampler.rs

@@ -3,7 +3,7 @@
 
 use std::{
     collections::HashMap,
-    sync::{Arc, RwLock},
+    sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard},
 };
 
 use crate::rate_sampler::RateSampler;
@@ -20,32 +20,48 @@ pub struct ServicesSampler {
 }
 
 impl ServicesSampler {
+    fn read_or_panic(&self) -> RwLockReadGuard<'_, HashMap<String, RateSampler>> {
+        #[allow(clippy::panic)]
+        match self.inner.read() {
+            Ok(guard) => guard,
+            Err(_) => panic!("ServicesSampler read lock poisoned"),
+        }
+    }
+
+    fn write_or_panic(&self) -> RwLockWriteGuard<'_, HashMap<String, RateSampler>> {
+        #[allow(clippy::panic)]
+        match self.inner.write() {
+            Ok(guard) => guard,
+            Err(_) => panic!("ServicesSampler write lock poisoned"),
+        }
+    }
+
     pub fn get(&self, service: &str) -> Option<RateSampler> {
-        self.inner.read().unwrap().get(service).cloned()
+        self.read_or_panic().get(service).cloned()
     }
 
     pub fn update_rates<I: IntoIterator<Item = (String, f64)>>(&self, rates: I) {
         let new_rates: HashMap<_, _> = rates
             .into_iter()
             .map(|(s, r)| (s, RateSampler::new(r)))
             .collect();
-        *self.inner.write().unwrap() = new_rates;
+        *self.write_or_panic() = new_rates;
     }
 
-    // used for testing purposes
+    // Test-only inspection helpers.
 
-    #[allow(dead_code)]
+    #[cfg(test)]
     pub(crate) fn is_empty(&self) -> bool {
-        self.inner.read().unwrap().is_empty()
+        self.read_or_panic().is_empty()
     }
 
-    #[allow(dead_code)]
+    #[cfg(test)]
     pub(crate) fn len(&self) -> usize {
-        self.inner.read().unwrap().len()
+        self.read_or_panic().len()
     }
 
-    #[allow(dead_code)]
+    #[cfg(test)]
     pub(crate) fn contains_key(&self, service: &str) -> bool {
-        self.inner.read().unwrap().contains_key(service)
+        self.read_or_panic().contains_key(service)
     }
 }

libdd-sampling/src/constants.rs

@@ -1,7 +1,7 @@
 // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
-//! Shared constants for the datadog_opentelemetry::sampling crate
+//! Shared constants for the libdd-sampling crate
 
 /// Sampling rate limits
 pub mod rate {

libdd-sampling/src/datadog_sampler.rs

@@ -45,8 +45,8 @@ impl DatadogSampler {
         }
     }
 
-    // used for tests
-    #[allow(dead_code)]
+    // Test-only helper that bypasses the agent-response parsing path.
+    #[cfg(test)]
     pub(crate) fn update_service_rates(&self, rates: impl IntoIterator<Item = (String, f64)>) {
         self.service_samplers.update_rates(rates);
     }
@@ -64,9 +64,11 @@ impl DatadogSampler {
         })
     }
 
-    /// Creates a callback for updating sampling rules from remote configuration
+    /// Creates a callback for updating sampling rules from remote configuration.
+    ///
     /// # Returns
-    /// A boxed function that takes a slice of SamplingRuleConfig and updates the sampling rules
+    ///
+    /// A boxed function that takes a slice of `SamplingRuleConfig` and updates the sampling rules.
     pub fn on_rules_update(&self) -> SamplingRulesCallback {
         let rules_sampler = self.rules.clone();
         Box::new(move |rule_configs: &[SamplingRuleConfig]| {
@@ -78,12 +80,9 @@ impl DatadogSampler {
 
     /// Computes a key for service-based sampling
     fn service_key(&self, span: &impl SpanProperties) -> String {
-        // Get service from span
-        let service = span.service().into_owned();
-        // Get env from span
-        let env = span.env();
-
-        format!("service:{service},env:{env}")
+        // `Cow<str>` implements `Display`, so no `into_owned()` allocation is needed here;
+        // `format!` will borrow directly from the span.
+        format!("service:{},env:{}", span.service(), span.env())
     }
 
     /// Finds the highest precedence rule that matches the span
@@ -98,40 +97,46 @@ impl DatadogSampler {
         used_agent_sampler: bool,
     ) -> SamplingMechanism {
         if let Some(rule) = rule {
+            // Provenance is set when rules come from remote configuration
+            // (see `on_rules_update`); locally configured rules use the default value.
             match rule.provenance.as_str() {
-                // Provenance will not be set for rules until we implement remote configuration
                 "customer" => mechanism::REMOTE_USER_TRACE_SAMPLING_RULE,
                 "dynamic" => mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE,
                 _ => mechanism::LOCAL_USER_TRACE_SAMPLING_RULE,
             }
         } else if used_agent_sampler {
-            // If using service-based sampling from the agent
             mechanism::AGENT_RATE_BY_SERVICE
         } else {
-            // Should not happen, but just in case
+            // Should not happen in practice: agent rates default to covering all services.
             mechanism::DEFAULT
         }
     }
 
-    /// Sample an incoming span based on the parent context and attributes
+    /// Sample an incoming span based on the parent context and attributes.
+    ///
+    /// If a parent sampling decision is present it is inherited; otherwise the root-span
+    /// sampling pipeline is run via [`Self::sample_root`].
     pub fn sample(&self, data: &impl SamplingData) -> DdSamplingResult {
         if let Some(is_parent_sampled) = data.is_parent_sampled() {
             let priority = match is_parent_sampled {
                 false => priority::AUTO_REJECT,
                 true => priority::AUTO_KEEP,
             };
-            // If a parent exists, inherit its sampling decision and trace state
             return DdSamplingResult {
                 priority,
                 trace_root_info: None,
             };
         }
 
-        // Apply rules-based sampling
         data.with_span_properties(self, |sampler, span| sampler.sample_root(data, span))
     }
 
-    /// Sample the root span of a trace
+    /// Sample the root span of a trace.
+    ///
+    /// Order of precedence:
+    /// 1. A matching local/remote sampling rule (with rate limiting on keep).
+    /// 2. Agent-provided per-service sampling rate.
+    /// 3. Default 100% keep.
     fn sample_root(
         &self,
         data: &impl SamplingData,
@@ -143,42 +148,32 @@ impl DatadogSampler {
         let mut rl_effective_rate: Option<f64> = None;
         let trace_id = data.trace_id();
 
-        // Find a matching rule
         let matching_rule = self.find_matching_rule(span);
 
-        // Apply sampling logic
         if let Some(rule) = &matching_rule {
-            // Get the sample rate from the rule
             sample_rate = rule.sample_rate;
 
-            // First check if the span should be sampled according to the rule
             if !rule.sample(trace_id) {
                 is_keep = false;
-            // If the span should be sampled, then apply rate limiting
             } else if !self.rate_limiter.is_allowed() {
+                // Rule kept the span, but the rate limiter dropped it.
                 is_keep = false;
                 rl_effective_rate = Some(self.rate_limiter.effective_rate());
             }
         } else {
-            // Try service-based sampling from Agent
             let service_key = self.service_key(span);
             if let Some(sampler) = self.service_samplers.get(&service_key) {
-                // Use the service-based sampler
                 used_agent_sampler = true;
-                sample_rate = sampler.sample_rate(); // Get rate for reporting
-
-                // Check if the service sampler decides to drop
+                sample_rate = sampler.sample_rate();
                 if !sampler.sample(trace_id) {
                     is_keep = false;
                 }
             } else {
-                // Default sample rate, should never happen in practice if agent provides rates
+                // No agent rate for this service yet; keep with rate 1.0 until rates arrive.
                 sample_rate = 1.0;
-                // Keep the default decision (RecordAndSample)
             }
         }
 
-        // Determine the sampling mechanism
         let mechanism = self.get_sampling_mechanism(matching_rule.as_ref(), used_agent_sampler);
 
         DdSamplingResult {
@@ -227,9 +222,7 @@ fn format_sampling_rate(rate: f64) -> Option<String> {
     let s = format!("{:.prec$}", rounded, prec = decimal_places);
     // Strip trailing zeros after decimal point
     Some(if s.contains('.') {
-        let s = s.trim_end_matches('0');
-        let s = s.trim_end_matches('.');
-        s.to_string()
+        s.trim_end_matches('0').trim_end_matches('.').to_string()
     } else {
         s
     })
@@ -381,15 +374,15 @@ mod tests {
     }
 
     impl ValueLike for TestValue {
-        fn extract_float(&self) -> Option<f64> {
+        fn as_float(&self) -> Option<f64> {
             match self {
                 TestValue::I64(i) => Some(*i as f64),
                 TestValue::F64(f) => Some(*f),
                 _ => None,
             }
         }
 
-        fn extract_string(&self) -> Option<Cow<'_, str>> {
+        fn as_str(&self) -> Option<Cow<'_, str>> {
             match self {
                 TestValue::String(s) => Some(Cow::Borrowed(s.as_str())),
                 TestValue::I64(i) => Some(Cow::Owned(i.to_string())),
@@ -476,23 +469,23 @@ mod tests {
             self.attributes
                 .iter()
                 .find(|attr| attr.key() == SERVICE_NAME)
-                .and_then(|attr| attr.value().extract_string())
+                .and_then(|attr| attr.value().as_str())
                 .unwrap_or(Cow::Borrowed(""))
         }
 
         fn env(&self) -> Cow<'_, str> {
             self.attributes
                 .iter()
                 .find(|attr| attr.key() == "datadog.env" || attr.key() == ENV_TAG)
-                .and_then(|attr| attr.value().extract_string())
+                .and_then(|attr| attr.value().as_str())
                 .unwrap_or(Cow::Borrowed(""))
         }
 
         fn resource(&self) -> Cow<'_, str> {
             self.attributes
                 .iter()
                 .find(|attr| attr.key() == RESOURCE_TAG)
-                .and_then(|attr| attr.value().extract_string())
+                .and_then(|attr| attr.value().as_str())
                 .unwrap_or(Cow::Borrowed(self.name))
         }
 
@@ -617,6 +610,20 @@ mod tests {
         ]
     }
 
+    // Helper function to create attributes with service plus arbitrary extra string tags.
+    fn create_attributes_with_extra(
+        service: &'static str,
+        resource: &'static str,
+        env: &'static str,
+        extra: &[(&'static str, &'static str)],
+    ) -> Vec<TestAttribute> {
+        let mut attrs = create_attributes_with_service(service.to_string(), resource, env);
+        for (k, v) in extra {
+            attrs.push(TestAttribute::new(*k, *v));
+        }
+        attrs
+    }
+
     // Helper function to create SamplingData for testing
     fn create_sampling_data<'a>(
         is_parent_sampled: Option<bool>,
@@ -710,18 +717,56 @@ mod tests {
 
     #[test]
     fn test_sampling_rule_matches() {
-        // Create a rule with specific service and name patterns
-        let _rule = SamplingRule::new(
+        // Rule constrained on service, operation name, and a required tag value.
+        // `TestSpan::operation_name()` returns "http.client.request" when the span
+        // carries an `http.request.method` attribute (see `get_operation_name`).
+        let rule = SamplingRule::new(
             0.5,
             Some("web-*".to_string()),
-            Some("http.*".to_string()),
+            Some("http.client.*".to_string()),
             None,
             Some(HashMap::from([(
                 "custom_key".to_string(),
                 "custom_value".to_string(),
             )])),
             None,
         );
+
+        // Matching span.
+        let attrs = create_attributes_with_extra(
+            "web-foo",
+            "resource",
+            "production",
+            &[(HTTP_REQUEST_METHOD, "GET"), ("custom_key", "custom_value")],
+        );
+        let span = TestSpan::new("span-name", attrs.as_slice());
+        assert!(rule.matches(&span), "rule should match qualifying span");
+
+        // Non-matching service.
+        let attrs_bad_service = create_attributes_with_extra(
+            "api-foo",
+            "resource",
+            "production",
+            &[(HTTP_REQUEST_METHOD, "GET"), ("custom_key", "custom_value")],
+        );
+        let span_bad_service = TestSpan::new("span-name", attrs_bad_service.as_slice());
+        assert!(
+            !rule.matches(&span_bad_service),
+            "rule should not match different service"
+        );
+
+        // Missing required tag.
+        let attrs_no_tag = create_attributes_with_extra(
+            "web-foo",
+            "resource",
+            "production",
+            &[(HTTP_REQUEST_METHOD, "GET")],
+        );
+        let span_no_tag = TestSpan::new("span-name", attrs_no_tag.as_slice());
+        assert!(
+            !rule.matches(&span_no_tag),
+            "rule should not match without required tag"
+        );
     }
 
     #[test]

libdd-sampling/src/dd_sampling.rs

@@ -142,6 +142,11 @@ impl SamplingMechanism {
     }
 
     /// Returns the string representation of the sampling mechanism.
+    ///
+    /// The format is `"-N"` (e.g. `"-4"` for manual sampling). The leading `-` comes from the
+    /// propagation tags RFC, which initially had a prefix component before the `-`; that prefix
+    /// was dropped, but the `-` was retained as-is for backwards compatibility with existing
+    /// existing tracers.
     pub fn to_cow(self) -> Cow<'static, str> {
         match self {
             mechanism::DEFAULT => Cow::Borrowed("-0"),