1use metrics::{Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit};
2use radix_trie::{Trie, TrieCommon};
3
4use crate::{MetricKind, MetricKindMask};
5
6pub 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 if !self.global_mask.matches(kind) {
28 self.default.as_ref()
29 } else {
30 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
73pub 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 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 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 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}