Skip to main content

saluki_components/common/datadog/
routing.rs

1//! Routing targets for metrics payloads.
2//!
3//! A routing target is one physical endpoint the Datadog forwarder delivers to: one intake URL paired with one API key.
4//! The [`RoutingTargetCatalog`] enumerates every target once and gives each a stable [`RoutingTargetId`], so an encoder
5//! can address a payload to a set of targets and the forwarder can deliver it to exactly those targets.
6//!
7//! The catalog is built once from configuration and shared by every encoder that addresses payloads and by the
8//! forwarder that delivers them. Sharing one catalog is what keeps the two sides agreeing on what each identifier means.
9//!
10//! A payload addressed with [`MetricsRoutingTargets`] reaches only the targets it names. A payload without it reaches
11//! every target, as it always has.
12
13use std::{collections::HashMap, fmt};
14
15use saluki_error::{ErrorContext as _, GenericError};
16use smallvec::SmallVec;
17
18use super::endpoints::resolve_additional_endpoints;
19
20/// Identifies one routing target within a [`RoutingTargetCatalog`].
21#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
22pub struct RoutingTargetId(u32);
23
24impl RoutingTargetId {
25    fn from_index(index: usize) -> Self {
26        Self(u32::try_from(index).expect("routing target count must fit in a u32"))
27    }
28
29    /// Returns the position of this target in its catalog.
30    pub const fn index(self) -> usize {
31        self.0 as usize
32    }
33}
34
35impl fmt::Display for RoutingTargetId {
36    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37        write!(f, "{}", self.0)
38    }
39}
40
41/// A set of routing targets.
42///
43/// The set is a bitset indexed by [`RoutingTargetId`]. It never holds trailing empty words, so two sets holding the
44/// same targets always compare and hash equal.
45#[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
46pub struct RoutingTargetSet {
47    words: SmallVec<[u64; 1]>,
48}
49
50impl RoutingTargetSet {
51    /// Creates an empty set.
52    pub fn new() -> Self {
53        Self::default()
54    }
55
56    /// Adds a target to the set.
57    pub fn insert(&mut self, id: RoutingTargetId) {
58        let (word, bit) = (id.index() / 64, id.index() % 64);
59        if self.words.len() <= word {
60            self.words.resize(word + 1, 0);
61        }
62        self.words[word] |= 1 << bit;
63    }
64
65    /// Returns `true` if the set holds the given target.
66    pub fn contains(&self, id: RoutingTargetId) -> bool {
67        let (word, bit) = (id.index() / 64, id.index() % 64);
68        self.words.get(word).is_some_and(|word| word & (1 << bit) != 0)
69    }
70
71    /// Returns `true` if the set holds no targets.
72    pub fn is_empty(&self) -> bool {
73        self.words.is_empty()
74    }
75
76    /// Returns the number of targets in the set.
77    pub fn len(&self) -> usize {
78        self.words.iter().map(|word| word.count_ones() as usize).sum()
79    }
80
81    /// Adds every target in `other` to this set.
82    pub fn union_with(&mut self, other: &Self) {
83        if self.words.len() < other.words.len() {
84            self.words.resize(other.words.len(), 0);
85        }
86        for (word, other) in self.words.iter_mut().zip(&other.words) {
87            *word |= other;
88        }
89    }
90
91    /// Returns the targets in the set, in ascending order.
92    pub fn iter(&self) -> impl Iterator<Item = RoutingTargetId> + '_ {
93        self.words.iter().enumerate().flat_map(|(word_index, word)| {
94            (0..64)
95                .filter(move |bit| word & (1 << bit) != 0)
96                .map(move |bit| RoutingTargetId::from_index(word_index * 64 + bit))
97        })
98    }
99}
100
101impl FromIterator<RoutingTargetId> for RoutingTargetSet {
102    fn from_iter<I: IntoIterator<Item = RoutingTargetId>>(iter: I) -> Self {
103        let mut set = Self::new();
104        for id in iter {
105            set.insert(id);
106        }
107        set
108    }
109}
110
111/// How the forwarder reaches a routing target.
112#[derive(Clone, Copy, Debug, Eq, PartialEq)]
113pub enum RoutingTargetKind {
114    /// The primary endpoint, built from `dd_url` or `site` and `api_key`.
115    ///
116    /// When an alternate metrics intake replaces the primary endpoint for metrics, the alternate intake receives the
117    /// payloads addressed to this target.
118    Primary,
119
120    /// One API key of one configured additional endpoint.
121    Additional,
122}
123
124/// One routing target in a [`RoutingTargetCatalog`].
125#[derive(Clone, Debug, Eq, PartialEq)]
126pub struct RoutingTarget {
127    id: RoutingTargetId,
128    kind: RoutingTargetKind,
129    configured_endpoint: String,
130    api_key_index: Option<usize>,
131}
132
133impl RoutingTarget {
134    /// Returns the target's identifier.
135    pub const fn id(&self) -> RoutingTargetId {
136        self.id
137    }
138
139    /// Returns how the forwarder reaches this target.
140    pub const fn kind(&self) -> RoutingTargetKind {
141        self.kind
142    }
143
144    /// Returns the target's endpoint as configuration spells it.
145    ///
146    /// For the primary target this is the effective primary endpoint. For an additional target it is the key of its
147    /// `additional_endpoints` entry.
148    pub fn configured_endpoint(&self) -> &str {
149        &self.configured_endpoint
150    }
151
152    /// Returns the position of this target's API key in its `additional_endpoints` entry.
153    ///
154    /// `None` for the primary target.
155    #[cfg(test)]
156    pub(crate) const fn api_key_index(&self) -> Option<usize> {
157        self.api_key_index
158    }
159}
160
161/// The routing targets of one Datadog forwarder.
162///
163/// Identifiers are assigned deterministically: the primary target comes first, followed by every additional endpoint
164/// API key sorted by configured endpoint and then by the key's position in that endpoint's list.
165#[derive(Clone, Debug, Eq, PartialEq)]
166pub struct RoutingTargetCatalog {
167    targets: Vec<RoutingTarget>,
168}
169
170impl RoutingTargetCatalog {
171    /// Creates a catalog holding the primary endpoint and every configured additional endpoint API key.
172    ///
173    /// Additional endpoint keys that the forwarder skips -- empty keys, repeated keys, and keys that cannot be sent as
174    /// an HTTP header -- are skipped here too, so the catalog holds exactly the endpoints the forwarder builds.
175    ///
176    /// # Errors
177    ///
178    /// Returns an error if an additional endpoint URL is invalid.
179    pub fn new(
180        primary_endpoint: &str, additional_endpoints: &HashMap<String, Vec<String>>,
181    ) -> Result<Self, GenericError> {
182        let mut additional = resolve_additional_endpoints(additional_endpoints)
183            .error_context("Failed parsing/resolving the additional destination endpoints.")?
184            .into_iter()
185            .filter_map(|endpoint| {
186                let (configured_endpoint, api_key_index) = endpoint.additional_endpoint_queue_key()?;
187                Some((configured_endpoint.to_string(), api_key_index))
188            })
189            .collect::<Vec<_>>();
190        additional.sort_unstable();
191
192        let mut targets = Vec::with_capacity(additional.len() + 1);
193        targets.push(RoutingTarget {
194            id: RoutingTargetId::from_index(0),
195            kind: RoutingTargetKind::Primary,
196            configured_endpoint: primary_endpoint.to_string(),
197            api_key_index: None,
198        });
199        for (configured_endpoint, api_key_index) in additional {
200            targets.push(RoutingTarget {
201                id: RoutingTargetId::from_index(targets.len()),
202                kind: RoutingTargetKind::Additional,
203                configured_endpoint,
204                api_key_index: Some(api_key_index),
205            });
206        }
207
208        Ok(Self { targets })
209    }
210
211    /// Returns the primary target.
212    pub fn primary(&self) -> RoutingTargetId {
213        self.targets[0].id
214    }
215
216    /// Returns every target, ordered by identifier.
217    pub fn targets(&self) -> &[RoutingTarget] {
218        &self.targets
219    }
220
221    /// Returns the target with the given identifier.
222    pub fn get(&self, id: RoutingTargetId) -> Option<&RoutingTarget> {
223        self.targets.get(id.index())
224    }
225
226    /// Returns the targets matching `predicate`.
227    pub fn select(&self, mut predicate: impl FnMut(&RoutingTarget) -> bool) -> RoutingTargetSet {
228        self.targets
229            .iter()
230            .filter(|target| predicate(target))
231            .map(RoutingTarget::id)
232            .collect()
233    }
234
235    /// Returns the additional target for the given configured endpoint and API key position, if any.
236    pub(crate) fn find_additional(&self, configured_endpoint: &str, api_key_index: usize) -> Option<RoutingTargetId> {
237        self.targets
238            .iter()
239            .find(|target| {
240                target.kind == RoutingTargetKind::Additional
241                    && target.api_key_index == Some(api_key_index)
242                    && target.configured_endpoint == configured_endpoint
243            })
244            .map(RoutingTarget::id)
245    }
246}
247
248/// The routing targets an encoded metrics payload is addressed to.
249///
250/// Encoders attach this to the payloads they emit, and the forwarder delivers a payload carrying it only to the targets
251/// it names. It is kept in memory only: a payload that is persisted for retry already belongs to one target.
252#[derive(Clone, Debug, Eq, PartialEq)]
253pub struct MetricsRoutingTargets {
254    targets: RoutingTargetSet,
255}
256
257impl MetricsRoutingTargets {
258    /// Creates routing metadata addressing the given targets.
259    pub const fn new(targets: RoutingTargetSet) -> Self {
260        Self { targets }
261    }
262
263    /// Returns the addressed targets.
264    pub const fn targets(&self) -> &RoutingTargetSet {
265        &self.targets
266    }
267
268    /// Returns `true` if the payload is addressed to the given target.
269    pub(crate) fn includes(&self, id: RoutingTargetId) -> bool {
270        self.targets.contains(id)
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use proptest::prelude::*;
277
278    use super::*;
279
280    fn additional(entries: &[(&str, &[&str])]) -> HashMap<String, Vec<String>> {
281        entries
282            .iter()
283            .map(|(url, keys)| (url.to_string(), keys.iter().map(|key| key.to_string()).collect()))
284            .collect()
285    }
286
287    #[test]
288    fn catalog_orders_primary_first_then_additional_targets_deterministically() {
289        let catalog = RoutingTargetCatalog::new(
290            "https://primary.example.com",
291            &additional(&[
292                ("https://b.example.com", &["key-1", "", "key-1", "key-2"]),
293                ("https://a.example.com", &["key-3"]),
294            ]),
295        )
296        .expect("catalog should build");
297
298        let targets = catalog
299            .targets()
300            .iter()
301            .map(|target| (target.kind(), target.configured_endpoint(), target.api_key_index))
302            .collect::<Vec<_>>();
303        assert_eq!(
304            targets,
305            [
306                (RoutingTargetKind::Primary, "https://primary.example.com", None),
307                (RoutingTargetKind::Additional, "https://a.example.com", Some(0)),
308                (RoutingTargetKind::Additional, "https://b.example.com", Some(0)),
309                (RoutingTargetKind::Additional, "https://b.example.com", Some(3)),
310            ]
311        );
312        assert_eq!(catalog.primary(), RoutingTargetId::from_index(0));
313        assert_eq!(
314            catalog.find_additional("https://b.example.com", 3),
315            Some(RoutingTargetId::from_index(3))
316        );
317        assert_eq!(catalog.find_additional("https://b.example.com", 1), None);
318    }
319
320    #[test]
321    fn catalog_rejects_an_invalid_additional_endpoint() {
322        let error = RoutingTargetCatalog::new("https://primary.example.com", &additional(&[("http://[::1", &["key"])]))
323            .expect_err("an invalid URL should be rejected");
324        assert!(error.to_string().contains("additional destination endpoints"));
325    }
326
327    #[test]
328    fn select_returns_matching_targets() {
329        let catalog = RoutingTargetCatalog::new(
330            "https://primary.example.com",
331            &additional(&[("https://a.example.com", &["key-1", "key-2"])]),
332        )
333        .unwrap();
334
335        let selected = catalog.select(|target| target.configured_endpoint() == "https://a.example.com");
336        assert_eq!(
337            selected.iter().collect::<Vec<_>>(),
338            [RoutingTargetId::from_index(1), RoutingTargetId::from_index(2)]
339        );
340        assert!(catalog.select(|_| false).is_empty());
341    }
342
343    #[test]
344    fn target_sets_compare_equal_regardless_of_growth() {
345        let mut grown = RoutingTargetSet::new();
346        grown.union_with(&[RoutingTargetId::from_index(100)].into_iter().collect());
347        let mut direct = RoutingTargetSet::new();
348        direct.insert(RoutingTargetId::from_index(100));
349        assert_eq!(grown, direct);
350        assert_ne!(direct, RoutingTargetSet::new());
351    }
352
353    proptest! {
354        #[test]
355        fn property_test_target_set_matches_a_btreeset(
356            left in proptest::collection::btree_set(0usize..200, 0..20),
357            right in proptest::collection::btree_set(0usize..200, 0..20),
358        ) {
359            let to_set = |ids: &std::collections::BTreeSet<usize>| {
360                ids.iter().copied().map(RoutingTargetId::from_index).collect::<RoutingTargetSet>()
361            };
362            let mut union = to_set(&left);
363            union.union_with(&to_set(&right));
364
365            let expected = left.union(&right).copied().collect::<Vec<_>>();
366            prop_assert_eq!(union.iter().map(RoutingTargetId::index).collect::<Vec<_>>(), expected.clone());
367            prop_assert_eq!(union.len(), expected.len());
368            prop_assert_eq!(union.is_empty(), expected.is_empty());
369            for id in 0..200 {
370                prop_assert_eq!(union.contains(RoutingTargetId::from_index(id)), expected.contains(&id));
371            }
372        }
373    }
374}