Skip to main content

shared_state/
counter.rs

1//! Atomic counter utilities.
2
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::Arc;
5
6/// A thread-safe atomic counter.
7pub struct SharedCounter {
8    inner: AtomicU64,
9}
10
11impl Default for SharedCounter {
12    fn default() -> Self {
13        Self::new()
14    }
15}
16
17impl SharedCounter {
18    /// Create a new counter at zero.
19    pub fn new() -> Self {
20        Self {
21            inner: AtomicU64::new(0),
22        }
23    }
24
25    /// Create a counter with an initial value.
26    pub fn with_value(value: u64) -> Self {
27        Self {
28            inner: AtomicU64::new(value),
29        }
30    }
31
32    /// Increment by 1 and return the new value.
33    pub fn increment(&self) -> u64 {
34        self.inner.fetch_add(1, Ordering::Relaxed) + 1
35    }
36
37    /// Increment by `n` and return the new value.
38    pub fn increment_by(&self, n: u64) -> u64 {
39        self.inner.fetch_add(n, Ordering::Relaxed) + n
40    }
41
42    /// Decrement by 1 and return the new value.
43    pub fn decrement(&self) -> u64 {
44        self.inner.fetch_sub(1, Ordering::Relaxed) - 1
45    }
46
47    /// Get the current value.
48    pub fn get(&self) -> u64 {
49        self.inner.load(Ordering::Relaxed)
50    }
51
52    /// Reset to zero and return the previous value.
53    pub fn reset(&self) -> u64 {
54        self.inner.swap(0, Ordering::Relaxed)
55    }
56}
57
58impl SharedCounter {
59    /// Wrap in Arc for shared access across tasks.
60    pub fn into_shared(self) -> Arc<Self> {
61        Arc::new(self)
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn increment_decrement() {
71        let c = SharedCounter::new();
72        assert_eq!(c.increment(), 1);
73        assert_eq!(c.increment(), 2);
74        assert_eq!(c.decrement(), 1);
75        assert_eq!(c.get(), 1);
76    }
77
78    #[test]
79    fn increment_by() {
80        let c = SharedCounter::new();
81        assert_eq!(c.increment_by(10), 10);
82        assert_eq!(c.increment_by(5), 15);
83    }
84
85    #[test]
86    fn reset() {
87        let c = SharedCounter::with_value(42);
88        assert_eq!(c.reset(), 42);
89        assert_eq!(c.get(), 0);
90    }
91
92    #[tokio::test]
93    async fn concurrent_increment() {
94        let counter = SharedCounter::new().into_shared();
95        let mut handles = vec![];
96        for _ in 0..100 {
97            let c = counter.clone();
98            handles.push(tokio::spawn(async move {
99                c.increment();
100            }));
101        }
102        for h in handles {
103            h.await.unwrap();
104        }
105        assert_eq!(counter.get(), 100);
106    }
107}