metrics_util/layers/
router.rs

1use std::fmt;
2
3use metrics::{Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit};
4use radix_trie::{Trie, TrieCommon};
5
6use crate::{MetricKind, MetricKindMask};
7
8/// Routes metrics to specific target recorders.
9///
10/// More information on the behavior of the layer can be found in [`RouterBuilder`].
11pub struct Router {
12    default: Box<dyn Recorder + Sync>,
13    global_mask: MetricKindMask,
14    targets: Vec<Box<dyn Recorder + Sync>>,
15    counter_routes: Trie<String, usize>,
16    gauge_routes: Trie<String, usize>,
17    histogram_routes: Trie<String, usize>,
18}
19
20impl fmt::Debug for Router {
21    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
22        f.debug_struct("Router")
23            .field("global_mask", &self.global_mask)
24            .field("targets_len", &self.targets.len())
25            .field("counter_routes", &self.counter_routes)
26            .field("gauge_routes", &self.gauge_routes)
27            .field("histogram_routes", &self.histogram_routes)
28            .finish_non_exhaustive()
29    }
30}
31impl Router {
32    fn route(
33        &self,
34        kind: MetricKind,
35        key: &str,
36        search_routes: &Trie<String, usize>,
37    ) -> &dyn Recorder {
38        // The global mask is essentially a Bloom filter of overridden route types.  If it doesn't
39        // match our metric, we know for a fact there's no route and must use the default recorder.
40        if !self.global_mask.matches(kind) {
41            self.default.as_ref()
42        } else {
43            // SAFETY: We derive the `idx` value that is inserted into our route maps by using the
44            // length of `targets` itself before adding a new target.  Ergo, the index is provably
45            // populated if the `idx` has been stored.
46            search_routes
47                .get_ancestor(key)
48                .map(|st| unsafe { self.targets.get_unchecked(*st.value().unwrap()).as_ref() })
49                .unwrap_or_else(|| self.default.as_ref())
50        }
51    }
52}
53
54impl Recorder for Router {
55    fn describe_counter(&self, key_name: KeyName, unit: Option<Unit>, description: SharedString) {
56        let target = self.route(MetricKind::Counter, key_name.as_str(), &self.counter_routes);
57        target.describe_counter(key_name, unit, description)
58    }
59
60    fn describe_gauge(&self, key_name: KeyName, unit: Option<Unit>, description: SharedString) {
61        let target = self.route(MetricKind::Gauge, key_name.as_str(), &self.gauge_routes);
62        target.describe_gauge(key_name, unit, description)
63    }
64
65    fn describe_histogram(&self, key_name: KeyName, unit: Option<Unit>, description: SharedString) {
66        let target = self.route(MetricKind::Histogram, key_name.as_str(), &self.histogram_routes);
67        target.describe_histogram(key_name, unit, description)
68    }
69
70    fn register_counter(&self, key: &Key, metadata: &Metadata<'_>) -> Counter {
71        let target = self.route(MetricKind::Counter, key.name(), &self.counter_routes);
72        target.register_counter(key, metadata)
73    }
74
75    fn register_gauge(&self, key: &Key, metadata: &Metadata<'_>) -> Gauge {
76        let target = self.route(MetricKind::Gauge, key.name(), &self.gauge_routes);
77        target.register_gauge(key, metadata)
78    }
79
80    fn register_histogram(&self, key: &Key, metadata: &Metadata<'_>) -> Histogram {
81        let target = self.route(MetricKind::Histogram, key.name(), &self.histogram_routes);
82        target.register_histogram(key, metadata)
83    }
84}
85
86/// Routes metrics to specific target recorders.
87///
88/// Routes are defined as a prefix to check against the metric name, and a mask for the metric type.
89/// For example,  a route with the pattern of "foo" would match "foo", "or "foo.submetric", but not
90/// "something.foo". Likewise, a metric mask of "all" would apply this route to counters, gauges,
91/// and histograms, while any specific mask would only apply to the given metric kind.
92///
93/// A default route (recorder) is always present and used in the case that no specific route exists.
94pub struct RouterBuilder {
95    default: Box<dyn Recorder + Sync>,
96    global_mask: MetricKindMask,
97    targets: Vec<Box<dyn Recorder + Sync>>,
98    counter_routes: Trie<String, usize>,
99    gauge_routes: Trie<String, usize>,
100    histogram_routes: Trie<String, usize>,
101}
102
103impl fmt::Debug for RouterBuilder {
104    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
105        f.debug_struct("RouterBuilder")
106            .field("global_mask", &self.global_mask)
107            .field("targets_len", &self.targets.len())
108            .field("counter_routes", &self.counter_routes)
109            .field("gauge_routes", &self.gauge_routes)
110            .field("histogram_routes", &self.histogram_routes)
111            .finish_non_exhaustive()
112    }
113}
114
115impl RouterBuilder {
116    /// Creates a [`RouterBuilder`] from a [`Recorder`].
117    ///
118    /// The given recorder is used as the default route when no other specific route exists.
119    pub fn from_recorder<R>(recorder: R) -> Self
120    where
121        R: Recorder + Sync + 'static,
122    {
123        RouterBuilder {
124            default: Box::new(recorder),
125            global_mask: MetricKindMask::NONE,
126            targets: Vec::new(),
127            counter_routes: Trie::new(),
128            gauge_routes: Trie::new(),
129            histogram_routes: Trie::new(),
130        }
131    }
132
133    /// Adds a route.
134    ///
135    /// `mask` defines which metric kinds will match the given route, and `pattern` is a prefix
136    /// string used to match against metric names.
137    ///
138    /// If a matching route already exists, it will be overwritten.
139    pub fn add_route<P, R>(
140        &mut self,
141        mask: MetricKindMask,
142        pattern: P,
143        recorder: R,
144    ) -> &mut RouterBuilder
145    where
146        P: AsRef<str>,
147        R: Recorder + Sync + 'static,
148    {
149        let target_idx = self.targets.len();
150        self.targets.push(Box::new(recorder));
151
152        self.global_mask = self.global_mask | mask;
153
154        match mask {
155            MetricKindMask::ALL => {
156                let _ = self.counter_routes.insert(pattern.as_ref().to_string(), target_idx);
157                let _ = self.gauge_routes.insert(pattern.as_ref().to_string(), target_idx);
158                let _ = self.histogram_routes.insert(pattern.as_ref().to_string(), target_idx);
159            }
160            MetricKindMask::COUNTER => {
161                let _ = self.counter_routes.insert(pattern.as_ref().to_string(), target_idx);
162            }
163            MetricKindMask::GAUGE => {
164                let _ = self.gauge_routes.insert(pattern.as_ref().to_string(), target_idx);
165            }
166            MetricKindMask::HISTOGRAM => {
167                let _ = self.histogram_routes.insert(pattern.as_ref().to_string(), target_idx);
168            }
169            _ => panic!("cannot add route for unknown or empty metric kind mask"),
170        };
171        self
172    }
173
174    /// Builds the configured [`Router`].
175    pub fn build(self) -> Router {
176        Router {
177            default: self.default,
178            global_mask: self.global_mask,
179            targets: self.targets,
180            counter_routes: self.counter_routes,
181            gauge_routes: self.gauge_routes,
182            histogram_routes: self.histogram_routes,
183        }
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use mockall::{
190        mock,
191        predicate::{always, eq},
192        Sequence,
193    };
194    use std::{borrow::Cow, sync::Arc};
195
196    use super::RouterBuilder;
197    use crate::MetricKindMask;
198    use metrics::{
199        Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit,
200    };
201
202    mock! {
203        #[derive(Debug)]
204        pub TestRecorder {
205        }
206
207        impl Recorder for TestRecorder {
208            fn describe_counter(&self, key_name: KeyName, unit: Option<Unit>, description: SharedString);
209            fn describe_gauge(&self, key_name: KeyName, unit: Option<Unit>, description: SharedString);
210            fn describe_histogram(&self, key_name: KeyName, unit: Option<Unit>, description: SharedString);
211            fn register_counter<'a>(&'a self, key: &'a Key, metadata: &'a Metadata<'a>) -> Counter;
212            fn register_gauge<'a>(&'a self, key: &'a Key, metadata: &'a Metadata<'a>) -> Gauge;
213            fn register_histogram<'a>(&'a self, key: &'a Key, metadata: &'a Metadata<'a>) -> Histogram;
214        }
215    }
216
217    #[test]
218    fn sync() {
219        #[allow(dead_code)]
220        fn assert_sync_recorder<T: Recorder + Sync>(_t: &T) {}
221
222        let recorder = RouterBuilder::from_recorder(MockTestRecorder::new()).build();
223        assert_sync_recorder(&recorder);
224    }
225
226    #[test]
227    fn test_construction() {
228        let _ = RouterBuilder::from_recorder(MockTestRecorder::new()).build();
229
230        let mut builder = RouterBuilder::from_recorder(MockTestRecorder::new());
231        // ensure that &str, String, and Cow<str> are all are accepted by the builder
232        builder
233            .add_route(MetricKindMask::COUNTER, "foo", MockTestRecorder::new())
234            .add_route(MetricKindMask::GAUGE, String::from("bar"), MockTestRecorder::new())
235            .add_route(MetricKindMask::HISTOGRAM, Cow::Borrowed("baz"), MockTestRecorder::new())
236            .add_route(MetricKindMask::ALL, "quux", MockTestRecorder::new());
237        let _ = builder.build();
238    }
239
240    #[test]
241    #[should_panic]
242    fn test_bad_construction() {
243        let mut builder = RouterBuilder::from_recorder(MockTestRecorder::new());
244        builder.add_route(MetricKindMask::NONE, "foo", MockTestRecorder::new());
245        let _ = builder.build();
246    }
247
248    #[test]
249    fn test_basic_functionality() {
250        let default_counter: Key = "counter_default.foo".into();
251        let override_counter: Key = "counter_override.foo".into();
252        let all_override: Key = "all_override.foo".into();
253
254        let mut default_mock = MockTestRecorder::new();
255        let mut counter_mock = MockTestRecorder::new();
256        let mut all_mock = MockTestRecorder::new();
257
258        let mut seq = Sequence::new();
259
260        static METADATA: metrics::Metadata =
261            metrics::Metadata::new(module_path!(), metrics::Level::INFO, Some(module_path!()));
262
263        default_mock
264            .expect_register_counter()
265            .times(1)
266            .in_sequence(&mut seq)
267            .with(eq(default_counter.clone()), always())
268            .returning(|_, _| Counter::noop());
269
270        counter_mock
271            .expect_register_counter()
272            .times(1)
273            .in_sequence(&mut seq)
274            .with(eq(override_counter.clone()), always())
275            .returning(|_, _| Counter::noop());
276
277        all_mock
278            .expect_register_counter()
279            .times(1)
280            .in_sequence(&mut seq)
281            .with(eq(all_override.clone()), always())
282            .returning(|_, _| Counter::noop());
283
284        all_mock
285            .expect_register_histogram()
286            .times(1)
287            .in_sequence(&mut seq)
288            .with(eq(all_override.clone()), always())
289            .returning(|_, _| Histogram::noop());
290
291        let mut builder = RouterBuilder::from_recorder(default_mock);
292        builder.add_route(MetricKindMask::COUNTER, "counter_override", counter_mock).add_route(
293            MetricKindMask::ALL,
294            "all_override",
295            all_mock,
296        );
297        let recorder = builder.build();
298
299        let _ = recorder.register_counter(&default_counter, &METADATA);
300        let _ = recorder.register_counter(&override_counter, &METADATA);
301        let _ = recorder.register_counter(&all_override, &METADATA);
302        let _ = recorder.register_histogram(&all_override, &METADATA);
303    }
304
305    #[test]
306    fn test_same_recorder_multiple_routes() {
307        let default_counter: Key = "default".into();
308        let foo_counter: Key = "foo.counter".into();
309        let bar_counter: Key = "bar.counter".into();
310
311        let mut default_mock = MockTestRecorder::new();
312        let mut foo_bar_mock = MockTestRecorder::new();
313
314        let mut seq = Sequence::new();
315
316        static METADATA: metrics::Metadata =
317            metrics::Metadata::new(module_path!(), metrics::Level::INFO, Some(module_path!()));
318
319        foo_bar_mock
320            .expect_register_counter()
321            .times(1)
322            .in_sequence(&mut seq)
323            .with(eq(foo_counter.clone()), always())
324            .returning(|_, _| Counter::noop());
325        foo_bar_mock
326            .expect_register_counter()
327            .times(1)
328            .in_sequence(&mut seq)
329            .with(eq(bar_counter.clone()), always())
330            .returning(|_, _| Counter::noop());
331        default_mock
332            .expect_register_counter()
333            .times(1)
334            .in_sequence(&mut seq)
335            .with(eq(default_counter.clone()), always())
336            .returning(|_, _| Counter::noop());
337
338        let foo_bar_mock = Arc::new(foo_bar_mock);
339
340        let mut builder = RouterBuilder::from_recorder(default_mock);
341        builder.add_route(MetricKindMask::COUNTER, "foo", foo_bar_mock.clone());
342        builder.add_route(MetricKindMask::COUNTER, "bar", foo_bar_mock);
343        let recorder = builder.build();
344
345        let _ = recorder.register_counter(&foo_counter, &METADATA);
346        let _ = recorder.register_counter(&bar_counter, &METADATA);
347        let _ = recorder.register_counter(&default_counter, &METADATA);
348    }
349}