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`. Panics on a malformed name/help.
22    #[must_use]
23    #[allow(clippy::missing_panics_doc)]
24    pub fn counter(&self, name: &str, help: &str) -> IntCounter {
25        let c =
26            IntCounter::new(name, help).unwrap_or_else(|e| panic!("invalid counter '{name}': {e}"));
27        self.registry.register(Box::new(c.clone())).ok();
28        c
29    }
30
31    /// Register an `IntGauge`. Panics on a malformed name/help.
32    #[must_use]
33    #[allow(clippy::missing_panics_doc)]
34    pub fn gauge(&self, name: &str, help: &str) -> IntGauge {
35        let g = IntGauge::new(name, help).unwrap_or_else(|e| panic!("invalid gauge '{name}': {e}"));
36        self.registry.register(Box::new(g.clone())).ok();
37        g
38    }
39
40    /// Register an `IntCounterVec`. Panics on malformed name/help/labels.
41    #[must_use]
42    #[allow(clippy::missing_panics_doc)]
43    pub fn counter_vec(&self, name: &str, help: &str, labels: &[&str]) -> IntCounterVec {
44        let v = IntCounterVec::new(Opts::new(name, help), labels)
45            .unwrap_or_else(|e| panic!("invalid counter_vec '{name}': {e}"));
46        self.registry.register(Box::new(v.clone())).ok();
47        v
48    }
49
50    /// Register an `IntGaugeVec`. Panics on malformed name/help/labels.
51    #[must_use]
52    #[allow(clippy::missing_panics_doc)]
53    pub fn gauge_vec(&self, name: &str, help: &str, labels: &[&str]) -> IntGaugeVec {
54        let v = IntGaugeVec::new(Opts::new(name, help), labels)
55            .unwrap_or_else(|e| panic!("invalid gauge_vec '{name}': {e}"));
56        self.registry.register(Box::new(v.clone())).ok();
57        v
58    }
59}
60
61/// `RegHandle` over `registry`, falling back to a fresh registry parked
62/// in `local` (which must outlive construction). Pattern:
63/// `let mut local = None; let reg = reg_or_local(registry, &mut local);`
64#[must_use]
65#[allow(clippy::missing_panics_doc)]
66pub fn reg_or_local<'a>(
67    registry: Option<&'a Registry>,
68    local: &'a mut Option<Registry>,
69) -> RegHandle<'a> {
70    if let Some(r) = registry {
71        RegHandle {
72            registry: r,
73            _local: None,
74        }
75    } else {
76        *local = Some(Registry::new());
77        RegHandle {
78            registry: local.as_ref().expect("just initialized"),
79            _local: None,
80        }
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    #[test]
89    fn counter_registers_and_counts() {
90        let r = Registry::new();
91        let h = RegHandle {
92            registry: &r,
93            _local: None,
94        };
95        let c = h.counter("test_total", "Test events");
96        c.inc();
97        c.inc();
98        assert_eq!(c.get(), 2);
99        // Registry should also see the counter.
100        let mfs = r.gather();
101        assert!(mfs.iter().any(|m| m.name() == "test_total"));
102    }
103
104    #[test]
105    fn reg_or_local_uses_provided() {
106        let provided = Registry::new();
107        let mut local = None;
108        let h = reg_or_local(Some(&provided), &mut local);
109        let c = h.counter("provided_total", "Counter on provided registry");
110        c.inc();
111        assert!(local.is_none());
112        assert!(provided
113            .gather()
114            .iter()
115            .any(|m| m.name() == "provided_total"));
116    }
117
118    #[test]
119    fn reg_or_local_falls_back_to_local() {
120        let mut local = None;
121        let h = reg_or_local(None, &mut local);
122        let c = h.counter("local_total", "Counter on local registry");
123        c.inc();
124        assert!(local.is_some());
125    }
126}