Skip to main content

laminar_connectors/
prom.rs

1//! Shared Prometheus counter/gauge construction for per-connector
2//! metric structs. Replaces the per-file `let local; let reg = …`
3//! boilerplate and `reg_c!` macros that used to be duplicated in
4//! eleven places.
5
6use prometheus::{IntCounter, IntCounterVec, IntGauge, IntGaugeVec, Opts, Registry};
7
8/// Borrowed registry handle used during metric-struct construction.
9pub struct RegHandle<'a> {
10    registry: &'a Registry,
11    _local: Option<Registry>,
12}
13
14impl RegHandle<'_> {
15    /// Borrow the inner registry for non-helper registrations (e.g. `Histogram`).
16    #[must_use]
17    pub fn registry(&self) -> &Registry {
18        self.registry
19    }
20
21    /// Register an `IntCounter`.
22    ///
23    /// # Panics
24    ///
25    /// Panics if `name`/`help` is not a valid Prometheus identifier.
26    #[must_use]
27    pub fn counter(&self, name: &str, help: &str) -> IntCounter {
28        let c =
29            IntCounter::new(name, help).unwrap_or_else(|e| panic!("invalid counter '{name}': {e}"));
30        self.registry.register(Box::new(c.clone())).ok();
31        c
32    }
33
34    /// Register an `IntGauge`.
35    ///
36    /// # Panics
37    ///
38    /// Panics if `name`/`help` is not a valid Prometheus identifier.
39    #[must_use]
40    pub fn gauge(&self, name: &str, help: &str) -> IntGauge {
41        let g = IntGauge::new(name, help).unwrap_or_else(|e| panic!("invalid gauge '{name}': {e}"));
42        self.registry.register(Box::new(g.clone())).ok();
43        g
44    }
45
46    /// Register an `IntCounterVec`.
47    ///
48    /// # Panics
49    ///
50    /// Panics if `name`/`help`/`labels` are not valid Prometheus identifiers.
51    #[must_use]
52    pub fn counter_vec(&self, name: &str, help: &str, labels: &[&str]) -> IntCounterVec {
53        let v = IntCounterVec::new(Opts::new(name, help), labels)
54            .unwrap_or_else(|e| panic!("invalid counter_vec '{name}': {e}"));
55        self.registry.register(Box::new(v.clone())).ok();
56        v
57    }
58
59    /// Register an `IntGaugeVec`.
60    ///
61    /// # Panics
62    ///
63    /// Panics if `name`/`help`/`labels` are not valid Prometheus identifiers.
64    #[must_use]
65    pub fn gauge_vec(&self, name: &str, help: &str, labels: &[&str]) -> IntGaugeVec {
66        let v = IntGaugeVec::new(Opts::new(name, help), labels)
67            .unwrap_or_else(|e| panic!("invalid gauge_vec '{name}': {e}"));
68        self.registry.register(Box::new(v.clone())).ok();
69        v
70    }
71}
72
73/// `RegHandle` over `registry`, falling back to a fresh registry parked
74/// in `local` (which must outlive construction). Pattern:
75/// `let mut local = None; let reg = reg_or_local(registry, &mut local);`
76#[must_use]
77#[allow(clippy::missing_panics_doc)] // `expect` follows an unconditional `*local = Some(..)`
78pub fn reg_or_local<'a>(
79    registry: Option<&'a Registry>,
80    local: &'a mut Option<Registry>,
81) -> RegHandle<'a> {
82    if let Some(r) = registry {
83        RegHandle {
84            registry: r,
85            _local: None,
86        }
87    } else {
88        *local = Some(Registry::new());
89        RegHandle {
90            registry: local.as_ref().expect("just initialized"),
91            _local: None,
92        }
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    #[test]
101    fn counter_registers_and_counts() {
102        let r = Registry::new();
103        let h = RegHandle {
104            registry: &r,
105            _local: None,
106        };
107        let c = h.counter("test_total", "Test events");
108        c.inc();
109        c.inc();
110        assert_eq!(c.get(), 2);
111        // Registry should also see the counter.
112        let mfs = r.gather();
113        assert!(mfs.iter().any(|m| m.name() == "test_total"));
114    }
115
116    #[test]
117    fn reg_or_local_uses_provided() {
118        let provided = Registry::new();
119        let mut local = None;
120        let h = reg_or_local(Some(&provided), &mut local);
121        let c = h.counter("provided_total", "Counter on provided registry");
122        c.inc();
123        assert!(local.is_none());
124        assert!(provided
125            .gather()
126            .iter()
127            .any(|m| m.name() == "provided_total"));
128    }
129
130    #[test]
131    fn reg_or_local_falls_back_to_local() {
132        let mut local = None;
133        let h = reg_or_local(None, &mut local);
134        let c = h.counter("local_total", "Counter on local registry");
135        c.inc();
136        assert!(local.is_some());
137    }
138}