mt_debug_counters/
counter.rs

1use lazy_static::lazy_static;
2use parking_lot::Mutex;
3use std::cmp::{max, min};
4use std::collections::HashMap;
5use std::marker::PhantomData;
6use std::sync::atomic::{AtomicI64, Ordering};
7use std::sync::Weak;
8
9#[doc(hidden)]
10#[macro_export]
11macro_rules! declare_counter_i64_impl {
12    ($name:expr, $mode:ty, $reset:expr, $extra:expr) => {
13        $crate::counter::AtomicCounter::<$mode> {
14            __get_counter: || {
15                thread_local! {
16                    static COUNTER: std::sync::Arc<std::sync::atomic::AtomicI64> = {
17                        let arc = std::sync::Arc::new(std::sync::atomic::AtomicI64::new(0));
18                        let mut list = $crate::counter::__COUNTERS_LIST.lock();
19                        let cvec = list.entry($name.to_string()).or_insert((Vec::new(), 0, <$crate::counter::AtomicCounter<$mode> as $crate::counter::__CounterType>::MODE, $reset));
20                        cvec.0.push(std::sync::Arc::downgrade(&arc));
21                        arc
22                    }
23                }
24                use std::ops::Deref;
25                COUNTER.with(|c| {
26                    unsafe {
27                        &*(c.deref() as *const std::sync::atomic::AtomicI64)
28                    }
29                })
30            },
31            __extra: $extra,
32            _phantom: std::marker::PhantomData
33        }
34    }
35}
36
37#[macro_export]
38macro_rules! declare_counter_i64 {
39    ($name:literal, $mode:ty, $reset:expr) => {
40        $crate::declare_counter_i64_impl!($name, $mode, $reset, ())
41    };
42}
43
44pub(crate) static COUNTER_SUFFIX: &str = "$COUNTER_83uRij";
45
46#[macro_export]
47macro_rules! declare_avg_counter_i64 {
48    ($name:literal, $reset:expr) => {
49        $crate::declare_counter_i64_impl!(
50            $name,
51            $crate::counter::AvgMode,
52            $reset,
53            $crate::declare_counter_i64_impl!(
54                concat!($name, "$COUNTER_83uRij"),
55                $crate::counter::SumMode,
56                $reset,
57                ()
58            )
59        )
60    };
61}
62
63lazy_static! {
64    #[doc(hidden)]
65    pub static ref __COUNTERS_LIST: Mutex<HashMap<String, (Vec<Weak<AtomicI64>>, i64, __AcMode, bool)>> = Mutex::new(HashMap::new());
66}
67
68#[doc(hidden)]
69pub trait AtomicCounterMode {
70    type Extra;
71}
72
73pub struct SumMode {}
74impl AtomicCounterMode for SumMode {
75    type Extra = ();
76}
77pub struct MaxMode {}
78impl AtomicCounterMode for MaxMode {
79    type Extra = ();
80}
81pub struct MinMode {}
82impl AtomicCounterMode for MinMode {
83    type Extra = ();
84}
85pub struct AvgMode {}
86impl AtomicCounterMode for AvgMode {
87    type Extra = AtomicCounter<SumMode>;
88}
89
90pub struct AtomicCounter<MODE: AtomicCounterMode> {
91    #[doc(hidden)]
92    pub __get_counter: fn() -> &'static AtomicI64,
93    #[doc(hidden)]
94    pub __extra: MODE::Extra,
95    #[doc(hidden)]
96    pub _phantom: PhantomData<MODE>,
97}
98
99#[doc(hidden)]
100#[derive(Eq, PartialEq)]
101pub enum __AcMode {
102    SUM,
103    MAX,
104    MIN,
105    AVG,
106}
107
108#[doc(hidden)]
109pub trait __CounterType {
110    const MODE: __AcMode;
111}
112
113impl AtomicCounter<SumMode> {
114    #[inline(always)]
115    pub fn inc(&self) {
116        self.inc_by(1);
117    }
118
119    #[inline(always)]
120    pub fn inc_by(&self, value: i64) {
121        (self.__get_counter)().fetch_add(value, Ordering::Relaxed);
122    }
123
124    #[inline(always)]
125    pub fn sub(&self, value: i64) {
126        (self.__get_counter)().fetch_sub(value, Ordering::Relaxed);
127    }
128}
129
130impl AtomicCounter<AvgMode> {
131    #[inline(always)]
132    pub fn add_value(&self, value: i64) {
133        (self.__get_counter)().fetch_add(value, Ordering::Relaxed);
134        self.__extra.inc();
135    }
136}
137
138impl __CounterType for AtomicCounter<SumMode> {
139    const MODE: __AcMode = __AcMode::SUM;
140}
141
142impl AtomicCounter<MaxMode> {
143    #[inline(always)]
144    pub fn max(&self, val: i64) {
145        (self.__get_counter)().fetch_max(val, Ordering::Relaxed);
146    }
147}
148
149impl __CounterType for AtomicCounter<MaxMode> {
150    const MODE: __AcMode = __AcMode::MAX;
151}
152
153impl AtomicCounter<MinMode> {
154    #[inline(always)]
155    pub fn min(&self, val: i64) {
156        (self.__get_counter)().fetch_min(val, Ordering::Relaxed);
157    }
158}
159
160impl __CounterType for AtomicCounter<MinMode> {
161    const MODE: __AcMode = __AcMode::MIN;
162}
163
164impl __CounterType for AtomicCounter<AvgMode> {
165    const MODE: __AcMode = __AcMode::AVG;
166}
167
168pub struct AtomicCounterGuardSum<'a> {
169    value: i64,
170    counter: &'a AtomicCounter<SumMode>,
171}
172
173impl<'a> AtomicCounterGuardSum<'a> {
174    pub fn new(counter: &'a AtomicCounter<SumMode>, value: i64) -> Self {
175        counter.inc_by(value);
176        Self { value, counter }
177    }
178}
179
180impl<'a> Drop for AtomicCounterGuardSum<'a> {
181    fn drop(&mut self) {
182        self.counter.sub(self.value);
183    }
184}
185
186pub fn get_counter_value(name: &str) -> (i64, i64) {
187    let mut counters = __COUNTERS_LIST.lock();
188
189    let (ref mut vec, part_value, mode, reset) = if let Some(val) = counters.get_mut(name) {
190        val
191    } else {
192        return (0, 0);
193    };
194
195    let reset_value = match mode {
196        __AcMode::SUM => 0,
197        __AcMode::MAX => 0,
198        __AcMode::MIN => i64::MAX,
199        __AcMode::AVG => 0,
200    };
201
202    if *reset {
203        *part_value = reset_value;
204    }
205
206    let mut result = *part_value;
207
208    vec.retain(|val| {
209        if val.strong_count() > 0 {
210            if let Some(value) = val.upgrade() {
211                let value = value.swap(reset_value, Ordering::Relaxed);
212
213                match mode {
214                    __AcMode::SUM => {
215                        result += value;
216                    }
217                    __AcMode::MAX => {
218                        result = max(result, value);
219                    }
220                    __AcMode::MIN => {
221                        result = min(result, value);
222                    }
223                    __AcMode::AVG => {
224                        result += value;
225                    }
226                }
227                true
228            } else {
229                false
230            }
231        } else {
232            false
233        }
234    });
235
236    *part_value = result;
237
238    let is_average = __AcMode::AVG == *mode;
239    drop(counters);
240
241    let counter = if is_average {
242        get_counter_value(&(name.to_string() + COUNTER_SUFFIX)).0
243    } else {
244        0
245    };
246
247    (result, counter)
248}
249
250#[cfg(test)]
251mod tests {
252    use super::SumMode;
253    #[test]
254    fn alloc_test() {
255        let _sum_counter = declare_counter_i64!("test_counter", SumMode, false);
256    }
257}
mt_debug_counters/counter.rs

mt_debug_counters/
counter.rs
