datadog_agent_apm/sampling/
remote_config.rs

1//! Loads trace sampling settings from the `APM_SAMPLING` Remote Configuration product.
2//!
3//! The product can change the priority and error samplers' target traces per second and turn rare sampling on or off.
4//! A configuration may set values for every environment or override individual fields for a specific environment.
5//! Unset fields use the consumer's local settings. This follows the Datadog Agent's
6//! `pkg/trace/remoteconfighandler/remote_config_handler.go`.
7
8use std::fmt;
9use std::marker::PhantomData;
10use std::sync::Arc;
11
12use datadog_agent_remote_config::{
13    decode_json, ApplyError, ConfigId, JsonError, ProductDecoder, RemoteConfigurationClient, Subscription,
14};
15use serde::de::value::MapAccessDeserializer;
16use serde::de::{self, IgnoredAny, MapAccess, Visitor};
17use serde::{Deserialize, Deserializer};
18
19/// Optional sampling settings for all environments or one named environment.
20///
21/// Missing and `null` fields leave that setting unset; `0` and `false` are explicit values.
22#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
23struct SamplerEnvConfig {
24    /// Target traces per second for the priority sampler.
25    #[serde(
26        default,
27        rename = "priority_sampler_target_TPS",
28        alias = "priority_sampler_target_tps"
29    )]
30    priority_sampler_target_tps: Option<f64>,
31
32    /// Target traces per second for the errors sampler.
33    #[serde(default, rename = "errors_sampler_target_TPS", alias = "errors_sampler_target_tps")]
34    errors_sampler_target_tps: Option<f64>,
35
36    /// Whether the rare sampler is enabled.
37    #[serde(default)]
38    rare_sampler_enabled: Option<bool>,
39}
40
41/// An environment name and its sampling overrides.
42#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
43struct EnvAndConfig {
44    /// Environment to match against the normalized local environment; the payload name is used as given.
45    #[serde(default, deserialize_with = "null_as_default")]
46    env: String,
47
48    /// The settings for `env`.
49    #[serde(default, deserialize_with = "null_or_object")]
50    config: SamplerEnvConfig,
51}
52
53/// Settings supplied by one `APM_SAMPLING` configuration.
54///
55/// The payload and each settings value in it must be a JSON object or `null`, as the Datadog Agent's JSON decoder
56/// requires; serde's derived visitors would also accept an array of field values, such as `[]`.
57///
58/// Call [`resolve`](Self::resolve) to combine the configuration with local settings for one environment.
59///
60/// ```json
61/// {
62///   "all_envs": { "priority_sampler_target_TPS": 10, "errors_sampler_target_TPS": 5, "rare_sampler_enabled": true },
63///   "by_env": [{ "env": "prod", "config": { "rare_sampler_enabled": false } }]
64/// }
65/// ```
66///
67/// Missing and `null` fields fall back through `all_envs` to local settings; `0` and `false` are explicit values. A
68/// payload that sets nothing, such as `{}` or `null`, is still an assigned configuration
69/// ([`RemoteSampling::Remote`]) and resolves entirely to local settings, not to previous remote values. This differs
70/// from an empty assignment, which leaves the last accepted configuration in place.
71///
72/// The Datadog Agent sends the `*_TPS` field names. Lowercase `*_tps` names also work;
73/// unknown fields are ignored.
74#[derive(Clone, Debug, Default, PartialEq)]
75pub struct SamplerConfig {
76    /// Settings for every environment.
77    all_envs: SamplerEnvConfig,
78
79    /// Settings for individual environments.
80    by_env: Vec<EnvAndConfig>,
81}
82
83impl<'de> Deserialize<'de> for SamplerConfig {
84    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
85        deserializer.deserialize_map(SamplerConfigVisitor)
86    }
87}
88
89/// Reads [`SamplerConfig`] from a JSON object only.
90struct SamplerConfigVisitor;
91
92impl<'de> Visitor<'de> for SamplerConfigVisitor {
93    type Value = SamplerConfig;
94
95    fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
96        f.write_str("a sampling configuration object")
97    }
98
99    fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Self::Value, A::Error> {
100        #[derive(Deserialize)]
101        #[serde(field_identifier, rename_all = "snake_case")]
102        enum Field {
103            AllEnvs,
104            ByEnv,
105            #[serde(other)]
106            Unknown,
107        }
108
109        let mut all_envs = None;
110        let mut by_env = None;
111        while let Some(field) = map.next_key()? {
112            match field {
113                Field::AllEnvs if all_envs.is_some() => return Err(de::Error::duplicate_field("all_envs")),
114                Field::AllEnvs => {
115                    all_envs = Some(
116                        map.next_value::<Option<Object<SamplerEnvConfig>>>()?
117                            .unwrap_or_default()
118                            .0,
119                    )
120                }
121                Field::ByEnv if by_env.is_some() => return Err(de::Error::duplicate_field("by_env")),
122                Field::ByEnv => {
123                    // A `null` list is empty and a `null` entry is an empty entry, as in the Datadog Agent.
124                    let entries = map.next_value::<Option<Vec<Option<Object<EnvAndConfig>>>>>()?;
125                    by_env = Some(
126                        entries
127                            .unwrap_or_default()
128                            .into_iter()
129                            .map(|entry| entry.unwrap_or_default().0)
130                            .collect(),
131                    );
132                }
133                Field::Unknown => {
134                    map.next_value::<IgnoredAny>()?;
135                }
136            }
137        }
138        Ok(SamplerConfig {
139            all_envs: all_envs.unwrap_or_default(),
140            by_env: by_env.unwrap_or_default(),
141        })
142    }
143}
144
145/// Deserializes `T` from a JSON object only, rejecting the array form that serde's derived visitors also accept.
146///
147/// `T`'s own `Deserialize` implementation reads the object's fields, so field names, aliases, and defaults still come
148/// from its derive.
149#[derive(Default)]
150struct Object<T>(T);
151
152impl<'de, T: Deserialize<'de>> Deserialize<'de> for Object<T> {
153    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
154        struct ObjectVisitor<T>(PhantomData<T>);
155
156        impl<'de, T: Deserialize<'de>> Visitor<'de> for ObjectVisitor<T> {
157            type Value = T;
158
159            fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
160                f.write_str("an object")
161            }
162
163            fn visit_map<A: MapAccess<'de>>(self, map: A) -> Result<Self::Value, A::Error> {
164                T::deserialize(MapAccessDeserializer::new(map))
165            }
166        }
167
168        deserializer.deserialize_map(ObjectVisitor(PhantomData)).map(Self)
169    }
170}
171
172/// Treats `null` like a missing field, as the Datadog Agent's JSON decoder does.
173fn null_as_default<'de, D, T>(deserializer: D) -> Result<T, D::Error>
174where
175    D: Deserializer<'de>,
176    T: Default + Deserialize<'de>,
177{
178    Ok(Option::<T>::deserialize(deserializer)?.unwrap_or_default())
179}
180
181/// Treats `null` like a missing field and requires any other value to be an object.
182fn null_or_object<'de, D, T>(deserializer: D) -> Result<T, D::Error>
183where
184    D: Deserializer<'de>,
185    T: Default + Deserialize<'de>,
186{
187    Ok(null_as_default::<D, Object<T>>(deserializer)?.0)
188}
189
190/// Effective sampling settings, either from local configuration or resolved remote overrides.
191#[derive(Clone, Copy, Debug, PartialEq)]
192pub struct SamplerSettings {
193    /// Target traces per second for the priority sampler.
194    pub target_traces_per_second: f64,
195
196    /// Target traces per second for the errors sampler.
197    pub errors_per_second: f64,
198
199    /// Whether the rare sampler is enabled.
200    pub rare_sampler_enabled: bool,
201}
202
203impl SamplerConfig {
204    /// Combines remote overrides with local sampling settings for the Agent's environment.
205    ///
206    /// The result contains target traces per second for priority and error sampling, plus an on/off switch for rare
207    /// sampling. Priority sampling honors existing decisions to keep or drop traces and uses its target to compute
208    /// sampling rates. Error sampling can keep error traces that would otherwise be dropped; rare sampling keeps
209    /// examples of rarely seen operations.
210    ///
211    /// `env` identifies the Agent's deployment environment, such as `prod`, not an individual trace's environment.
212    /// The caller supplies the Agent's resolved, normalized environment. This method neither derives that environment
213    /// nor normalizes names: a remote name of `Prod` does not match `prod`.
214    ///
215    /// This method selects the last `by_env` entry whose name exactly matches `env`. Earlier matching entries
216    /// contribute nothing, even if the selected entry leaves fields unset. For each setting independently, it uses:
217    ///
218    /// 1. The value in the selected entry's `config`, if set.
219    /// 2. Otherwise, the value in `all_envs`, the remote settings shared by all environments.
220    /// 3. Otherwise, the value in `static_settings`, the original local settings, not previous remote values.
221    ///
222    /// Missing or `null` fields are unset; `0` and `false` are explicit overrides. If no entry matches `env`, each
223    /// setting falls back to `all_envs`, then `static_settings`. Numeric targets pass through without validation or
224    /// clamping.
225    ///
226    /// For example, suppose the selected `prod` entry sets only the priority target to 20 and `all_envs` sets only
227    /// the error target to 5. The result uses 20 for priority sampling, 5 for error sampling, and the local setting
228    /// for rare sampling.
229    pub fn resolve(&self, env: &str, static_settings: &SamplerSettings) -> SamplerSettings {
230        let for_env = self
231            .by_env
232            .iter()
233            .rfind(|entry| entry.env == env)
234            .map(|entry| &entry.config);
235
236        let pick = |field: fn(&SamplerEnvConfig) -> Option<f64>, fallback: f64| {
237            for_env
238                .and_then(field)
239                .or_else(|| field(&self.all_envs))
240                .unwrap_or(fallback)
241        };
242        SamplerSettings {
243            target_traces_per_second: pick(
244                |c| c.priority_sampler_target_tps,
245                static_settings.target_traces_per_second,
246            ),
247            errors_per_second: pick(|c| c.errors_sampler_target_tps, static_settings.errors_per_second),
248            rare_sampler_enabled: for_env
249                .and_then(|c| c.rare_sampler_enabled)
250                .or(self.all_envs.rare_sampler_enabled)
251                .unwrap_or(static_settings.rare_sampler_enabled),
252        }
253    }
254}
255
256/// The latest accepted `APM_SAMPLING` state.
257#[derive(Clone, Debug, PartialEq)]
258pub enum RemoteSampling {
259    /// An empty assignment accepted before any concrete configuration; use local sampling settings.
260    ///
261    /// After accepting a concrete configuration, the subscription retains it on withdrawal rather than publishing
262    /// this marker.
263    Unassigned,
264
265    /// The last accepted concrete sampling configuration and its ID.
266    ///
267    /// Retained after withdrawal, it may no longer be remotely assigned.
268    Remote {
269        /// The configuration's ID.
270        id: ConfigId,
271
272        /// The configuration.
273        config: SamplerConfig,
274    },
275}
276
277/// Why an `APM_SAMPLING` assignment was rejected.
278///
279/// The client reports this error to the remote configuration service.
280#[derive(Debug)]
281pub(super) enum SamplingError {
282    /// A configuration's payload is not a valid sampling configuration.
283    InvalidPayload(JsonError),
284
285    /// More than one configuration was assigned.
286    TooManyConfigurations {
287        /// How many configurations were assigned, including ones that failed to decode.
288        assigned: usize,
289    },
290
291    /// The one assigned configuration failed to decode.
292    NoValidConfiguration,
293}
294
295impl fmt::Display for SamplingError {
296    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
297        match self {
298            Self::InvalidPayload(error) => error.fmt(f),
299            Self::TooManyConfigurations { assigned } => write!(
300                f,
301                "APM_SAMPLING carries {assigned} configurations; expected at most one."
302            ),
303            Self::NoValidConfiguration => f.write_str("The assigned APM_SAMPLING configuration is not valid."),
304        }
305    }
306}
307
308impl std::error::Error for SamplingError {}
309
310impl ApplyError for SamplingError {
311    fn apply_error(&self) -> String {
312        self.to_string()
313    }
314}
315
316/// Watches the latest accepted remote trace sampling configuration.
317///
318/// Create with [`new`](Self::new) and read [`current`](Self::current) before processing a batch of traces. This is a
319/// latest-state watch, not a queue: successive concrete configurations can coalesce before you read them.
320///
321/// Once accepted, a concrete configuration remains available across withdrawal, cache expiry, polling-worker restart,
322/// and rejected updates. An initially accepted empty assignment publishes [`RemoteSampling::Unassigned`], for which
323/// the consumer uses local settings. Later withdrawals never replace an accepted concrete configuration with that
324/// marker.
325///
326/// The default subscription is inert: it receives no updates and contacts no service. Clones share a registration;
327/// dropping the last clone unsubscribes.
328#[derive(Clone, Debug)]
329pub struct TraceSamplingSubscription {
330    subscription: Subscription<RemoteSampling, SamplingError>,
331}
332
333impl TraceSamplingSubscription {
334    /// Subscribes to `APM_SAMPLING` on `client`.
335    ///
336    /// # Errors
337    ///
338    /// Returns an error if `client` already has a live subscription to `APM_SAMPLING`.
339    pub fn new(client: &RemoteConfigurationClient) -> Result<Self, datadog_agent_remote_config::Error> {
340        Ok(Self::from_subscription(client.subscribe::<SamplingDecoder>()?))
341    }
342
343    /// Wraps a subscription to the sampling decoder.
344    pub(super) fn from_subscription(subscription: Subscription<RemoteSampling, SamplingError>) -> Self {
345        Self { subscription }
346    }
347
348    /// Returns the current accepted snapshot immediately.
349    ///
350    /// Returns `None` until the first accepted snapshot, not merely the first poll: polling may fail or reject an
351    /// update. An initially accepted empty assignment publishes [`RemoteSampling::Unassigned`].
352    ///
353    /// Each newly accepted concrete configuration gets a new `Arc`, even if its values are unchanged. Later withdrawals
354    /// and rejected updates preserve the current `Arc`; use [`Arc::ptr_eq`] to compare snapshot identity.
355    pub fn current(&self) -> Option<Arc<RemoteSampling>> {
356        self.subscription.current()
357    }
358}
359
360impl Default for TraceSamplingSubscription {
361    fn default() -> Self {
362        Self::from_subscription(Subscription::inert())
363    }
364}
365
366/// Accepts at most one `APM_SAMPLING` configuration per update.
367///
368/// No assignment produces [`RemoteSampling::Unassigned`], which the client publishes only if nothing was accepted
369/// before; see [`keeps_previous`](ProductDecoder::keeps_previous). An invalid or multiple-configuration assignment is
370/// rejected, leaving subscribers with the last accepted update.
371#[derive(Default)]
372pub(super) struct SamplingDecoder {
373    assigned: usize,
374    chosen: Option<(ConfigId, SamplerConfig)>,
375}
376
377impl ProductDecoder for SamplingDecoder {
378    const PRODUCT: &'static str = "APM_SAMPLING";
379
380    type Snapshot = RemoteSampling;
381    type Error = SamplingError;
382
383    fn decode(&mut self, id: &ConfigId, payload: &[u8]) -> Result<(), Self::Error> {
384        // Count before parsing so invalid configurations still count toward the one-config limit.
385        self.assigned += 1;
386        // The Datadog Agent treats a `null` payload as a configuration with no fields set.
387        let config = decode_json::<Option<SamplerConfig>>(payload)
388            .map_err(SamplingError::InvalidPayload)?
389            .unwrap_or_default();
390        self.chosen = Some((id.clone(), config));
391        Ok(())
392    }
393
394    fn build(self) -> Result<Self::Snapshot, Self::Error> {
395        if self.assigned > 1 {
396            // The client reports this rejection.
397            return Err(SamplingError::TooManyConfigurations {
398                assigned: self.assigned,
399            });
400        }
401        match self.chosen {
402            Some((id, config)) => Ok(RemoteSampling::Remote { id, config }),
403            None if self.assigned == 0 => Ok(RemoteSampling::Unassigned),
404            None => Err(SamplingError::NoValidConfiguration),
405        }
406    }
407
408    fn keeps_previous(snapshot: &Self::Snapshot) -> bool {
409        matches!(snapshot, RemoteSampling::Unassigned)
410    }
411}
412
413#[cfg(test)]
414mod tests {
415    use datadog_agent_remote_config::TestPublisher;
416
417    use super::*;
418    use crate::sampling::test_util::TestSamplingPublisher;
419
420    const STATIC: SamplerSettings = SamplerSettings {
421        target_traces_per_second: 41.0,
422        errors_per_second: 41.0,
423        rare_sampler_enabled: true,
424    };
425
426    fn decode_all(assignment: &[(&str, &str)]) -> (Vec<Result<(), String>>, Result<RemoteSampling, String>) {
427        let mut decoder = SamplingDecoder::default();
428        let verdicts = assignment
429            .iter()
430            .map(|(id, payload)| {
431                decoder
432                    .decode(&ConfigId::new(*id), payload.as_bytes())
433                    .map_err(|e| e.apply_error())
434            })
435            .collect();
436        (verdicts, decoder.build().map_err(|e| e.apply_error()))
437    }
438
439    fn parse(json: &str) -> SamplerConfig {
440        decode_json(json.as_bytes()).expect("payload should decode")
441    }
442
443    fn env_config(priority: Option<f64>, errors: Option<f64>, rare: Option<bool>) -> SamplerEnvConfig {
444        SamplerEnvConfig {
445            priority_sampler_target_tps: priority,
446            errors_sampler_target_tps: errors,
447            rare_sampler_enabled: rare,
448        }
449    }
450
451    #[test]
452    fn payload_uses_the_agent_field_names() {
453        let config = parse(
454            r#"{"all_envs":{"priority_sampler_target_TPS":1.5,"errors_sampler_target_TPS":2.5,
455                "rare_sampler_enabled":true},
456                "by_env":[{"env":"prod","config":{"priority_sampler_target_TPS":3}}]}"#,
457        );
458        assert_eq!(config.all_envs, env_config(Some(1.5), Some(2.5), Some(true)));
459        assert_eq!(
460            config.by_env,
461            [EnvAndConfig {
462                env: "prod".to_owned(),
463                config: env_config(Some(3.0), None, None)
464            }]
465        );
466    }
467
468    #[test]
469    fn payload_accepts_lowercase_tps_field_names() {
470        let config = parse(r#"{"all_envs":{"priority_sampler_target_tps":1,"errors_sampler_target_tps":2}}"#);
471        assert_eq!(config.all_envs, env_config(Some(1.0), Some(2.0), None));
472    }
473
474    #[test]
475    fn payload_treats_null_and_absent_fields_as_unset() {
476        assert_eq!(parse(r#"{"all_envs":null,"by_env":null}"#), SamplerConfig::default());
477        assert_eq!(
478            parse(r#"{"by_env":[{"env":"prod","config":null}]}"#).by_env,
479            [EnvAndConfig {
480                env: "prod".to_owned(),
481                config: SamplerEnvConfig::default()
482            }]
483        );
484        assert_eq!(
485            parse(
486                r#"{"all_envs":{"priority_sampler_target_TPS":null,"errors_sampler_target_TPS":null,
487                    "rare_sampler_enabled":null}}"#
488            ),
489            SamplerConfig::default()
490        );
491        assert_eq!(parse("{}"), SamplerConfig::default());
492    }
493
494    #[test]
495    fn payload_honors_explicit_zero_and_false() {
496        let config = parse(
497            r#"{"all_envs":{"priority_sampler_target_TPS":0,"errors_sampler_target_TPS":0,
498                "rare_sampler_enabled":false}}"#,
499        );
500        assert_eq!(config.all_envs, env_config(Some(0.0), Some(0.0), Some(false)));
501        assert_eq!(
502            config.resolve("", &STATIC),
503            SamplerSettings {
504                target_traces_per_second: 0.0,
505                errors_per_second: 0.0,
506                rare_sampler_enabled: false,
507            }
508        );
509    }
510
511    #[test]
512    fn payload_ignores_unknown_fields() {
513        let config = parse(r#"{"future":1,"all_envs":{"rare_sampler_enabled":true,"other":"x"}}"#);
514        assert_eq!(config.all_envs, env_config(None, None, Some(true)));
515    }
516
517    // Based on the Datadog Agent's `TestEnvPrecedence`, with additional fallback cases:
518    // https://github.com/DataDog/datadog-agent/blob/17ecddf4e3e83ccbb0e68aeb99461d1a1d902927/pkg/trace/remoteconfighandler/remote_config_handler_test.go#L289-L322
519    #[test]
520    fn resolve_prefers_env_then_all_envs_then_static() {
521        let all_envs = env_config(Some(42.0), Some(42.0), Some(true));
522        let for_env = |env: &str, config: SamplerEnvConfig| EnvAndConfig {
523            env: env.to_owned(),
524            config,
525        };
526        let settings = |priority, errors, rare| SamplerSettings {
527            target_traces_per_second: priority,
528            errors_per_second: errors,
529            rare_sampler_enabled: rare,
530        };
531
532        let cases = [
533            (
534                "env entry beats all_envs",
535                all_envs.clone(),
536                vec![for_env("agent-env", env_config(Some(43.0), Some(43.0), Some(false)))],
537                settings(43.0, 43.0, false),
538            ),
539            (
540                "all_envs applies without an env entry",
541                all_envs.clone(),
542                vec![for_env("other-env", env_config(Some(43.0), Some(43.0), Some(false)))],
543                settings(42.0, 42.0, true),
544            ),
545            (
546                "static applies when neither sets a field",
547                SamplerEnvConfig::default(),
548                vec![],
549                STATIC,
550            ),
551            (
552                "a field the env entry leaves unset falls back to all_envs",
553                all_envs.clone(),
554                vec![for_env("agent-env", env_config(Some(43.0), None, None))],
555                settings(43.0, 42.0, true),
556            ),
557            (
558                "a field neither sets falls back to static per field",
559                env_config(None, Some(42.0), None),
560                vec![for_env("agent-env", env_config(Some(43.0), None, Some(false)))],
561                settings(43.0, 42.0, false),
562            ),
563            (
564                "the last matching env entry wins",
565                SamplerEnvConfig::default(),
566                vec![
567                    for_env("agent-env", env_config(Some(43.0), Some(43.0), Some(false))),
568                    for_env("agent-env", env_config(Some(44.0), None, None)),
569                ],
570                settings(44.0, 41.0, true),
571            ),
572        ];
573
574        for (name, all_envs, by_env, expected) in cases {
575            let config = SamplerConfig { all_envs, by_env };
576            assert_eq!(config.resolve("agent-env", &STATIC), expected, "{name}");
577        }
578    }
579
580    #[test]
581    fn resolve_compares_the_payload_env_as_given() {
582        let config = SamplerConfig {
583            all_envs: SamplerEnvConfig::default(),
584            by_env: vec![EnvAndConfig {
585                env: "Agent-Env".to_owned(),
586                config: env_config(Some(43.0), None, None),
587            }],
588        };
589        assert_eq!(config.resolve("agent-env", &STATIC).target_traces_per_second, 41.0);
590        assert_eq!(config.resolve("Agent-Env", &STATIC).target_traces_per_second, 43.0);
591    }
592
593    #[test]
594    fn empty_assignment_is_unassigned() {
595        assert_eq!(decode_all(&[]).1, Ok(RemoteSampling::Unassigned));
596    }
597
598    #[test]
599    fn single_valid_configuration_is_chosen() {
600        let (verdicts, snapshot) = decode_all(&[("a", r#"{"all_envs":{"rare_sampler_enabled":false}}"#)]);
601        assert_eq!(verdicts, [Ok(())]);
602        assert_eq!(
603            snapshot,
604            Ok(RemoteSampling::Remote {
605                id: ConfigId::new("a"),
606                config: SamplerConfig {
607                    all_envs: env_config(None, None, Some(false)),
608                    by_env: vec![],
609                },
610            })
611        );
612    }
613
614    #[test]
615    fn null_payload_is_a_configuration_with_nothing_set() {
616        let (verdicts, snapshot) = decode_all(&[("a", "null")]);
617        assert_eq!(verdicts, [Ok(())]);
618        assert_eq!(
619            snapshot,
620            Ok(RemoteSampling::Remote {
621                id: ConfigId::new("a"),
622                config: SamplerConfig::default(),
623            })
624        );
625    }
626
627    #[test]
628    fn null_env_is_the_empty_env() {
629        let (_, snapshot) = decode_all(&[(
630            "a",
631            r#"{"by_env":[{"env":null,"config":{"priority_sampler_target_TPS":3}}]}"#,
632        )]);
633        let Ok(RemoteSampling::Remote { config, .. }) = snapshot else {
634            panic!("the configuration should be chosen: {snapshot:?}");
635        };
636        assert_eq!(
637            config.by_env,
638            [EnvAndConfig {
639                env: String::new(),
640                config: env_config(Some(3.0), None, None)
641            }]
642        );
643    }
644
645    #[test]
646    fn null_by_env_element_is_an_empty_entry() {
647        let (_, snapshot) = decode_all(&[(
648            "a",
649            r#"{"by_env":[null,{"env":"prod","config":{"priority_sampler_target_TPS":3}}]}"#,
650        )]);
651        let Ok(RemoteSampling::Remote { config, .. }) = snapshot else {
652            panic!("the configuration should be chosen: {snapshot:?}");
653        };
654        assert_eq!(
655            config.by_env,
656            [
657                EnvAndConfig::default(),
658                EnvAndConfig {
659                    env: "prod".to_owned(),
660                    config: env_config(Some(3.0), None, None)
661                }
662            ]
663        );
664    }
665
666    #[test]
667    fn invalid_payload_error_gives_only_its_position() {
668        let (verdicts, snapshot) = decode_all(&[("a", r#"{"all_envs":{"rare_sampler_enabled":"secret"}}"#)]);
669        assert_eq!(
670            verdicts,
671            [Err("Payload is malformed JSON at line 1, column 44.".to_owned())]
672        );
673        assert_eq!(
674            snapshot,
675            Err("The assigned APM_SAMPLING configuration is not valid.".to_owned())
676        );
677    }
678
679    #[test]
680    fn more_than_one_configuration_is_rejected_with_the_count() {
681        let (verdicts, snapshot) = decode_all(&[("a", "{}"), ("b", "not json"), ("c", "{}")]);
682        assert_eq!(
683            verdicts,
684            [
685                Ok(()),
686                Err("Payload is malformed JSON at line 1, column 2.".to_owned()),
687                Ok(())
688            ]
689        );
690        assert_eq!(
691            snapshot,
692            Err("APM_SAMPLING carries 3 configurations; expected at most one.".to_owned())
693        );
694    }
695
696    #[test]
697    fn configurations_that_fail_to_decode_count_toward_the_limit() {
698        let (_, snapshot) = decode_all(&[("a", "not json"), ("b", "{}")]);
699        assert_eq!(
700            snapshot,
701            Err("APM_SAMPLING carries 2 configurations; expected at most one.".to_owned())
702        );
703    }
704
705    #[test]
706    fn default_subscription_never_has_an_update() {
707        assert_eq!(TraceSamplingSubscription::default().current(), None);
708    }
709
710    #[test]
711    fn new_subscribes_to_apm_sampling_until_the_last_clone_drops() {
712        let client = RemoteConfigurationClient::without_agent();
713
714        let subscription = TraceSamplingSubscription::new(&client).expect("the product is not yet subscribed");
715        assert_eq!(client.subscribed_products(), ["APM_SAMPLING"]);
716        assert!(matches!(
717            TraceSamplingSubscription::new(&client),
718            Err(datadog_agent_remote_config::Error::AlreadySubscribed { product }) if product == "APM_SAMPLING"
719        ));
720
721        // A surviving clone keeps the registration alive.
722        let clone = subscription.clone();
723        drop(subscription);
724        assert!(TraceSamplingSubscription::new(&client).is_err());
725
726        drop(clone);
727        assert!(client.subscribed_products().is_empty());
728        let _resubscribed = TraceSamplingSubscription::new(&client).expect("the product is free again");
729        assert_eq!(client.subscribed_products(), ["APM_SAMPLING"]);
730    }
731
732    #[test]
733    fn subscription_has_no_update_before_one_is_accepted() {
734        let (_publisher, subscription) = TestSamplingPublisher::new();
735        assert_eq!(subscription.current(), None);
736    }
737
738    #[test]
739    fn clones_observe_accepted_updates() {
740        let (publisher, subscription) = TestSamplingPublisher::new();
741        let clone = subscription.clone();
742
743        publisher.assign([("a", r#"{"all_envs":{"rare_sampler_enabled":true}}"#)]);
744        let first = subscription.current().expect("a configuration is accepted");
745        assert!(matches!(&*first, RemoteSampling::Remote { id, .. } if id == &ConfigId::new("a")));
746        assert!(Arc::ptr_eq(
747            &first,
748            &clone.current().expect("clones see the same update")
749        ));
750    }
751
752    // The Datadog Agent decodes into structs with `json.Unmarshal`, which rejects an array where an object belongs:
753    // https://github.com/DataDog/datadog-agent/blob/17ecddf4e3e83ccbb0e68aeb99461d1a1d902927/pkg/remoteconfig/state/products/apmsampling/sampler_config.go
754    // Serde's derived visitors accept one, so `[]` would decode as a configuration with nothing set and
755    // `{"all_envs":[0,0,false]}` would set both targets to zero and turn off the rare sampler.
756    #[test]
757    fn payload_rejects_non_objects_where_the_agent_requires_an_object() {
758        let invalid = [
759            "[]",
760            "[null,null]",
761            "5",
762            r#""text""#,
763            "true",
764            r#"{"all_envs":[0,0,false]}"#,
765            r#"{"all_envs":[]}"#,
766            r#"{"all_envs":5}"#,
767            r#"{"by_env":{}}"#,
768            r#"{"by_env":[[]]}"#,
769            r#"{"by_env":[["prod",{}]]}"#,
770            r#"{"by_env":[5]}"#,
771            r#"{"by_env":[{"env":"prod","config":[0,0,false]}]}"#,
772            r#"{"by_env":[{"env":"prod","config":"x"}]}"#,
773            r#"{"by_env":[{"env":5}]}"#,
774        ];
775        for payload in invalid {
776            let (verdicts, snapshot) = decode_all(&[("a", payload)]);
777            assert!(
778                matches!(&verdicts[..], [Err(reason)] if reason.starts_with("Payload is malformed JSON at line 1")),
779                "{payload}: {verdicts:?}"
780            );
781            assert!(snapshot.is_err(), "{payload}: {snapshot:?}");
782        }
783
784        let valid = [
785            "null",
786            "{}",
787            r#"{"all_envs":null,"by_env":null}"#,
788            r#"{"all_envs":{}}"#,
789            r#"{"by_env":[]}"#,
790            r#"{"by_env":[null,{},{"env":null,"config":null},{"env":"other","config":{}}]}"#,
791        ];
792        for payload in valid {
793            let (verdicts, snapshot) = decode_all(&[("a", payload)]);
794            assert_eq!(verdicts, [Ok(())], "{payload}");
795            let Ok(RemoteSampling::Remote { config, .. }) = snapshot else {
796                panic!("{payload} should be chosen: {snapshot:?}");
797            };
798            assert_eq!(
799                config.resolve("prod", &STATIC),
800                STATIC,
801                "{payload} sets nothing for `prod`"
802            );
803        }
804    }
805
806    #[test]
807    fn payload_rejects_a_duplicate_field() {
808        let (verdicts, _) = decode_all(&[("a", r#"{"by_env":[],"by_env":[]}"#)]);
809        assert!(matches!(&verdicts[..], [Err(_)]), "{verdicts:?}");
810    }
811
812    // A latest-state subscription coalesces publications, so an empty assignment published before the consumer reads
813    // must not replace a configuration it has not applied. The Datadog Agent applies each update as it arrives and
814    // ignores an empty one:
815    // https://github.com/DataDog/datadog-agent/blob/17ecddf4e3e83ccbb0e68aeb99461d1a1d902927/pkg/trace/remoteconfighandler/remote_config_handler.go#L234-L295
816    #[test]
817    fn withdrawal_and_rejection_keep_the_last_accepted_configuration() {
818        let empty = std::iter::empty::<(&str, &str)>;
819        let current_id = |subscription: &TraceSamplingSubscription| match &*subscription.current().expect("accepted") {
820            RemoteSampling::Remote { id, .. } => Some(id.clone()),
821            RemoteSampling::Unassigned => None,
822        };
823        let (publisher, subscription) = TestSamplingPublisher::new();
824
825        publisher.assign(empty());
826        assert_eq!(current_id(&subscription), None, "nothing has been assigned yet");
827
828        publisher.assign([("a", "{}")]);
829        let accepted = subscription.current().expect("a configuration is accepted");
830        for (name, assignment) in [
831            ("an empty assignment", vec![]),
832            ("an array payload", vec![("b", "[]")]),
833            ("an array settings value", vec![("b", r#"{"all_envs":[0,0,false]}"#)]),
834            ("more than one configuration", vec![("b", "{}"), ("c", "{}")]),
835        ] {
836            publisher.assign(assignment);
837            assert!(
838                Arc::ptr_eq(&accepted, &subscription.current().expect("still accepted")),
839                "{name} keeps the last accepted configuration"
840            );
841        }
842
843        publisher.assign([("b", "{}")]);
844        publisher.assign(empty());
845        assert_eq!(
846            current_id(&subscription),
847            Some(ConfigId::new("b")),
848            "the newest configuration is kept"
849        );
850    }
851
852    #[test]
853    fn rejected_assignment_keeps_the_last_update() {
854        let (publisher, subscription) = TestSamplingPublisher::new();
855        publisher.assign([("a", "{}")]);
856        let accepted = subscription.current().expect("a configuration is accepted");
857
858        publisher.assign([("a", "{}"), ("b", "{}")]);
859        assert!(Arc::ptr_eq(
860            &accepted,
861            &subscription.current().expect("the earlier update is retained")
862        ));
863    }
864
865    #[tokio::test]
866    async fn an_empty_assignment_after_a_rejection_clears_it_and_notifies_with_the_last_configuration() {
867        let (publisher, mut subscription) = TestPublisher::<RemoteSampling, SamplingError>::new();
868        publisher.assign::<SamplingDecoder>([("a", "{}")]);
869        let accepted = subscription
870            .changed()
871            .await
872            .expect("one configuration should be accepted");
873
874        publisher.assign::<SamplingDecoder>([("a", "[]")]);
875        assert!(subscription.changed().await.is_err(), "the array payload is rejected");
876
877        publisher.assign::<SamplingDecoder>(std::iter::empty::<(&str, &str)>());
878        let kept = subscription.changed().await.expect("the empty assignment is accepted");
879        assert!(Arc::ptr_eq(&accepted, &kept));
880    }
881}
datadog_agent_apm/sampling/remote_config.rs

datadog_agent_apm/sampling/
remote_config.rs
