feat(sampling): OTel consistent-probability rv/th derivation (APMAPI-2181)

libdd-sampling/src/datadog_sampler.rs

@@ -1,6 +1,7 @@
 // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+use crate::constants::numeric::KNUTH_FACTOR;
 use crate::dd_constants::{
     RL_EFFECTIVE_RATE, SAMPLING_AGENT_RATE_TAG_KEY, SAMPLING_DECISION_MAKER_TAG_KEY,
     SAMPLING_KNUTH_RATE_TAG_KEY, SAMPLING_PRIORITY_TAG_KEY, SAMPLING_RULE_RATE_TAG_KEY,
@@ -12,7 +13,7 @@ use crate::sampling_rule_config::SamplingRuleConfig;
 /// Consolidated callback type used across crates for remote config sampling updates
 pub type SamplingRulesCallback = Box<dyn for<'a> Fn(&'a [SamplingRuleConfig]) + Send + Sync>;
 
-use crate::types::{SamplingData, SpanProperties};
+use crate::types::{SamplingData, SpanProperties, TraceIdLike};
 
 use super::agent_service_sampler::{AgentRates, ServicesSampler};
 use super::rate_limiter::RateLimiter;
@@ -185,6 +186,7 @@ impl DatadogSampler {
                 mechanism,
                 rate: sample_rate,
                 rl_effective_rate,
+                is_keep,
             }),
         }
     }
@@ -213,13 +215,63 @@ fn format_sampling_rate(rate: f64) -> Option<String> {
     Some(s.trim_end_matches('0').trim_end_matches('.').to_string())
 }
 
+/// OTel consistent-probability tracestate values (`ot.rv`, `ot.th`), 56 bits each.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct OtelConsistentSampling {
+    /// 56-bit random value derived from the trace id.
+    pub rv: u64,
+    /// 56-bit rejection threshold derived from the sample rate.
+    pub th: u64,
+}
+
 pub struct TraceRootSamplingInfo {
     mechanism: SamplingMechanism,
     rate: f64,
     rl_effective_rate: Option<f64>,
+    is_keep: bool,
+}
+
+/// Derives the 56-bit OTel `rv` from the low 64 bits of a trace id.
+fn derive_rv<T: TraceIdLike>(trace_id: &T) -> u64 {
+    let low64 = trace_id.to_u128() as u64;
+    (!low64.wrapping_mul(KNUTH_FACTOR)) >> 8
+}
+
+/// Derives the 56-bit OTel `th` (rejection threshold) from a DD sample rate.
+fn derive_th(rate: f64) -> u64 {
+    const U56_MOD: u64 = 1u64 << 56;
+    const U56_MAX: u64 = U56_MOD - 1;
+    (((1.0 - rate) * U56_MOD as f64).round() as u64).clamp(0, U56_MAX)
+}
+
+/// Nudges `rv` across `th` when compressing to 56 bits flipped the `rv >= th`
+/// comparison relative to `kept`.
+fn reconcile_rv(rv: u64, th: u64, kept: bool) -> u64 {
+    if kept && rv < th {
+        th
+    } else if !kept && rv >= th {
+        th.saturating_sub(1)
+    } else {
+        rv
+    }
 }
 
 impl TraceRootSamplingInfo {
+    /// Derives the reconciled OTel `rv`/`th` pair, or `None` for a
+    /// non-probability mechanism or a rate-limiter-overturned keep.
+    pub fn otel_consistent_sampling<T: TraceIdLike>(
+        &self,
+        trace_id: &T,
+    ) -> Option<OtelConsistentSampling> {
+        if !self.mechanism.is_probability() || (self.rl_effective_rate.is_some() && !self.is_keep) {
+            return None;
+        }
+        let rv = derive_rv(trace_id);
+        let th = derive_th(self.rate);
+        let rv = reconcile_rv(rv, th, self.is_keep);
+        Some(OtelConsistentSampling { rv, th })
+    }
+
     /// Returns the sampling mechanism used for this trace root
     pub fn mechanism(&self) -> SamplingMechanism {
         self.mechanism
@@ -322,6 +374,130 @@ mod tests {
     use std::borrow::Cow;
     use std::collections::HashMap;
 
+    #[test]
+    fn test_otel_consistent_sampling() {
+        use crate::dd_sampling::mechanism;
+        // rv depends only on trace id
+        let cases = [
+            (1u128, 0xf0948a54d43b8eu64),
+            (10, 0x65cd67504a538e),
+            (100, 0xfa060922e7438e),
+            (1000, 0xc43c5b5d08a38e),
+            (18446744073709551615, 0x0f6b75ab2bc471),
+            (83, 0x0028d980cf4f1c),
+            (18444899399302180863, 0xef284ace7a91e1),
+        ];
+        for (tid, want_rv) in cases {
+            // Raw derivation only, independent of the keep/drop reconciliation
+            // step exercised separately below.
+            assert_eq!(derive_rv(&tid), want_rv, "rv for {tid}");
+        }
+        // th depends only on rate. See system-tests #7372.
+        let th_cases = [
+            (0.01f64, 0xfd70a3d70a3d70u64),
+            (0.1, 0xe6666666666668),
+            (0.2, 0xccccccccccccd0),
+            (0.5, 0x80000000000000),
+            (0.99, 0x28f5c28f5c290),
+            (1.0, 0x0),
+        ];
+        for (rate, want_th) in th_cases {
+            assert_eq!(derive_th(rate), want_th, "th for rate {rate}");
+        }
+
+        // non-probability mechanism -> None
+        let manual = TraceRootSamplingInfo {
+            mechanism: mechanism::MANUAL,
+            rate: 1.0,
+            rl_effective_rate: None,
+            is_keep: true,
+        };
+        assert!(manual.otel_consistent_sampling(&1u128).is_none());
+
+        // probability mechanism, but the rate limiter overturned the keep ->
+        // None (th erased; caller still forwards an inherited rv).
+        let rate_limited = TraceRootSamplingInfo {
+            mechanism: mechanism::LOCAL_USER_TRACE_SAMPLING_RULE,
+            rate: 0.5,
+            rl_effective_rate: Some(0.25),
+            is_keep: false,
+        };
+        assert!(rate_limited.otel_consistent_sampling(&1u128).is_none());
+
+        let rate_limiter_allowed = TraceRootSamplingInfo {
+            mechanism: mechanism::LOCAL_USER_TRACE_SAMPLING_RULE,
+            rate: 0.1,
+            rl_effective_rate: Some(1.0),
+            is_keep: true,
+        };
+        assert_eq!(
+            rate_limiter_allowed.otel_consistent_sampling(&1u128),
+            Some(OtelConsistentSampling {
+                rv: 0xf0948a54d43b8e,
+                th: 0xe6666666666668,
+            })
+        );
+    }
+
+    #[test]
+    fn test_reconcile_rv_nudges_only_on_disagreement() {
+        // Worked example (rate 0.1): DD keeps, but rv < th, so a downstream
+        // re-derivation of `rv >= th` would disagree (drop).
+        assert_eq!(
+            reconcile_rv(0xe6666666666666, 0xe6666666666668, true),
+            0xe6666666666668
+        );
+        // Symmetric case: DD drops, but rv >= th, so a re-derivation would
+        // disagree (keep).
+        assert_eq!(
+            reconcile_rv(0xe6666666666668, 0xe6666666666668, false),
+            0xe6666666666667
+        );
+        // Already agrees (kept, rv >= th) -> untouched.
+        assert_eq!(
+            reconcile_rv(0xef284ace7a91e1, 0xe6666666666666, true),
+            0xef284ace7a91e1
+        );
+        // Already agrees (dropped, rv < th) -> untouched.
+        assert_eq!(
+            reconcile_rv(0x1000000000000, 0xe6666666666666, false),
+            0x1000000000000
+        );
+        // th == 0 (rate 1.0, keep-all): a drop decision here is inconsistent
+        // with th, but the saturating subtraction must not underflow.
+        assert_eq!(reconcile_rv(0, 0, false), 0);
+    }
+
+    #[test]
+    fn test_otel_consistent_sampling_reconciles_disagreement() {
+        // rate 0.1 -> th = 0xe6666666666668. Pick a trace id whose raw rv
+        // lands just under th, then check that a DD keep nudges rv up to
+        // agree, while a DD drop leaves it alone (already agrees).
+        let trace_id = 18446744073709551615u128; // rv = 0x0f6b75ab2bc471 (see golden cases)
+        let raw_rv = derive_rv(&trace_id);
+        let th = derive_th(0.1);
+        assert!(raw_rv < th, "test assumes a naturally-disagreeing pair");
+
+        let kept = TraceRootSamplingInfo {
+            mechanism: mechanism::DEFAULT,
+            rate: 0.1,
+            rl_effective_rate: None,
+            is_keep: true,
+        };
+        let got = kept.otel_consistent_sampling(&trace_id).unwrap();
+        assert_eq!(got.th, th);
+        assert_eq!(got.rv, th, "kept but rv < th must be nudged up to th");
+
+        let dropped = TraceRootSamplingInfo {
+            mechanism: mechanism::DEFAULT,
+            rate: 0.1,
+            rl_effective_rate: None,
+            is_keep: false,
+        };
+        let got = dropped.otel_consistent_sampling(&trace_id).unwrap();
+        assert_eq!(got.rv, raw_rv, "dropped and rv < th already agree");
+    }
+
     // Test-only semantic convention constants
     const HTTP_REQUEST_METHOD: &str = "http.request.method";
     const SERVICE_NAME: &str = "service.name";
@@ -991,6 +1167,7 @@ mod tests {
                 mechanism,
                 rate: 0.5,
                 rl_effective_rate: None,
+                is_keep: is_sampled,
             }),
         };
 
@@ -1063,6 +1240,7 @@ mod tests {
                 mechanism,
                 rate: 0.5,
                 rl_effective_rate: Some(rate_limit),
+                is_keep: is_sampled,
             }),
         };
         let attrs_with_limit = sampling_result
@@ -1099,6 +1277,7 @@ mod tests {
                 mechanism,
                 rate: agent_rate,
                 rl_effective_rate: None,
+                is_keep: is_sampled,
             }),
         };
 

libdd-sampling/src/dd_sampling.rs

@@ -141,6 +141,25 @@ impl SamplingMechanism {
         }
     }
 
+    /// Returns whether this mechanism is a probability (rate-based) sampling
+    /// decision, for OTel consistent-probability tracestate (`ot.th`).
+    ///
+    /// `OTLP_INGEST_PROBABILISTIC_SAMPLING` is excluded: it's the sender's own OTel
+    /// decision, with no rate or trace-id-hash contract available here to derive `th`.
+    pub fn is_probability(&self) -> bool {
+        matches!(
+            *self,
+            mechanism::DEFAULT
+                | mechanism::AGENT_RATE_BY_SERVICE
+                | mechanism::REMOTE_RATE
+                | mechanism::REMOTE_RATE_USER
+                | mechanism::REMOTE_RATE_DATADOG
+                | mechanism::LOCAL_USER_TRACE_SAMPLING_RULE
+                | mechanism::REMOTE_USER_TRACE_SAMPLING_RULE
+                | mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE
+        )
+    }
+
     /// Returns the string representation of the sampling mechanism.
     ///
     /// The format is `"-N"` (e.g. `"-4"` for manual sampling). The leading `-` comes from the
@@ -399,6 +418,32 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_mechanism_is_probability() {
+        use mechanism::*;
+        for m in [
+            DEFAULT,
+            AGENT_RATE_BY_SERVICE,
+            REMOTE_RATE,
+            REMOTE_RATE_USER,
+            REMOTE_RATE_DATADOG,
+            LOCAL_USER_TRACE_SAMPLING_RULE,
+            REMOTE_USER_TRACE_SAMPLING_RULE,
+            REMOTE_DYNAMIC_TRACE_SAMPLING_RULE,
+        ] {
+            assert!(m.is_probability(), "{m:?} should be probability");
+        }
+        for m in [
+            MANUAL,
+            APPSEC,
+            SPAN_SAMPLING_RULE,
+            DATA_JOBS_MONITORING,
+            OTLP_INGEST_PROBABILISTIC_SAMPLING,
+        ] {
+            assert!(!m.is_probability(), "{m:?} should not be probability");
+        }
+    }
+
     #[test]
     fn test_mechanism_from_str_errors() {
         assert!("not_a_number".parse::<SamplingMechanism>().is_err());

libdd-sampling/src/lib.rs

@@ -47,7 +47,7 @@ pub mod v04_span;
 
 // Re-export key types for convenience
 pub use agent_service_sampler::ServicesSampler;
-pub use datadog_sampler::{DatadogSampler, SamplingRulesCallback};
+pub use datadog_sampler::{DatadogSampler, OtelConsistentSampling, SamplingRulesCallback};
 pub use dd_sampling::{mechanism, priority, SamplingDecision, SamplingMechanism, SamplingPriority};
 pub use sampling_rule::SamplingRule;
 pub use sampling_rule_config::{ParsedSamplingRules, SamplingRuleConfig};

libdd-sampling/src/types.rs

@@ -33,6 +33,13 @@ pub trait TraceIdLike: Eq {
     fn to_u128(&self) -> u128;
 }
 
+/// `u128` is the canonical numeric trace-id representation.
+impl TraceIdLike for u128 {
+    fn to_u128(&self) -> u128 {
+        *self
+    }
+}
+
 /// A trait for accessing span attribute key-value pairs.
 ///
 /// Provides methods for retrieving the key and value of a span attribute.

libdd-sampling/src/v04_span.rs

@@ -3,8 +3,8 @@
 
 //! Sampling trait implementations for the v04 [`Span<T>`] type.
 //!
-//! This module implements all six sampling traits on the v04 span representation:
-//! [`TraceIdLike`] for `u128`, [`ValueLike`]/[`AttributeLike`] for attributes
+//! This module implements the sampling traits on the v04 span representation:
+//! [`ValueLike`]/[`AttributeLike`] for attributes
 //! borrowed from `span.meta` and `span.metrics`, [`SpanProperties`] via the
 //! [`V04SpanProperties`] wrapper, [`SamplingData`] via [`V04SamplingData`], and
 //! [`AttributeFactory`] via [`V04AttributeFactory`].
@@ -48,16 +48,7 @@ use std::borrow::{Borrow, Cow};
 
 use libdd_trace_utils::span::{v04::Span, TraceData};
 
-use crate::types::{
-    AttributeFactory, AttributeLike, SamplingData, SpanProperties, TraceIdLike, ValueLike,
-};
-
-/// `u128` is the native type for v04 trace IDs.
-impl TraceIdLike for u128 {
-    fn to_u128(&self) -> u128 {
-        *self
-    }
-}
+use crate::types::{AttributeFactory, AttributeLike, SamplingData, SpanProperties, ValueLike};
 
 /// A span attribute value sourced from either `span.meta` (string) or `span.metrics` (f64).
 pub enum SpanAttributeValue<'a> {
@@ -267,16 +258,6 @@ mod tests {
         }
     }
 
-    #[test]
-    fn test_trace_id_like_u128() {
-        let id: u128 = 42;
-        assert_eq!(id.to_u128(), 42);
-        let zero: u128 = 0;
-        assert_eq!(zero.to_u128(), 0);
-        let max = u128::MAX;
-        assert_eq!(max.to_u128(), u128::MAX);
-    }
-
     #[test]
     fn test_span_attribute_value_meta() {
         let val = SpanAttributeValue::Meta("hello");