rmqtt_utils/
rate_counter.rs

1//! Lightweight rate counter for tracking cumulative throughput,
2//! per-second processing speed, current in-flight count, and peak.
3//!
4//! Uses `AtomicU64` / `AtomicI64` fields with Relaxed ordering
5//! throughout. The caller drives the sampling by calling
6//! [`RateCounter::tick`] at a known interval.
7//!
8//! [`inc`](RateCounter::inc) / [`incs`](RateCounter::incs) increment
9//! **both** the cumulative `total` **and** the current in-flight count,
10//! and update the peak (`max`) via atomc `fetch_max`.
11//! [`dec`](RateCounter::dec) / [`decs`](RateCounter::decs) only
12//! decrement `current` — matching the pattern of a
13//! producer-consumer pipeline.
14//!
15//! # Example
16//! ```
17//! use std::time::Duration;
18//! use rmqtt_utils::RateCounter;
19//!
20//! let rc = RateCounter::new();
21//!
22//! // Task arrives: track throughput, in-flight, and peak
23//! rc.incs(42);
24//! assert_eq!(rc.total(), 42);
25//! assert_eq!(rc.current(), 42);
26//! assert_eq!(rc.max(), 42);
27//!
28//! // Higher peak
29//! rc.incs(10);
30//! assert_eq!(rc.max(), 52);
31//!
32//! // Task completes: in-flight decreases, peak unchanged
33//! rc.decs(20);
34//! assert_eq!(rc.current(), 32);
35//! assert_eq!(rc.max(), 52);
36//!
37//! // Compute per-second rate over a 3 s interval
38//! rc.tick(Duration::from_secs(3));
39//! assert!((rc.speed() - 17.333333333333332).abs() < 1e-12);
40//! ```
41
42use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
43use std::sync::Arc;
44use std::time::Duration;
45
46use serde::{Deserialize, Deserializer, Serialize, Serializer};
47use serde_json::json;
48
49use crate::StatsMergeMode;
50
51/// A lock-free, thread-safe rate counter with current in-flight tracking
52/// and peak water mark.
53///
54/// Wraps five atomic values in `Arc` so the counter is [`Clone`] —
55/// clones share the same underlying counters, making it easy to pass
56/// into concurrent tasks (e.g. [`tokio::spawn`]).
57///
58/// The `speed` field stores a true **per-second rate** as an `f64`
59/// (encoded via `f64::to_bits` / `f64::from_bits` in the `AtomicU64`),
60/// normalised by the sampling interval passed to [`tick`](RateCounter::tick).
61///
62/// The `current` field tracks the current in-flight or active count.
63/// [`inc`](RateCounter::inc) / [`incs`](RateCounter::incs) increment **both**
64/// `total` and `current` simultaneously and update `max` to the new peak.
65/// [`dec`](RateCounter::dec) / [`decs`](RateCounter::decs) only decrement
66/// `current` without affecting `max`.
67/// [`tick`](RateCounter::tick) does **not** affect `current` or `max`.
68///
69/// # Serialisation
70///
71/// `RateCounter` serialises as a snapshot of the **current** counter values
72/// (total, speed, current, max). Deserialisation produces a fresh,
73/// **independent** counter initialised to those values — it does NOT share
74/// atomics with the original. This makes it safe for cross-node transfer
75/// (e.g. via gRPC).
76#[derive(Debug)]
77pub struct RateCounter {
78    /// Total cumulative count since construction (or last reset).
79    total: Arc<AtomicU64>,
80    /// Per-second rate, stored as `f64::to_bits` in an `AtomicU64`.
81    speed: Arc<AtomicU64>,
82    /// Snapshot of `total` at the time of the last [`tick`](RateCounter::tick).
83    last_total: Arc<AtomicU64>,
84    /// Current (in-flight / active) count.
85    current: Arc<AtomicI64>,
86    /// Peak (historical maximum) of the current count.
87    max: Arc<AtomicI64>,
88    /// Merge behaviour for [`Stats::add`] aggregation.
89    mode: StatsMergeMode,
90}
91
92impl Clone for RateCounter {
93    fn clone(&self) -> Self {
94        Self {
95            total: Arc::clone(&self.total),
96            speed: Arc::clone(&self.speed),
97            last_total: Arc::clone(&self.last_total),
98            current: Arc::clone(&self.current),
99            max: Arc::clone(&self.max),
100            mode: self.mode.clone(),
101        }
102    }
103}
104
105impl Default for RateCounter {
106    fn default() -> Self {
107        Self::new()
108    }
109}
110
111impl RateCounter {
112    /// Creates a new `RateCounter` with all values initialised to zero
113    /// and [`StatsMergeMode::None`].
114    #[inline]
115    pub fn new() -> Self {
116        Self::new_with_mode(StatsMergeMode::None)
117    }
118
119    /// Creates a new `RateCounter` with a given merge mode.
120    #[inline]
121    pub fn new_with_mode(mode: StatsMergeMode) -> Self {
122        Self {
123            total: Arc::new(AtomicU64::new(0)),
124            speed: Arc::new(AtomicU64::new(0)),
125            last_total: Arc::new(AtomicU64::new(0)),
126            current: Arc::new(AtomicI64::new(0)),
127            max: Arc::new(AtomicI64::new(0)),
128            mode,
129        }
130    }
131
132    /// Creates an independent deep copy (new atomics) with the same values.
133    #[inline]
134    pub fn snapshot(&self) -> Self {
135        Self {
136            total: Arc::new(AtomicU64::new(self.total())),
137            speed: Arc::new(AtomicU64::new(f64::to_bits(self.speed()))),
138            last_total: Arc::new(AtomicU64::new(self.total())),
139            current: Arc::new(AtomicI64::new(self.current())),
140            max: Arc::new(AtomicI64::new(self.max())),
141            mode: self.mode.clone(),
142        }
143    }
144
145    /// Increments total, current, and potentially updates max by 1.
146    #[inline]
147    pub fn inc(&self) {
148        self.total.fetch_add(1, Ordering::Relaxed);
149        let old = self.current.fetch_add(1, Ordering::Relaxed);
150        self.max.fetch_max(old + 1, Ordering::Relaxed);
151    }
152
153    /// Increments total, current, and potentially updates max by `n`.
154    #[inline]
155    pub fn incs(&self, n: u64) {
156        self.total.fetch_add(n, Ordering::Relaxed);
157        let old = self.current.fetch_add(n as i64, Ordering::Relaxed);
158        self.max.fetch_max(old + n as i64, Ordering::Relaxed);
159    }
160
161    /// Returns the cumulative total count.
162    #[inline]
163    pub fn total(&self) -> u64 {
164        self.total.load(Ordering::Relaxed)
165    }
166
167    /// Returns the per-second rate computed by the most recent [`tick`].
168    ///
169    /// Returns `0.0` if [`tick`] has never been called.
170    #[inline]
171    pub fn speed(&self) -> f64 {
172        f64::from_bits(self.speed.load(Ordering::Relaxed))
173    }
174
175    /// Computes the per-second rate: `speed = (total - last_total) / interval`.
176    ///
177    /// Call this periodically at a known `interval` (e.g. every 3 s) to get
178    /// a true per-second throughput, regardless of the sampling duration.
179    #[inline]
180    pub fn tick(&self, interval: Duration) {
181        let curr = self.total.load(Ordering::Relaxed);
182        let prev = self.last_total.swap(curr, Ordering::Relaxed);
183        let delta = curr.wrapping_sub(prev) as f64;
184        let rate = delta / interval.as_secs_f64();
185        self.speed.store(f64::to_bits(rate), Ordering::Relaxed);
186    }
187
188    /// Resets all counters to zero.
189    #[inline]
190    pub fn reset(&self) {
191        self.total.store(0, Ordering::Relaxed);
192        self.speed.store(f64::to_bits(0.0), Ordering::Relaxed);
193        self.last_total.store(0, Ordering::Relaxed);
194        self.current.store(0, Ordering::Relaxed);
195        self.max.store(0, Ordering::Relaxed);
196    }
197
198    // ── current (in-flight) counter ──
199
200    /// Returns the current (in-flight / active) count.
201    #[inline]
202    pub fn current(&self) -> i64 {
203        self.current.load(Ordering::Relaxed)
204    }
205
206    /// Returns the peak (historical maximum) of the current count.
207    #[inline]
208    pub fn max(&self) -> i64 {
209        self.max.load(Ordering::Relaxed)
210    }
211
212    /// Converts the rate counter to JSON format.
213    #[inline]
214    pub fn to_json(&self) -> serde_json::Value {
215        json!({
216            "total": self.total(),
217            "speed": self.speed(),
218            "current": self.current(),
219            "max": self.max(),
220        })
221    }
222
223    // ── Aggregation helpers (mirroring Counter) ──
224
225    /// Sums `total`, `current` and `max` from `other` into `self`.
226    #[inline]
227    pub fn add(&self, other: &Self) {
228        self.total.fetch_add(other.total(), Ordering::Relaxed);
229        self.current.fetch_add(other.current(), Ordering::Relaxed);
230        self.max.fetch_add(other.max(), Ordering::Relaxed);
231    }
232
233    /// Replaces all fields with the values from `other`.
234    #[inline]
235    pub fn set(&self, other: &Self) {
236        self.total.store(other.total(), Ordering::Relaxed);
237        self.speed.store(f64::to_bits(other.speed()), Ordering::Relaxed);
238        self.last_total.store(other.total(), Ordering::Relaxed);
239        self.current.store(other.current(), Ordering::Relaxed);
240        self.max.store(other.max(), Ordering::Relaxed);
241    }
242
243    /// Merges `other` into `self` according to [`self.mode`](StatsMergeMode).
244    ///
245    /// * [`None`](StatsMergeMode::None) — no-op.
246    /// * [`Sum`](StatsMergeMode::Sum) — totals are summed,
247    ///   `current` / `max` take the larger value.
248    /// * [`Max`](StatsMergeMode::Max) / [`Min`](StatsMergeMode::Min) —
249    ///   only `current` and `max` are affected.
250    #[inline]
251    pub fn merge(&self, other: &Self) {
252        match self.mode {
253            StatsMergeMode::None => {}
254            StatsMergeMode::Sum => {
255                self.add(other);
256            }
257            StatsMergeMode::Max => {
258                self.current.fetch_max(other.current(), Ordering::Relaxed);
259                self.max.fetch_max(other.max(), Ordering::Relaxed);
260                self.total.fetch_max(other.total(), Ordering::Relaxed);
261            }
262            StatsMergeMode::Min => {
263                self.current.fetch_min(other.current(), Ordering::Relaxed);
264                self.max.fetch_min(other.max(), Ordering::Relaxed);
265                self.total.fetch_min(other.total(), Ordering::Relaxed);
266            }
267            _ => {}
268        }
269    }
270
271    /// Decrements the current count by 1 (total and max unchanged).
272    #[inline]
273    pub fn dec(&self) {
274        self.current.fetch_sub(1, Ordering::Relaxed);
275    }
276
277    /// Decrements the current count by `n` (total and max unchanged).
278    #[inline]
279    pub fn decs(&self, n: i64) {
280        self.current.fetch_sub(n, Ordering::Relaxed);
281    }
282}
283
284// ── Serde: snapshot-based serialisation ──
285
286/// Helper struct carrying the fields we serialise.
287#[derive(Serialize, Deserialize)]
288struct RateCounterData {
289    total: u64,
290    speed: f64,
291    current: i64,
292    max: i64,
293    mode: StatsMergeMode,
294}
295
296impl Serialize for RateCounter {
297    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
298        RateCounterData {
299            total: self.total(),
300            speed: self.speed(),
301            current: self.current(),
302            max: self.max(),
303            mode: self.mode.clone(),
304        }
305        .serialize(serializer)
306    }
307}
308
309impl<'de> Deserialize<'de> for RateCounter {
310    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
311        let data = RateCounterData::deserialize(deserializer)?;
312        Ok(RateCounter {
313            total: Arc::new(AtomicU64::new(data.total)),
314            speed: Arc::new(AtomicU64::new(f64::to_bits(data.speed))),
315            last_total: Arc::new(AtomicU64::new(data.total)),
316            current: Arc::new(AtomicI64::new(data.current)),
317            max: Arc::new(AtomicI64::new(data.max)),
318            mode: data.mode,
319        })
320    }
321}
322
323#[cfg(test)]
324mod tests {
325    use super::*;
326    use std::time::Duration;
327
328    #[test]
329    fn new_is_zero() {
330        let rc = RateCounter::new();
331        assert_eq!(rc.total(), 0);
332        assert_eq!(rc.speed(), 0.0);
333        assert_eq!(rc.current(), 0);
334        assert_eq!(rc.max(), 0);
335    }
336
337    #[test]
338    fn inc_updates_current_and_max() {
339        let rc = RateCounter::new();
340        rc.inc();
341        assert_eq!(rc.total(), 1);
342        assert_eq!(rc.current(), 1);
343        assert_eq!(rc.max(), 1);
344    }
345
346    #[test]
347    fn incs_updates_current_and_max() {
348        let rc = RateCounter::new();
349        rc.incs(42);
350        assert_eq!(rc.total(), 42);
351        assert_eq!(rc.current(), 42);
352        assert_eq!(rc.max(), 42);
353    }
354
355    #[test]
356    fn max_tracks_peak() {
357        let rc = RateCounter::new();
358        rc.incs(10);
359        assert_eq!(rc.max(), 10);
360
361        rc.decs(5);
362        assert_eq!(rc.current(), 5);
363        assert_eq!(rc.max(), 10, "decrease should not lower max");
364
365        rc.incs(20);
366        assert_eq!(rc.current(), 25);
367        assert_eq!(rc.max(), 25, "new higher peak should update max");
368    }
369
370    #[test]
371    fn dec_does_not_affect_total_nor_max() {
372        let rc = RateCounter::new();
373        rc.incs(10);
374        assert_eq!(rc.total(), 10);
375        assert_eq!(rc.max(), 10);
376
377        rc.dec();
378        assert_eq!(rc.total(), 10, "dec should not change total");
379        assert_eq!(rc.current(), 9);
380        assert_eq!(rc.max(), 10, "dec should not change max");
381
382        rc.decs(5);
383        assert_eq!(rc.total(), 10, "decs should not change total");
384        assert_eq!(rc.current(), 4);
385        assert_eq!(rc.max(), 10, "decs should not change max");
386    }
387
388    #[test]
389    fn tick_does_not_affect_current_nor_max() {
390        let rc = RateCounter::new();
391        rc.incs(100);
392        assert_eq!(rc.current(), 100);
393        assert_eq!(rc.max(), 100);
394
395        rc.tick(Duration::from_secs(1));
396        assert!((rc.speed() - 100.0).abs() < f64::EPSILON);
397        assert_eq!(rc.total(), 100);
398        assert_eq!(rc.current(), 100, "tick should not affect current");
399        assert_eq!(rc.max(), 100, "tick should not affect max");
400    }
401
402    #[test]
403    fn reset_clears_all_including_max() {
404        let rc = RateCounter::new();
405        rc.incs(100);
406        rc.decs(20);
407        rc.tick(Duration::from_secs(1));
408        rc.reset();
409        assert_eq!(rc.total(), 0);
410        assert_eq!(rc.speed(), 0.0);
411        assert_eq!(rc.current(), 0);
412        assert_eq!(rc.max(), 0);
413
414        // tick after reset should yield 0
415        rc.tick(Duration::from_secs(1));
416        assert_eq!(rc.speed(), 0.0);
417    }
418
419    #[test]
420    fn concurrent_incs_max() {
421        use std::thread;
422
423        let rc = Arc::new(RateCounter::new());
424        let mut handles = Vec::new();
425
426        for _ in 0..8 {
427            let rc = rc.clone();
428            handles.push(thread::spawn(move || {
429                for _ in 0..1_000 {
430                    rc.inc();
431                }
432            }));
433        }
434
435        for h in handles {
436            h.join().unwrap();
437        }
438
439        assert_eq!(rc.total(), 8_000);
440        assert_eq!(rc.current(), 8_000);
441        // With 8 threads racing, the max should be 8_000 (all incs at once)
442        // or less (some threads may have started before others completed decs
443        // in other tests, but here there are no decs, so max == current)
444        assert_eq!(rc.max(), 8_000);
445    }
446
447    #[test]
448    fn clone_shares_max() {
449        let a = RateCounter::new();
450        let b = a.clone();
451
452        a.incs(10);
453        assert_eq!(b.total(), 10, "clone should see the same total");
454        assert_eq!(b.current(), 10, "clone should see the same current");
455        assert_eq!(b.max(), 10, "clone should see the same max");
456
457        b.dec();
458        assert_eq!(a.current(), 9, "original should see b's dec");
459
460        a.incs(20);
461        assert_eq!(b.current(), 29, "clone should see a's new current");
462        assert_eq!(b.max(), 29, "clone should see a's new max");
463
464        a.tick(Duration::from_secs(1));
465        assert!((b.speed() - 30.0).abs() < f64::EPSILON, "speed computed on a should also be visible on b");
466    }
467
468    #[test]
469    fn concurrent_inc_dec_pair_max() {
470        use std::thread;
471
472        let rc = Arc::new(RateCounter::new());
473        let mut handles = Vec::new();
474
475        // 8 threads each inc and dec 1_000 times → net current should be 0
476        for _ in 0..8 {
477            let rc = rc.clone();
478            handles.push(thread::spawn(move || {
479                for _ in 0..1_000 {
480                    rc.inc();
481                    rc.dec();
482                }
483            }));
484        }
485
486        for h in handles {
487            h.join().unwrap();
488        }
489
490        // total should reflect all incs
491        assert_eq!(rc.total(), 8_000);
492        // net current should be 0
493        assert_eq!(rc.current(), 0);
494        // max should be > 0 (at some point there were concurrent incs)
495        assert!(rc.max() > 0, "concurrent inc/dec should have produced a peak");
496    }
497
498    #[test]
499    fn to_json_includes_all_fields() {
500        let rc = RateCounter::new();
501        rc.incs(100);
502        rc.decs(20);
503        rc.tick(Duration::from_secs(5));
504
505        let json = rc.to_json();
506        let obj = json.as_object().expect("to_json should return an object");
507
508        assert_eq!(obj["total"].as_u64(), Some(100));
509        assert!((obj["speed"].as_f64().unwrap() - 20.0).abs() < f64::EPSILON);
510        assert_eq!(obj["current"].as_i64(), Some(80));
511        assert_eq!(obj["max"].as_i64(), Some(100));
512    }
513}
rmqtt_utils/rate_counter.rs

rmqtt_utils/
rate_counter.rs
