Skip to main content

moirai_utils/
atomic.rs

1//! Atomic operations and counters for lock-free programming.
2//!
3//! This module provides atomic utilities for building lock-free data structures
4//! and implementing efficient counters and statistics tracking.
5
6use core::sync::atomic::{AtomicUsize, Ordering};
7
8/// A thread-safe atomic counter with increment and decrement operations.
9///
10/// This counter uses relaxed ordering for maximum performance while
11/// maintaining atomicity across threads.
12#[derive(Debug)]
13pub struct AtomicCounter {
14    value: AtomicUsize,
15}
16
17impl AtomicCounter {
18    /// Create a new atomic counter with initial value 0.
19    pub const fn new() -> Self {
20        Self {
21            value: AtomicUsize::new(0),
22        }
23    }
24
25    /// Create a new atomic counter with the given initial value.
26    pub const fn with_initial(initial: usize) -> Self {
27        Self {
28            value: AtomicUsize::new(initial),
29        }
30    }
31
32    /// Get the current value of the counter.
33    pub fn get(&self) -> usize {
34        self.value.load(Ordering::Relaxed)
35    }
36
37    /// Set the counter to a specific value.
38    pub fn set(&self, value: usize) {
39        self.value.store(value, Ordering::Relaxed);
40    }
41
42    /// Increment the counter by 1 and return the previous value.
43    pub fn increment(&self) -> usize {
44        self.value.fetch_add(1, Ordering::Relaxed)
45    }
46
47    /// Increment the counter by a specific amount and return the previous value.
48    pub fn add(&self, amount: usize) -> usize {
49        self.value.fetch_add(amount, Ordering::Relaxed)
50    }
51
52    /// Atomic fetch-and-add operation compatible with std::sync::atomic interface.
53    /// This provides a lower-level interface for compatibility.
54    pub fn fetch_add(&self, amount: usize, ordering: Ordering) -> usize {
55        self.value.fetch_add(amount, ordering)
56    }
57
58    /// Decrement the counter by 1 and return the previous value.
59    /// Note: This will wrap around on underflow.
60    pub fn decrement(&self) -> usize {
61        self.value.fetch_sub(1, Ordering::Relaxed)
62    }
63
64    /// Decrement the counter by a specific amount and return the previous value.
65    /// Note: This will wrap around on underflow.
66    pub fn subtract(&self, amount: usize) -> usize {
67        self.value.fetch_sub(amount, Ordering::Relaxed)
68    }
69
70    /// Atomically set the counter to the maximum of its current value and the given value.
71    pub fn max(&self, value: usize) -> usize {
72        self.value.fetch_max(value, Ordering::Relaxed)
73    }
74
75    /// Atomically set the counter to the minimum of its current value and the given value.
76    pub fn min(&self, value: usize) -> usize {
77        self.value.fetch_min(value, Ordering::Relaxed)
78    }
79
80    /// Reset the counter to 0 and return the previous value.
81    pub fn reset(&self) -> usize {
82        self.value.swap(0, Ordering::Relaxed)
83    }
84
85    /// Compare and swap the counter value.
86    /// Returns the previous value and whether the swap was successful.
87    pub fn compare_and_swap(&self, current: usize, new: usize) -> (usize, bool) {
88        match self
89            .value
90            .compare_exchange_weak(current, new, Ordering::Relaxed, Ordering::Relaxed)
91        {
92            Ok(prev) => (prev, true),
93            Err(prev) => (prev, false),
94        }
95    }
96}
97
98impl Default for AtomicCounter {
99    fn default() -> Self {
100        Self::new()
101    }
102}
103
104impl Clone for AtomicCounter {
105    fn clone(&self) -> Self {
106        Self::with_initial(self.get())
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[test]
115    fn test_atomic_counter_basic() {
116        let counter = AtomicCounter::new();
117        assert_eq!(counter.get(), 0);
118
119        counter.increment();
120        assert_eq!(counter.get(), 1);
121
122        counter.decrement();
123        assert_eq!(counter.get(), 0);
124    }
125
126    #[test]
127    fn test_atomic_counter_add_subtract() {
128        let counter = AtomicCounter::new();
129
130        counter.add(10);
131        assert_eq!(counter.get(), 10);
132
133        counter.subtract(5);
134        assert_eq!(counter.get(), 5);
135    }
136
137    #[test]
138    fn test_atomic_counter_max_min() {
139        let counter = AtomicCounter::with_initial(5);
140
141        counter.max(3); // Should remain 5
142        assert_eq!(counter.get(), 5);
143
144        counter.max(10); // Should become 10
145        assert_eq!(counter.get(), 10);
146
147        counter.min(15); // Should remain 10
148        assert_eq!(counter.get(), 10);
149
150        counter.min(7); // Should become 7
151        assert_eq!(counter.get(), 7);
152    }
153
154    #[test]
155    fn test_atomic_counter_reset() {
156        let counter = AtomicCounter::with_initial(42);
157        let prev = counter.reset();
158        assert_eq!(prev, 42);
159        assert_eq!(counter.get(), 0);
160    }
161}