saluki_components/common/datadog/
routing.rs1use std::{collections::HashMap, fmt};
14
15use saluki_error::{ErrorContext as _, GenericError};
16use smallvec::SmallVec;
17
18use super::endpoints::resolve_additional_endpoints;
19
20#[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 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#[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
46pub struct RoutingTargetSet {
47 words: SmallVec<[u64; 1]>,
48}
49
50impl RoutingTargetSet {
51 pub fn new() -> Self {
53 Self::default()
54 }
55
56 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 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 pub fn is_empty(&self) -> bool {
73 self.words.is_empty()
74 }
75
76 pub fn len(&self) -> usize {
78 self.words.iter().map(|word| word.count_ones() as usize).sum()
79 }
80
81 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 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#[derive(Clone, Copy, Debug, Eq, PartialEq)]
113pub enum RoutingTargetKind {
114 Primary,
119
120 Additional,
122}
123
124#[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 pub const fn id(&self) -> RoutingTargetId {
136 self.id
137 }
138
139 pub const fn kind(&self) -> RoutingTargetKind {
141 self.kind
142 }
143
144 pub fn configured_endpoint(&self) -> &str {
149 &self.configured_endpoint
150 }
151
152 #[cfg(test)]
156 pub(crate) const fn api_key_index(&self) -> Option<usize> {
157 self.api_key_index
158 }
159}
160
161#[derive(Clone, Debug, Eq, PartialEq)]
166pub struct RoutingTargetCatalog {
167 targets: Vec<RoutingTarget>,
168}
169
170impl RoutingTargetCatalog {
171 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 pub fn primary(&self) -> RoutingTargetId {
213 self.targets[0].id
214 }
215
216 pub fn targets(&self) -> &[RoutingTarget] {
218 &self.targets
219 }
220
221 pub fn get(&self, id: RoutingTargetId) -> Option<&RoutingTarget> {
223 self.targets.get(id.index())
224 }
225
226 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 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#[derive(Clone, Debug, Eq, PartialEq)]
253pub struct MetricsRoutingTargets {
254 targets: RoutingTargetSet,
255}
256
257impl MetricsRoutingTargets {
258 pub const fn new(targets: RoutingTargetSet) -> Self {
260 Self { targets }
261 }
262
263 pub const fn targets(&self) -> &RoutingTargetSet {
265 &self.targets
266 }
267
268 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}