metrics_prometheus_client/
collector.rs

1use std::{
2    collections::HashMap,
3    fmt,
4    sync::{Arc, RwLock},
5};
6
7use metrics::{Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit};
8use prometheus_client::{
9    collector::Collector,
10    encoding::{DescriptorEncoder, EncodeMetric},
11    registry::Unit as PrometheusUnit,
12};
13
14use crate::{metrics::*, utils::*};
15
16pub type Map<K, V> = Arc<RwLock<HashMap<K, V>>>;
17pub type Label = (String, String);
18
19#[derive(Debug, Default)]
20struct Descriptor {
21    help: String,
22    unit: Option<PrometheusUnit>,
23}
24
25impl Descriptor {
26    fn new(help: String, unit: Option<Unit>) -> Self {
27        Self {
28            help,
29            unit: unit.map(convert_unit_to_prometheus),
30        }
31    }
32}
33
34/// This module provides compatibility with the `metrics` crate.
35/// It registers as a `Collector` with the `prometheus_client` crate
36/// and implements the `Recorder` trait of the `metrics` crate.
37#[derive(Debug, Default)]
38pub struct MetricsCollector {
39    metrics: Map<KeyName, Vec<(Vec<Label>, Metric)>>,
40    descriptors: Map<KeyName, Descriptor>,
41}
42
43impl Clone for MetricsCollector {
44    fn clone(&self) -> Self {
45        Self {
46            metrics: Arc::clone(&self.metrics),
47            descriptors: Arc::clone(&self.descriptors),
48        }
49    }
50}
51
52impl Collector for MetricsCollector {
53    fn encode(&self, mut encoder: DescriptorEncoder) -> Result<(), fmt::Error> {
54        for (key_name, metrics) in self.metrics.read().unwrap().iter() {
55            // Find descriptor for the metric.
56            let descriptors = self.descriptors.read().unwrap();
57            let (help, unit) = descriptors
58                .get(key_name)
59                .map(|d| (d.help.as_str(), d.unit.as_ref()))
60                .unwrap_or_else(|| ("", None));
61
62            // Gather statistics about the metrics.
63            if metrics.is_empty() {
64                continue;
65            }
66            let metric_type = metrics[0].1.metric_type();
67            // If there is more than one entry, this is always true.
68            let has_labels = !metrics[0].0.is_empty();
69
70            // Encode descriptor and metric.
71            let mut descriptor_encoder =
72                encoder.encode_descriptor(key_name.as_str(), help, unit, metric_type)?;
73
74            // Encode metrics for this key.
75            // If labels are present, encode the metric as a family.
76            if has_labels {
77                for (labels, metric) in metrics {
78                    let metric_encoder = descriptor_encoder.encode_family(labels)?;
79                    metric.encode(metric_encoder)?;
80                }
81            } else {
82                let metric = &metrics[0].1;
83                metric.encode(descriptor_encoder)?;
84            }
85        }
86        Ok(())
87    }
88}
89
90impl MetricsCollector {
91    fn register(&self, key: &Key, metric: Metric) {
92        let (key_name, labels) = key.clone().into_parts();
93        let labels = convert_labels_to_prometheus(labels);
94
95        let mut metrics = self.metrics.write().unwrap();
96        let entry = metrics.entry(key_name).or_default();
97
98        // Make sure that all metrics for a key have the same type
99        // and that labels are set on duplicate entries..
100        assert!(
101            entry.is_empty()
102                || (entry[0].1.metric_type().as_str() == metric.metric_type().as_str()
103                    && !entry[0].0.is_empty()
104                    && !labels.is_empty()),
105            "Registering a metric with a different type or missing labels: `{:?}`",
106            key
107        );
108        entry.push((labels, metric));
109    }
110
111    fn describe(&self, key: &KeyName, unit: Option<Unit>, description: SharedString) {
112        assert!(
113            self.descriptors
114                .write()
115                .unwrap()
116                .insert(key.clone(), Descriptor::new(description.into_owned(), unit))
117                .is_none(),
118            "Registering a duplicate metric descriptor: `{:?}`",
119            key
120        );
121    }
122}
123
124impl Recorder for MetricsCollector {
125    fn describe_counter(&self, key: KeyName, unit: Option<Unit>, description: SharedString) {
126        self.describe(&key, unit, description)
127    }
128
129    fn describe_gauge(&self, key: KeyName, unit: Option<Unit>, description: SharedString) {
130        self.describe(&key, unit, description)
131    }
132
133    fn describe_histogram(&self, key: KeyName, unit: Option<Unit>, description: SharedString) {
134        self.describe(&key, unit, description)
135    }
136
137    fn register_counter(&self, key: &Key, _metadata: &Metadata<'_>) -> Counter {
138        let counter = Arc::new(MetricsCounter::default());
139        self.register(key, Metric::Counter(Arc::clone(&counter)));
140        Counter::from_arc(counter)
141    }
142
143    fn register_gauge(&self, key: &Key, _metadata: &Metadata<'_>) -> Gauge {
144        let gauge = Arc::new(MetricsGauge::default());
145        self.register(key, Metric::Gauge(Arc::clone(&gauge)));
146        Gauge::from_arc(gauge)
147    }
148
149    fn register_histogram(&self, key: &Key, _metadata: &Metadata<'_>) -> Histogram {
150        let hist = Arc::new(MetricsHistogram::default());
151        self.register(key, Metric::Histogram(Arc::clone(&hist)));
152        Histogram::from_arc(hist)
153    }
154}