Skip to main content

libdd_sampling/
sampling_rule_config.rs

1// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::fmt::Display;
7use std::ops::Deref;
8use std::str::FromStr;
9
10/// Configuration for a single sampling rule
11#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
12pub struct SamplingRuleConfig {
13    /// The sample rate to apply (0.0-1.0)
14    pub sample_rate: f64,
15
16    /// Optional service name pattern to match
17    #[serde(default)]
18    pub service: Option<String>,
19
20    /// Optional span name pattern to match
21    #[serde(default)]
22    pub name: Option<String>,
23
24    /// Optional resource name pattern to match
25    #[serde(default)]
26    pub resource: Option<String>,
27
28    /// Tags that must match (key-value pairs)
29    #[serde(default)]
30    pub tags: HashMap<String, String>,
31
32    /// Where this rule comes from (customer, dynamic, default).
33    /// Not exposed in the public `datadog-opentelemetry` API — set automatically
34    /// during conversion from the public `SamplingRuleConfig` type.
35    #[serde(default = "default_provenance")]
36    pub provenance: String,
37}
38
39impl Default for SamplingRuleConfig {
40    fn default() -> Self {
41        // Keep `Default` in sync with the serde defaults so that constructing a config
42        // with `..Default::default()` matches what deserialization would produce.
43        Self {
44            sample_rate: 0.0,
45            service: None,
46            name: None,
47            resource: None,
48            tags: HashMap::new(),
49            provenance: default_provenance(),
50        }
51    }
52}
53
54impl Display for SamplingRuleConfig {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        write!(f, "{}", serde_json::json!(self))
57    }
58}
59
60fn default_provenance() -> String {
61    "default".to_string()
62}
63
64#[derive(Debug, Default, Clone, PartialEq)]
65pub struct ParsedSamplingRules {
66    pub rules: Vec<SamplingRuleConfig>,
67}
68
69impl Deref for ParsedSamplingRules {
70    type Target = [SamplingRuleConfig];
71
72    fn deref(&self) -> &Self::Target {
73        &self.rules
74    }
75}
76
77impl From<ParsedSamplingRules> for Vec<SamplingRuleConfig> {
78    fn from(parsed: ParsedSamplingRules) -> Self {
79        parsed.rules
80    }
81}
82
83impl FromStr for ParsedSamplingRules {
84    type Err = serde_json::Error;
85
86    fn from_str(s: &str) -> Result<Self, Self::Err> {
87        if s.trim().is_empty() {
88            return Ok(ParsedSamplingRules::default());
89        }
90        // DD_TRACE_SAMPLING_RULES is expected to be a JSON array of SamplingRuleConfig objects.
91        let rules_vec: Vec<SamplingRuleConfig> = serde_json::from_str(s)?;
92        Ok(ParsedSamplingRules { rules: rules_vec })
93    }
94}
95
96impl Display for ParsedSamplingRules {
97    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
98        write!(
99            f,
100            "{}",
101            serde_json::to_string(&self.rules).unwrap_or_default()
102        )
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    // --- SamplingRuleConfig ---
111
112    #[test]
113    fn test_sampling_rule_config_defaults() {
114        let config = SamplingRuleConfig::default();
115        assert_eq!(config.sample_rate, 0.0);
116        assert!(config.service.is_none());
117        assert!(config.name.is_none());
118        assert!(config.resource.is_none());
119        assert!(config.tags.is_empty());
120        // `Default` matches the serde default for `provenance`.
121        assert_eq!(config.provenance, "default");
122    }
123
124    #[test]
125    fn test_sampling_rule_config_default_matches_serde_default() {
126        // Constructing from an empty-but-valid JSON object must yield the same value
127        // as `Default::default()`.
128        let from_serde: SamplingRuleConfig =
129            serde_json::from_str(r#"{"sample_rate": 0.0}"#).unwrap();
130        assert_eq!(from_serde, SamplingRuleConfig::default());
131    }
132
133    #[test]
134    fn test_sampling_rule_config_serde_default_provenance() {
135        // When provenance is absent from JSON, serde fills it in as "default"
136        let json = r#"{"sample_rate": 0.5}"#;
137        let config: SamplingRuleConfig = serde_json::from_str(json).unwrap();
138        assert_eq!(config.provenance, "default");
139    }
140
141    #[test]
142    fn test_sampling_rule_config_deserialize_full() {
143        let json = r#"{
144            "sample_rate": 0.5,
145            "service": "my-service",
146            "name": "http.*",
147            "resource": "/api/*",
148            "tags": {"env": "prod"},
149            "provenance": "customer"
150        }"#;
151        let config: SamplingRuleConfig = serde_json::from_str(json).unwrap();
152        assert_eq!(config.sample_rate, 0.5);
153        assert_eq!(config.service.as_deref(), Some("my-service"));
154        assert_eq!(config.name.as_deref(), Some("http.*"));
155        assert_eq!(config.resource.as_deref(), Some("/api/*"));
156        assert_eq!(config.tags.get("env").map(String::as_str), Some("prod"));
157        assert_eq!(config.provenance, "customer");
158    }
159
160    #[test]
161    fn test_sampling_rule_config_deserialize_minimal() {
162        let json = r#"{"sample_rate": 1.0}"#;
163        let config: SamplingRuleConfig = serde_json::from_str(json).unwrap();
164        assert_eq!(config.sample_rate, 1.0);
165        assert!(config.service.is_none());
166        assert_eq!(config.provenance, "default");
167    }
168
169    #[test]
170    fn test_sampling_rule_config_roundtrip() {
171        let original = SamplingRuleConfig {
172            sample_rate: 0.25,
173            service: Some("svc".into()),
174            name: Some("op".into()),
175            resource: Some("/res".into()),
176            tags: HashMap::from([("k".into(), "v".into())]),
177            provenance: "dynamic".into(),
178        };
179        let json = serde_json::to_string(&original).unwrap();
180        let restored: SamplingRuleConfig = serde_json::from_str(&json).unwrap();
181        assert_eq!(original, restored);
182    }
183
184    #[test]
185    fn test_sampling_rule_config_display() {
186        let config = SamplingRuleConfig {
187            sample_rate: 1.0,
188            service: Some("svc".into()),
189            ..Default::default()
190        };
191        let s = config.to_string();
192        assert!(s.contains("sample_rate"));
193        assert!(s.contains("svc"));
194    }
195
196    // --- ParsedSamplingRules ---
197
198    #[test]
199    fn test_parsed_sampling_rules_empty_string() {
200        let parsed: ParsedSamplingRules = "".parse().unwrap();
201        assert!(parsed.rules.is_empty());
202    }
203
204    #[test]
205    fn test_parsed_sampling_rules_whitespace_only() {
206        let parsed: ParsedSamplingRules = "   ".parse().unwrap();
207        assert!(parsed.rules.is_empty());
208    }
209
210    #[test]
211    fn test_parsed_sampling_rules_valid_json() {
212        let json = r#"[{"sample_rate": 0.5, "service": "svc"}, {"sample_rate": 1.0}]"#;
213        let parsed: ParsedSamplingRules = json.parse().unwrap();
214        assert_eq!(parsed.rules.len(), 2);
215        assert_eq!(parsed.rules[0].sample_rate, 0.5);
216        assert_eq!(parsed.rules[1].sample_rate, 1.0);
217    }
218
219    #[test]
220    fn test_parsed_sampling_rules_invalid_json() {
221        let result: Result<ParsedSamplingRules, _> = "not json".parse();
222        assert!(result.is_err());
223    }
224
225    #[test]
226    fn test_parsed_sampling_rules_deref() {
227        let json = r#"[{"sample_rate": 0.5}]"#;
228        let parsed: ParsedSamplingRules = json.parse().unwrap();
229        // Deref to &[SamplingRuleConfig]
230        assert_eq!(parsed.len(), 1);
231        assert_eq!(parsed[0].sample_rate, 0.5);
232    }
233
234    #[test]
235    fn test_parsed_sampling_rules_into_vec() {
236        let json = r#"[{"sample_rate": 0.5}, {"sample_rate": 1.0}]"#;
237        let parsed: ParsedSamplingRules = json.parse().unwrap();
238        let vec: Vec<SamplingRuleConfig> = parsed.into();
239        assert_eq!(vec.len(), 2);
240    }
241
242    #[test]
243    fn test_parsed_sampling_rules_display() {
244        let json = r#"[{"sample_rate":0.5}]"#;
245        let parsed: ParsedSamplingRules = json.parse().unwrap();
246        let s = parsed.to_string();
247        assert!(s.contains("sample_rate"));
248        assert!(s.contains("0.5"));
249    }
250
251    #[test]
252    fn test_parsed_sampling_rules_default_is_empty() {
253        let parsed = ParsedSamplingRules::default();
254        assert!(parsed.rules.is_empty());
255    }
256}