metrics_prometheus_client/
collector.rs1use 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#[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 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 if metrics.is_empty() {
64 continue;
65 }
66 let metric_type = metrics[0].1.metric_type();
67 let has_labels = !metrics[0].0.is_empty();
69
70 let mut descriptor_encoder =
72 encoder.encode_descriptor(key_name.as_str(), help, unit, metric_type)?;
73
74 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 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}