witchcraft_metrics/
registry.rs

1// Copyright 2019 Palantir Technologies, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14use crate::{
15    Clock, Counter, ExponentiallyDecayingReservoir, Gauge, Histogram, Meter, MetricId, Timer,
16};
17use parking_lot::Mutex;
18use std::collections::hash_map::Entry;
19use std::collections::{hash_map, HashMap};
20use std::sync::Arc;
21
22/// An enum of all metric types.
23#[derive(Clone)]
24pub enum Metric {
25    /// A counter metric.
26    Counter(Arc<Counter>),
27    /// A meter metric.
28    Meter(Arc<Meter>),
29    /// A gauge metric.
30    Gauge(Arc<dyn Gauge>),
31    /// A histogram metric.
32    Histogram(Arc<Histogram>),
33    /// A timer metric.
34    Timer(Arc<Timer>),
35}
36
37/// A collection of metrics.
38///
39/// Many of the registry's methods take a `T: Into<MetricId>` rather than just a [`MetricId`]. This allows you to pass
40/// either a full `MetricId` or just a `&str` for more convenient use:
41///
42/// ```
43/// use witchcraft_metrics::{MetricRegistry, MetricId};
44///
45/// let registry = MetricRegistry::new();
46///
47/// let requests_meter = registry.meter("server.requests");
48/// let yak_shavings = registry.counter(MetricId::new("shavings").with_tag("animal", "yak"));
49/// ```
50pub struct MetricRegistry {
51    metrics: Mutex<Arc<HashMap<Arc<MetricId>, Metric>>>,
52    clock: Arc<dyn Clock>,
53}
54
55impl Default for MetricRegistry {
56    fn default() -> Self {
57        MetricRegistry {
58            metrics: Mutex::new(Arc::new(HashMap::new())),
59            clock: crate::SYSTEM_CLOCK.clone(),
60        }
61    }
62}
63
64impl MetricRegistry {
65    /// Creates a new, empty registry.
66    #[inline]
67    pub fn new() -> MetricRegistry {
68        MetricRegistry::default()
69    }
70
71    /// Sets the [`Clock`] used as a time source for new metrics created by the registry.
72    ///
73    /// Defaults to [`SystemClock`](crate::SystemClock).
74    #[inline]
75    pub fn set_clock(&mut self, clock: Arc<dyn Clock>) {
76        self.clock = clock;
77    }
78
79    /// Returns a reference to the [`Clock`] used as a time source for new metrics created by the registry.
80    #[inline]
81    pub fn clock(&self) -> &Arc<dyn Clock> {
82        &self.clock
83    }
84
85    /// Returns the counter with the specified ID, using make_counter to create it if absent.
86    ///
87    /// # Panics
88    ///
89    /// Panics if a metric is registered with the ID that is not a counter.
90    pub fn counter_with<T, F>(&self, id: T, make_counter: F) -> Arc<Counter>
91    where
92        T: Into<MetricId>,
93        F: FnOnce() -> Counter,
94    {
95        match Arc::make_mut(&mut self.metrics.lock()).entry(Arc::new(id.into())) {
96            Entry::Occupied(e) => match e.get() {
97                Metric::Counter(c) => c.clone(),
98                _ => panic!("metric already registered as a non-counter: {:?}", e.key()),
99            },
100            Entry::Vacant(e) => {
101                let counter = Arc::new(make_counter());
102                e.insert(Metric::Counter(counter.clone()));
103                counter
104            }
105        }
106    }
107
108    /// Returns the counter with the specified ID, creating a default instance if absent.
109    ///
110    /// # Panics
111    ///
112    /// Panics if a metric is registered with the ID that is not a counter.
113    pub fn counter<T>(&self, id: T) -> Arc<Counter>
114    where
115        T: Into<MetricId>,
116    {
117        self.counter_with(id, Counter::default)
118    }
119
120    /// Returns the meter with the specified ID, using make_meter to create it if absent.
121    ///
122    /// # Panics
123    ///
124    /// Panics if a metric is registered with the ID that is not a meter.
125    pub fn meter_with<T, F>(&self, id: T, make_meter: F) -> Arc<Meter>
126    where
127        T: Into<MetricId>,
128        F: FnOnce() -> Meter,
129    {
130        match Arc::make_mut(&mut self.metrics.lock()).entry(Arc::new(id.into())) {
131            Entry::Occupied(e) => match e.get() {
132                Metric::Meter(m) => m.clone(),
133                _ => panic!("metric already registered as a non-meter: {:?}", e.key()),
134            },
135            Entry::Vacant(e) => {
136                let meter = Arc::new(make_meter());
137                e.insert(Metric::Meter(meter.clone()));
138                meter
139            }
140        }
141    }
142
143    /// Returns the meter with the specified ID, creating a default instance if absent.
144    ///
145    /// # Panics
146    ///
147    /// Panics if a metric is registered with the ID that is not a meter.
148    pub fn meter<T>(&self, id: T) -> Arc<Meter>
149    where
150        T: Into<MetricId>,
151    {
152        self.meter_with(id, || Meter::new_with(self.clock.clone()))
153    }
154
155    /// Returns the gauge with the specified ID, using make_gauge to register a new one if absent.
156    ///
157    /// # Panics
158    ///
159    /// Panics if a metric is registered with the ID that is not a gauge.
160    pub fn gauge_with<T, F, G>(&self, id: T, make_gauge: F) -> Arc<dyn Gauge>
161    where
162        T: Into<MetricId>,
163        F: FnOnce() -> G,
164        G: Gauge,
165    {
166        match Arc::make_mut(&mut self.metrics.lock()).entry(Arc::new(id.into())) {
167            Entry::Occupied(e) => match e.get() {
168                Metric::Gauge(m) => m.clone(),
169                _ => panic!("metric already registered as a non-gauge: {:?}", e.key()),
170            },
171            Entry::Vacant(e) => {
172                let gauge = Arc::new(make_gauge());
173                e.insert(Metric::Gauge(gauge.clone()));
174                gauge
175            }
176        }
177    }
178
179    /// Returns the gauge with the specified ID, registering a new one if absent.
180    ///
181    /// # Panics
182    ///
183    /// Panics if a metric is registered with the ID that is not a gauge.
184    pub fn gauge<T, G>(&self, id: T, gauge: G) -> Arc<dyn Gauge>
185    where
186        T: Into<MetricId>,
187        G: Gauge,
188    {
189        self.gauge_with(id, || gauge)
190    }
191
192    /// Adds a gauge to the registry, overwriting the previous metric with that ID if present.
193    pub fn replace_gauge<T, G>(&self, id: T, gauge: G)
194    where
195        T: Into<MetricId>,
196        G: Gauge,
197    {
198        Arc::make_mut(&mut self.metrics.lock())
199            .insert(Arc::new(id.into()), Metric::Gauge(Arc::new(gauge)));
200    }
201
202    /// Returns the histogram with the specified ID, using make_histogram to create it if absent.
203    ///
204    /// # Panics
205    ///
206    /// Panics if a metric is registered with the ID that is not a histogram.
207    pub fn histogram_with<T, F>(&self, id: T, make_histogram: F) -> Arc<Histogram>
208    where
209        T: Into<MetricId>,
210        F: FnOnce() -> Histogram,
211    {
212        match Arc::make_mut(&mut self.metrics.lock()).entry(Arc::new(id.into())) {
213            Entry::Occupied(e) => match e.get() {
214                Metric::Histogram(m) => m.clone(),
215                _ => panic!(
216                    "metric already registered as a non-histogram: {:?}",
217                    e.key()
218                ),
219            },
220            Entry::Vacant(e) => {
221                let histogram = Arc::new(make_histogram());
222                e.insert(Metric::Histogram(histogram.clone()));
223                histogram
224            }
225        }
226    }
227
228    /// Returns the histogram with the specified ID, creating a default instance if absent.
229    ///
230    /// # Panics
231    ///
232    /// Panics if a metric is registered with the ID that is not a histogram.
233    pub fn histogram<T>(&self, id: T) -> Arc<Histogram>
234    where
235        T: Into<MetricId>,
236    {
237        self.histogram_with(id, || {
238            Histogram::new(ExponentiallyDecayingReservoir::new_with(self.clock.clone()))
239        })
240    }
241
242    /// Returns the timer with the specified ID, using make_timer to create it if absent.
243    ///
244    /// # Panics
245    ///
246    /// Panics if a metric is registered with the ID that is not a timer.
247    pub fn timer_with<T, F>(&self, id: T, make_timer: F) -> Arc<Timer>
248    where
249        T: Into<MetricId>,
250        F: FnOnce() -> Timer,
251    {
252        match Arc::make_mut(&mut self.metrics.lock()).entry(Arc::new(id.into())) {
253            Entry::Occupied(e) => match e.get() {
254                Metric::Timer(m) => m.clone(),
255                _ => panic!("metric already registered as a non-timer: {:?}", e.key()),
256            },
257            Entry::Vacant(e) => {
258                let timer = Arc::new(make_timer());
259                e.insert(Metric::Timer(timer.clone()));
260                timer
261            }
262        }
263    }
264
265    /// Returns the timer with the specified ID, creating a default instance if absent.
266    ///
267    /// # Panics
268    ///
269    /// Panics if a metric is registered with the ID that is not a timer.
270    pub fn timer<T>(&self, id: T) -> Arc<Timer>
271    where
272        T: Into<MetricId>,
273    {
274        self.timer_with(id, || {
275            Timer::new_with(
276                ExponentiallyDecayingReservoir::new_with(self.clock.clone()),
277                self.clock.clone(),
278            )
279        })
280    }
281
282    /// Removes a metric from the registry, returning it if present.
283    pub fn remove<T>(&self, id: T) -> Option<Metric>
284    where
285        T: Into<MetricId>,
286    {
287        Arc::make_mut(&mut self.metrics.lock()).remove(&id.into())
288    }
289
290    /// Returns a snapshot of the metrics in the registry.
291    ///
292    /// Modifications to the registry after this method is called will not affect the state of the returned `Metrics`.
293    pub fn metrics(&self) -> Metrics {
294        Metrics(self.metrics.lock().clone())
295    }
296}
297
298/// A snapshot of the metrics in a registry.
299pub struct Metrics(Arc<HashMap<Arc<MetricId>, Metric>>);
300
301impl Metrics {
302    /// Returns an iterator over the metrics.
303    pub fn iter(&self) -> MetricsIter<'_> {
304        MetricsIter(self.0.iter())
305    }
306}
307
308impl<'a> IntoIterator for &'a Metrics {
309    type Item = (&'a MetricId, &'a Metric);
310    type IntoIter = MetricsIter<'a>;
311
312    fn into_iter(self) -> MetricsIter<'a> {
313        self.iter()
314    }
315}
316
317/// An iterator over metrics and their IDs.
318pub struct MetricsIter<'a>(hash_map::Iter<'a, Arc<MetricId>, Metric>);
319
320impl<'a> Iterator for MetricsIter<'a> {
321    type Item = (&'a MetricId, &'a Metric);
322
323    #[inline]
324    fn next(&mut self) -> Option<(&'a MetricId, &'a Metric)> {
325        self.0.next().map(|(k, v)| (&**k, v))
326    }
327
328    #[inline]
329    fn size_hint(&self) -> (usize, Option<usize>) {
330        self.0.size_hint()
331    }
332}
333
334impl<'a> ExactSizeIterator for MetricsIter<'a> {}
335
336#[cfg(test)]
337mod test {
338    use crate::{MetricId, MetricRegistry};
339    use serde_value::Value;
340    use std::time::Duration;
341
342    #[test]
343    fn first_metric_wins() {
344        let registry = MetricRegistry::new();
345
346        let a = registry.counter("counter");
347        let b = registry.counter("counter");
348        a.add(1);
349        assert_eq!(b.count(), 1);
350
351        registry.gauge("gauge", || 1);
352        let b = registry.gauge("gauge", || 2);
353        assert_eq!(b.value(), Value::I32(1));
354
355        let a = registry.histogram("histogram");
356        let b = registry.histogram("histogram");
357        a.update(0);
358        assert_eq!(b.count(), 1);
359
360        let a = registry.meter("meter");
361        let b = registry.meter("meter");
362        a.mark(1);
363        assert_eq!(b.count(), 1);
364
365        let a = registry.timer("timer");
366        let b = registry.timer("timer");
367        a.update(Duration::from_secs(0));
368        assert_eq!(b.count(), 1);
369    }
370
371    #[test]
372    fn metrics_returns_snapshot() {
373        let registry = MetricRegistry::new();
374
375        registry.counter("counter");
376
377        let metrics = registry.metrics();
378
379        registry.timer("timer");
380
381        let metrics = metrics.iter().collect::<Vec<_>>();
382        assert_eq!(metrics.len(), 1);
383        assert_eq!(metrics[0].0, &MetricId::new("counter"));
384    }
385
386    #[test]
387    fn tagged_distinct_from_untagged() {
388        let registry = MetricRegistry::new();
389
390        let a = registry.counter("counter");
391        let b = registry.counter(MetricId::new("counter").with_tag("foo", "bar"));
392        a.inc();
393        assert_eq!(b.count(), 0);
394    }
395}
witchcraft_metrics/registry.rs

witchcraft_metrics/
registry.rs
