clickhouse_connection_pool/
metrics.rs

1use actix_web_prom::{PrometheusMetrics, PrometheusMetricsBuilder};
2use prometheus::{CounterVec, GaugeVec, IntCounter, IntCounterVec, IntGaugeVec, Opts};
3use std::{
4    collections::HashMap,
5    sync::{Arc, RwLock},
6};
7
8#[derive(Debug, thiserror::Error)]
9pub enum Error {
10    #[error(transparent)]
11    Prometheus(#[from] prometheus::Error),
12}
13
14#[derive(Debug, Default, Clone)]
15pub enum Kind {
16    #[default]
17    Default,
18    GaugeVec,
19    CounterVec,
20    IntGaugeVec,
21    IntCounterVec,
22    IntCounter,
23}
24
25#[derive(Debug, Default, Clone)]
26pub struct MetricConfig<'a> {
27    pub kind: Kind,
28    pub name: &'a str,
29    pub help: &'a str,
30    pub label_names: &'a [&'a str],
31}
32
33pub type SharedRegistrar = Arc<Registrar>;
34
35/// An abstracted metrics registrar for Prometheus.
36#[derive(Clone)]
37pub struct Registrar {
38    prometheus: Arc<PrometheusMetrics>,
39    int_counters_vecs: Arc<RwLock<HashMap<String, IntCounterVec>>>,
40    int_counters: Arc<RwLock<HashMap<String, IntCounter>>>,
41    int_gauges_vecs: Arc<RwLock<HashMap<String, IntGaugeVec>>>,
42    counters_vecs: Arc<RwLock<HashMap<String, CounterVec>>>,
43    gauges_vecs: Arc<RwLock<HashMap<String, GaugeVec>>>,
44}
45
46impl std::fmt::Debug for Registrar {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        f.debug_struct("Registrar")
49            .field("prometheus", &format!("{:?}", self.prometheus.registry))
50            .field(
51                "int_counters_vecs",
52                &format!("{:?}", self.int_counters_vecs),
53            )
54            .field("int_gauges_vecs", &format!("{:?}", self.int_gauges_vecs))
55            .field("counters_vecs", &format!("{:?}", self.counters_vecs))
56            .field("gauges_vecs", &format!("{:?}", self.gauges_vecs))
57            .finish()
58    }
59}
60
61impl Default for Registrar {
62    fn default() -> Self {
63        Self {
64            prometheus: Arc::new(PrometheusMetricsBuilder::new("default").build().unwrap()),
65            int_counters_vecs: Default::default(),
66            int_counters: Default::default(),
67            int_gauges_vecs: Default::default(),
68            counters_vecs: Default::default(),
69            gauges_vecs: Default::default(),
70        }
71    }
72}
73
74pub trait Registry {
75    fn with_metric_configs<'a>(&self, metrics: &'a [MetricConfig<'a>]) -> Result<(), Error>;
76    fn with_metric_config<'a>(&self, metric: &'a MetricConfig<'a>) -> Result<(), Error>;
77}
78
79impl Registry for Registrar {
80    fn with_metric_config<'a>(&self, metric: &'a MetricConfig<'a>) -> Result<(), Error> {
81        log::info!(
82            "Attempting to register metric with name {:?} and labels {:?}",
83            metric.name,
84            metric.label_names
85        );
86        match &metric.kind {
87            Kind::Default => panic!("Registrar metric kind is `default`, so panicking."),
88            Kind::GaugeVec => {
89                let gauge =
90                    GaugeVec::new(Opts::new(metric.name, metric.help), metric.label_names).unwrap();
91                match self.prometheus.registry.register(Box::new(gauge.clone())) {
92                    Ok(()) => {
93                        self.gauges_vecs
94                            .write()
95                            .unwrap()
96                            .insert(metric.name.to_string(), gauge);
97                        Ok(())
98                    }
99                    Err(e) => {
100                        if let prometheus::Error::AlreadyReg = e {
101                            log::info!("Metric {:?} is already registered.", metric.name);
102                            Ok(())
103                        } else {
104                            log::error!("Failed to register metric {:?}. {:?}", metric.name, e);
105                            Err(Error::Prometheus(e))
106                        }
107                    }
108                }
109            }
110            Kind::CounterVec => {
111                let counter =
112                    CounterVec::new(Opts::new(metric.name, metric.help), metric.label_names)
113                        .unwrap();
114                match self.prometheus.registry.register(Box::new(counter.clone())) {
115                    Ok(()) => {
116                        self.counters_vecs
117                            .write()
118                            .unwrap()
119                            .insert(metric.name.to_string(), counter);
120                        Ok(())
121                    }
122                    Err(e) => {
123                        if let prometheus::Error::AlreadyReg = e {
124                            log::info!("Metric {:?} is already registered.", metric.name);
125                            Ok(())
126                        } else {
127                            log::error!("Failed to register metric {:?}. {:?}", metric.name, e);
128                            Err(Error::Prometheus(e))
129                        }
130                    }
131                }
132            }
133            Kind::IntGaugeVec => {
134                let gauge =
135                    IntGaugeVec::new(Opts::new(metric.name, metric.help), metric.label_names)
136                        .unwrap();
137                match self.prometheus.registry.register(Box::new(gauge.clone())) {
138                    Ok(()) => {
139                        self.int_gauges_vecs
140                            .write()
141                            .unwrap()
142                            .insert(metric.name.to_string(), gauge);
143                        Ok(())
144                    }
145                    Err(e) => {
146                        if let prometheus::Error::AlreadyReg = e {
147                            log::info!("Metric {:?} is already registered.", metric.name);
148                            Ok(())
149                        } else {
150                            log::error!("Failed to register metric {:?}. {:?}", metric.name, e);
151                            Err(Error::Prometheus(e))
152                        }
153                    }
154                }
155            }
156            Kind::IntCounterVec => {
157                let counter =
158                    IntCounterVec::new(Opts::new(metric.name, metric.help), metric.label_names)
159                        .unwrap();
160                match self.prometheus.registry.register(Box::new(counter.clone())) {
161                    Ok(()) => {
162                        self.int_counters_vecs
163                            .write()
164                            .unwrap()
165                            .insert(metric.name.to_string(), counter);
166                        Ok(())
167                    }
168                    Err(e) => {
169                        if let prometheus::Error::AlreadyReg = e {
170                            log::info!("Metric {:?} is already registered.", metric.name);
171                            Ok(())
172                        } else {
173                            log::error!("Failed to register metric {:?}. {:?}", metric.name, e);
174                            Err(Error::Prometheus(e))
175                        }
176                    }
177                }
178            }
179            Kind::IntCounter => {
180                let counter = IntCounter::new(metric.name, metric.help).unwrap();
181                match self.prometheus.registry.register(Box::new(counter.clone())) {
182                    Ok(()) => {
183                        self.int_counters
184                            .write()
185                            .unwrap()
186                            .insert(metric.name.to_string(), counter);
187                        Ok(())
188                    }
189                    Err(e) => {
190                        if let prometheus::Error::AlreadyReg = e {
191                            log::info!("Metric {:?} is already registered.", metric.name);
192                            Ok(())
193                        } else {
194                            log::error!("Failed to register metric {:?}. {:?}", metric.name, e);
195                            Err(Error::Prometheus(e))
196                        }
197                    }
198                }
199            }
200        }
201    }
202    fn with_metric_configs<'a>(&self, metrics: &'a [MetricConfig<'a>]) -> Result<(), Error> {
203        for metric in metrics {
204            match self.with_metric_config(metric) {
205                Ok(()) => (),
206                Err(e) => return Err(e),
207            }
208        }
209        Ok(())
210    }
211}
212
213impl Registrar {
214    pub fn new(prometheus: Arc<PrometheusMetrics>) -> Self {
215        Self {
216            prometheus,
217            ..Default::default()
218        }
219    }
220    // Int
221    pub fn inc_int_counter(&self, key: &str) {
222        let mut counters = self.int_counters.write().unwrap();
223        let counter = match counters.get_mut(key) {
224            Some(r) => r,
225            None => return,
226        };
227
228        counter.inc()
229    }
230    pub fn inc_int_counter_vec_mut(&self, key: &str, labels: &[&str]) {
231        let mut counters = self.int_counters_vecs.write().unwrap();
232        let counter = match counters.get_mut(key) {
233            Some(r) => r,
234            None => return,
235        };
236
237        counter.with_label_values(labels).inc()
238    }
239    pub fn inc_by_int_counter_vec_mut(&self, key: &str, labels: &[&str], value: u64) {
240        let mut counters = self.int_counters_vecs.write().unwrap();
241        let counter = match counters.get_mut(key) {
242            Some(r) => r,
243            None => return,
244        };
245
246        counter.with_label_values(labels).inc_by(value)
247    }
248
249    pub fn set_int_gauge_vec_mut(&self, key: &str, labels: &[&str], value: i64) {
250        let mut gauges = self.int_gauges_vecs.write().unwrap();
251        let gauge = match gauges.get_mut(key) {
252            Some(r) => r,
253            None => return,
254        };
255
256        gauge.with_label_values(labels).set(value)
257    }
258    // Floating
259    pub fn inc_counter_vec_mut(&self, key: &str, labels: &[&str]) {
260        let mut counters = self.counters_vecs.write().unwrap();
261        let counter = match counters.get_mut(key) {
262            Some(r) => r,
263            None => return,
264        };
265
266        counter.with_label_values(labels).inc()
267    }
268    pub fn set_gauge_vec_mut(&self, key: &str, labels: &[&str], value: f64) {
269        let mut gauges = self.gauges_vecs.write().unwrap();
270        let gauge = match gauges.get_mut(key) {
271            Some(r) => r,
272            None => return,
273        };
274
275        gauge.with_label_values(labels).set(value)
276    }
277}
clickhouse_connection_pool/metrics.rs

clickhouse_connection_pool/
metrics.rs
