shared_state/
counter.rs

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

shared_state/
counter.rs
