metrics_util/layers/
router.rs

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