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
8pub 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 if !self.global_mask.matches(kind) {
41 self.default.as_ref()
42 } else {
43 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
86pub 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 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 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 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 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}