kumo_prometheus/
counter.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2use std::sync::{Arc, Weak};
3
4/// This trait enables having prunable and non-pruning value types
5/// in the CounterRegistry.
6pub trait AtomicCounterEntry: Send + Sync {
7    /// resolve this entry to an AtomicCounter instance
8    fn resolve(&self) -> Option<AtomicCounter>;
9    /// Given a new strong AtomicCounter reference, return Self
10    /// suitable for storing in the counter registry
11    fn make_storable(strong: &AtomicCounter) -> Self;
12    /// Indicate whether this type of value requires pruning the
13    /// containing counter registry
14    fn needs_pruning() -> bool;
15}
16
17#[derive(Clone)]
18pub struct WeakAtomicCounter(Weak<AtomicUsize>);
19
20impl WeakAtomicCounter {
21    pub fn upgrade(&self) -> Option<AtomicCounter> {
22        Some(AtomicCounter(self.0.upgrade()?))
23    }
24}
25
26/// WeakAtomicCounter stores values as weak references and thus
27/// requires pruning.
28impl AtomicCounterEntry for WeakAtomicCounter {
29    fn resolve(&self) -> Option<AtomicCounter> {
30        self.upgrade()
31    }
32
33    fn make_storable(strong: &AtomicCounter) -> WeakAtomicCounter {
34        strong.weak()
35    }
36
37    fn needs_pruning() -> bool {
38        true
39    }
40}
41
42#[derive(Clone)]
43pub struct AtomicCounter(Arc<AtomicUsize>);
44
45/// AtomicCounter is a direct store of the underlying counter value.
46/// No pruning is required for this type of value.
47impl AtomicCounterEntry for AtomicCounter {
48    fn resolve(&self) -> Option<AtomicCounter> {
49        Some(self.clone())
50    }
51
52    fn make_storable(strong: &AtomicCounter) -> AtomicCounter {
53        strong.clone()
54    }
55
56    fn needs_pruning() -> bool {
57        false
58    }
59}
60
61impl std::fmt::Debug for AtomicCounter {
62    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
63        fmt.debug_struct("AtomicCounter").finish()
64    }
65}
66
67impl AtomicCounter {
68    pub fn new() -> Self {
69        Self(Arc::new(AtomicUsize::new(0)))
70    }
71
72    pub fn weak(&self) -> WeakAtomicCounter {
73        WeakAtomicCounter(Arc::downgrade(&self.0))
74    }
75
76    #[inline]
77    pub fn get(&self) -> usize {
78        self.0.load(Ordering::Relaxed)
79    }
80
81    #[inline]
82    pub fn set(&self, v: usize) {
83        self.0.store(v, Ordering::Relaxed)
84    }
85
86    #[inline]
87    pub fn inc(&self) {
88        self.inc_by(1);
89    }
90
91    #[inline]
92    pub fn inc_by(&self, n: usize) {
93        self.0.fetch_add(n, Ordering::Relaxed);
94    }
95
96    #[inline]
97    pub fn dec(&self) {
98        self.sub(1);
99    }
100
101    #[inline]
102    pub fn sub(&self, n: usize) {
103        self.0.fetch_sub(n, Ordering::Relaxed);
104    }
105}
106
107impl Default for AtomicCounter {
108    fn default() -> Self {
109        Self::new()
110    }
111}