arc_metrics/
lib.rs

1use std::{
2    any::Any,
3    borrow::Cow,
4    fmt::Display,
5    ops::Deref,
6    sync::{
7        atomic::{AtomicU64, Ordering},
8        Arc,
9    },
10};
11
12use helpers::RegisterableMetric;
13
14#[derive(Default, Debug)]
15pub struct IntCounter(pub AtomicU64);
16
17#[derive(Default, Debug)]
18pub struct IntGauge(pub AtomicU64);
19
20pub mod helpers;
21
22pub struct ChildMetric<T, C: 'static> {
23    arc: Arc<T>,
24    child: &'static C,
25}
26
27impl<T, C: 'static> Deref for ChildMetric<T, C> {
28    type Target = C;
29
30    fn deref(&self) -> &Self::Target {
31        self.child
32    }
33}
34
35impl<T: 'static, C: 'static> Clone for ChildMetric<T, C> {
36    fn clone(&self) -> Self {
37        Self {
38            arc: self.arc.clone(),
39            child: self.child,
40        }
41    }
42}
43
44impl<T: 'static, C: 'static> ChildMetric<T, C> {
45    pub fn create<F: Fn(&'static T) -> &'static C>(arc: &Arc<T>, get: F) -> Self {
46        let cloned = arc.clone();
47        let item = get(unsafe { std::mem::transmute::<&T, &'static T>(&cloned) });
48        Self {
49            arc: cloned,
50            child: item,
51        }
52    }
53}
54
55impl IntCounter {
56    pub fn owned_inc(&self) {
57        self.owned_inc_by(1);
58    }
59
60    pub fn inc(&self) {
61        self.shared_inc();
62    }
63
64    pub fn inc_by(&self, amount: u64) {
65        self.shared_inc_by(amount);
66    }
67
68    pub fn shared_inc(&self) {
69        self.shared_inc_by(1);
70    }
71
72    pub fn owned_inc_by(&self, amount: u64) {
73        self.0.fetch_add(amount, Ordering::Relaxed);
74    }
75
76    pub fn shared_inc_by(&self, amount: u64) {
77        self.0.fetch_add(amount, Ordering::AcqRel);
78    }
79
80    pub fn load(&self) -> u64 {
81        self.shared_load()
82    }
83
84    pub fn shared_load(&self) -> u64 {
85        self.0.load(Ordering::Acquire)
86    }
87
88    pub fn owned_load(&self) -> u64 {
89        self.0.load(Ordering::Relaxed)
90    }
91}
92
93impl IntGauge {
94    pub fn set(&self, value: u64) {
95        self.0.store(value, Ordering::Relaxed);
96    }
97
98    pub fn owned_dec(&self) {
99        self.owned_dec_by(1);
100    }
101
102    pub fn dec(&self) {
103        self.shared_dec();
104    }
105
106    pub fn shared_dec(&self) {
107        self.shared_dec_by(1);
108    }
109
110    pub fn owned_dec_by(&self, amount: u64) {
111        self.0.fetch_sub(amount, Ordering::Relaxed);
112    }
113
114    pub fn shared_dec_by(&self, amount: u64) {
115        self.0.fetch_sub(amount, Ordering::AcqRel);
116    }
117
118    pub fn inc(&self) {
119        self.shared_inc();
120    }
121
122    pub fn shared_inc(&self) {
123        self.shared_inc_by(1);
124    }
125
126    pub fn owned_inc_by(&self, amount: u64) {
127        self.0.fetch_add(amount, Ordering::Relaxed);
128    }
129
130    pub fn shared_inc_by(&self, amount: u64) {
131        self.0.fetch_add(amount, Ordering::AcqRel);
132    }
133
134    pub fn load(&self) -> u64 {
135        self.shared_load()
136    }
137
138    pub fn shared_load(&self) -> u64 {
139        self.0.load(Ordering::Acquire)
140    }
141
142    pub fn owned_load(&self) -> u64 {
143        self.0.load(Ordering::Relaxed)
144    }
145}
146
147pub struct PromMetricRegistry {
148    /* note: keep reference to Arc to ensure it doesn't drop */
149    metric_holders: Vec<Arc<dyn Any>>,
150    metrics: Vec<RegisteredMetric>,
151    base_attributes: Vec<[Cow<'static, str>; 2]>,
152}
153
154impl Default for PromMetricRegistry {
155    fn default() -> Self {
156        let base_attributes = if let Some(details) = pkg_details::try_get() {
157            vec![
158                [Cow::Borrowed("program"), Cow::Borrowed(details.pkg_name)],
159                [
160                    Cow::Borrowed("pkg_version"),
161                    Cow::Borrowed(details.pkg_version),
162                ],
163            ]
164        } else {
165            Vec::new()
166        };
167
168        PromMetricRegistry {
169            metric_holders: Vec::new(),
170            metrics: Vec::new(),
171            base_attributes,
172        }
173    }
174}
175
176unsafe impl Send for PromMetricRegistry {}
177unsafe impl Sync for PromMetricRegistry {}
178
179struct RegisteredMetric {
180    metric_type: MetricType,
181    name: Cow<'static, str>,
182    value: &'static AtomicU64,
183    attributes: Vec<[Cow<'static, str>; 2]>,
184    skip_zero: bool,
185}
186
187#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
188pub enum MetricType {
189    IntCounter,
190    IntGauge,
191}
192
193impl Display for MetricType {
194    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
195        match self {
196            Self::IntCounter => write!(f, "counter"),
197            Self::IntGauge => write!(f, "gauge"),
198        }
199    }
200}
201
202impl Display for PromMetricRegistry {
203    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
204        let mut last = None;
205
206        for metric in &self.metrics {
207            let matches = if let Some((last, ty)) = &last {
208                last == &metric.name && *ty == metric.metric_type
209            } else {
210                false
211            };
212
213            if metric.skip_zero && metric.value.load(Ordering::Relaxed) == 0 {
214                continue;
215            }
216
217            if !matches {
218                writeln!(f, "# HELP {}", metric.name)?;
219                writeln!(f, "# TYPE {} {}", metric.name, metric.metric_type)?;
220                last = Some((metric.name.clone(), metric.metric_type));
221            }
222
223            write!(f, "{}", metric.name)?;
224            let end = metric.attributes.len();
225            for (i, [key, value]) in metric.attributes.iter().enumerate() {
226                if i == 0 {
227                    write!(f, "{{{}=\"{}\"", key, value)?;
228                    if end == 1 {
229                        write!(f, "}}")?;
230                    }
231                } else if i + 1 == end {
232                    write!(f, ",{}=\"{}\"}}", key, value)?;
233                } else {
234                    write!(f, ",{}=\"{}\"", key, value)?;
235                }
236            }
237
238            writeln!(f, " {}", metric.value.load(Ordering::Relaxed))?;
239        }
240
241        Ok(())
242    }
243}
244
245impl PromMetricRegistry {
246    pub fn new() -> Self {
247        Self::default()
248    }
249
250    pub fn register<M: RegisterableMetric + 'static>(&mut self, metrics: &Arc<M>) {
251        self.register_fn(metrics, |m, reg| {
252            m.register(reg);
253        });
254    }
255
256    pub fn register_fn<'a, T: 'static>(
257        &'a mut self,
258        metrics: &Arc<T>,
259        register: impl FnOnce(&'static T, &mut RegisterAction<'a>),
260    ) {
261        /* allows us to keep static references as we own an Arc copy */
262        self.metric_holders
263            .push(Arc::clone(metrics) as Arc<dyn Any>);
264
265        let mut action = RegisterAction {
266            name_prefix: None,
267            metrics: &mut self.metrics,
268            base_attributes: self.base_attributes.clone(),
269        };
270
271        let metric_ref = unsafe { std::mem::transmute::<&T, &'static T>(metrics) };
272        register(metric_ref, &mut action);
273    }
274}
275
276pub struct RegisterAction<'a> {
277    metrics: &'a mut Vec<RegisteredMetric>,
278    name_prefix: Option<String>,
279    base_attributes: Vec<[Cow<'static, str>; 2]>,
280}
281
282impl RegisterAction<'_> {
283    pub fn child(&mut self) -> RegisterAction<'_> {
284        RegisterAction {
285            metrics: self.metrics,
286            name_prefix: self.name_prefix.clone(),
287            base_attributes: self.base_attributes.clone(),
288        }
289    }
290
291    pub fn name_prefix<S: Into<String>>(&mut self, prefix: S) -> &mut Self {
292        self.name_prefix = Some(prefix.into());
293        self
294    }
295
296    pub fn base_attr<K: Into<Cow<'static, str>>, V: Into<Cow<'static, str>>>(
297        &mut self,
298        key: K,
299        value: V,
300    ) -> &mut Self {
301        let key = key.into();
302        let value = value.into();
303        self.base_attributes.push([key, value]);
304        self
305    }
306
307    pub fn count<N: Into<Cow<'static, str>>>(
308        &mut self,
309        name: N,
310        count: &'static IntCounter,
311    ) -> RegisterHelper<'_> {
312        self.metric(name, &count.0, MetricType::IntCounter)
313    }
314
315    pub fn gauge<N: Into<Cow<'static, str>>>(
316        &mut self,
317        name: N,
318        gauge: &'static IntGauge,
319    ) -> RegisterHelper<'_> {
320        self.metric(name, &gauge.0, MetricType::IntGauge)
321    }
322
323    fn metric<N: Into<Cow<'static, str>>>(
324        &mut self,
325        name: N,
326        value: &'static AtomicU64,
327        metric_type: MetricType,
328    ) -> RegisterHelper<'_> {
329        let mut helper = self.empty();
330        helper.metric(name, value, metric_type);
331        helper
332    }
333
334    pub fn group<N: Into<Cow<'static, str>>>(&mut self, prefix: N) -> RegisterHelper<'_> {
335        self.start(Some(prefix))
336    }
337
338    pub fn empty(&mut self) -> RegisterHelper<'_> {
339        self.start::<String>(None)
340    }
341
342    fn start<N: Into<Cow<'static, str>>>(&mut self, prefix: Option<N>) -> RegisterHelper<'_> {
343        let attributes = self.base_attributes.clone();
344
345        let name_prefix = match (&self.name_prefix, prefix) {
346            (Some(prefix), None) => Some(Cow::Owned(prefix.clone())),
347            (None, Some(prefix)) => Some(prefix.into()),
348            (Some(a), Some(b)) => {
349                let b = b.into();
350                Some(Cow::Owned(format!("{}_{}", a, b)))
351            }
352            (None, None) => None,
353        };
354
355        RegisterHelper {
356            metrics: self.metrics,
357            name_prefix,
358            attributes,
359            registered: Vec::new(),
360        }
361    }
362}
363
364pub struct RegisterHelper<'a> {
365    name_prefix: Option<Cow<'static, str>>,
366    metrics: &'a mut Vec<RegisteredMetric>,
367    attributes: Vec<[Cow<'static, str>; 2]>,
368    registered: Vec<RegisteredMetric>,
369}
370
371impl RegisterHelper<'_> {
372    pub fn attr<K: Into<Cow<'static, str>>, V: Into<Cow<'static, str>>>(
373        &mut self,
374        key: K,
375        value: V,
376    ) -> &mut Self {
377        let key = key.into();
378        let value = value.into();
379        self.attributes.push([key, value]);
380        self
381    }
382
383    pub fn count<N: Into<Cow<'static, str>>>(
384        &mut self,
385        name: N,
386        count: &'static IntCounter,
387    ) -> &mut Self {
388        self.metric(name, &count.0, MetricType::IntCounter)
389    }
390
391    pub fn gauge<N: Into<Cow<'static, str>>>(
392        &mut self,
393        name: N,
394        gauge: &'static IntGauge,
395    ) -> &mut Self {
396        self.metric(name, &gauge.0, MetricType::IntGauge)
397    }
398
399    pub fn metric<N: Into<Cow<'static, str>>>(
400        &mut self,
401        name: N,
402        value: &'static AtomicU64,
403        metric_type: MetricType,
404    ) -> &mut Self {
405        self.metric_opt(name, value, metric_type, false)
406    }
407
408    pub fn metric_opt<N: Into<Cow<'static, str>>>(
409        &mut self,
410        name: N,
411        value: &'static AtomicU64,
412        metric_type: MetricType,
413        skip_zero: bool,
414    ) -> &mut Self {
415        let name = match &self.name_prefix {
416            Some(prefix) => Cow::Owned(format!("{}_{}", prefix, name.into())),
417            None => name.into(),
418        };
419
420        self.registered.push(RegisteredMetric {
421            metric_type,
422            name,
423            value,
424            attributes: Vec::new(),
425            skip_zero,
426        });
427
428        self
429    }
430}
431
432impl Drop for RegisterHelper<'_> {
433    fn drop(&mut self) {
434        for mut reg in self.registered.drain(..) {
435            reg.attributes = self.attributes.clone();
436            self.metrics.push(reg);
437        }
438        self.metrics.sort_by_key(|item| SortKey {
439            name: item.name.clone(),
440            metric: item.metric_type,
441        });
442    }
443}
444
445#[derive(PartialEq, Eq, PartialOrd, Ord)]
446struct SortKey {
447    name: Cow<'static, str>,
448    metric: MetricType,
449}
450
451#[cfg(test)]
452mod test {
453    use std::sync::Arc;
454
455    use crate::{IntCounter, IntGauge, PromMetricRegistry};
456
457    #[derive(Debug, Default)]
458    struct Met {
459        a: IntCounter,
460        b: IntCounter,
461        c: IntGauge,
462    }
463
464    #[test]
465    fn metrics_test() {
466        let met = Arc::new(Met::default());
467        let mut reg = PromMetricRegistry::new();
468        reg.base_attributes.push(["prefix".into(), "set".into()]);
469
470        reg.register_fn(&met, |m, reg| {
471            reg.name_prefix("base_prefix");
472
473            reg.group("prefix")
474                .count("a", &m.a)
475                .metric_opt("b", &m.b.0, crate::MetricType::IntCounter, true)
476                .attr("test", "2");
477
478            reg.gauge("c", &m.c);
479        });
480
481        println!("{}", reg);
482
483        met.b.inc();
484        println!("{}", reg);
485    }
486}
arc_metrics/lib.rs

arc_metrics/
lib.rs
