Skip to main content

libdd_sampling/
agent_service_sampler.rs

1// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{
5    collections::HashMap,
6    sync::{Arc, RwLock},
7};
8
9use libdd_common::RwLockExt;
10
11use crate::rate_sampler::RateSampler;
12
13#[derive(Debug, serde::Deserialize)]
14pub struct AgentRates<'a> {
15    #[serde(borrow)]
16    pub rate_by_service: Option<HashMap<&'a str, f64>>,
17}
18
19#[derive(Debug, Default, Clone)]
20pub struct ServicesSampler {
21    inner: Arc<RwLock<HashMap<String, RateSampler>>>,
22}
23
24impl ServicesSampler {
25    pub fn get(&self, service: &str) -> Option<RateSampler> {
26        self.inner.read_or_panic().get(service).cloned()
27    }
28
29    pub fn update_rates<I: IntoIterator<Item = (String, f64)>>(&self, rates: I) {
30        let new_rates: HashMap<_, _> = rates
31            .into_iter()
32            .map(|(s, r)| (s, RateSampler::new(r)))
33            .collect();
34        *self.inner.write_or_panic() = new_rates;
35    }
36
37    // Test-only inspection helpers.
38
39    #[cfg(test)]
40    pub(crate) fn is_empty(&self) -> bool {
41        self.inner.read_or_panic().is_empty()
42    }
43
44    #[cfg(test)]
45    pub(crate) fn len(&self) -> usize {
46        self.inner.read_or_panic().len()
47    }
48
49    #[cfg(test)]
50    pub(crate) fn contains_key(&self, service: &str) -> bool {
51        self.inner.read_or_panic().contains_key(service)
52    }
53}