1
use std::sync::atomic::AtomicU64;
2
use std::sync::atomic::Ordering;
3

            
4
use crossbeam_utils::CachePadded;
5
use thread_local::ThreadLocal;
6

            
7
/// A monotonic counter that avoids the cross-thread contention of a single
8
/// shared [`AtomicU64`].
9
///
10
/// Increments are cheap and uncontended; reading the total is `O(threads that
11
/// have touched this counter)` and is intended for occasional reporting rather
12
/// than the hot path. Like a relaxed atomic, a [`get`](Self::get) concurrent
13
/// with increments observes a value somewhere between the pre- and post-state
14
/// and carries no ordering guarantees with respect to other memory.
15
#[derive(Default)]
16
pub struct ShardedCounter {
17
    /// One cache-line-padded atomic per thread that has touched this counter.
18
    shards: ThreadLocal<CachePadded<AtomicU64>>,
19
}
20

            
21
impl ShardedCounter {
22
    /// Creates a new counter starting at zero.
23
    #[must_use]
24
2323
    pub fn new() -> Self {
25
2323
        ShardedCounter {
26
2323
            shards: ThreadLocal::new(),
27
2323
        }
28
2323
    }
29

            
30
    /// Adds `value` to the calling thread's shard. Uncontended with respect to
31
    /// other threads.
32
    #[inline]
33
800003
    pub fn add(&self, value: u64) {
34
800003
        self.shard().fetch_add(value, Ordering::Relaxed);
35
800003
    }
36

            
37
    /// Increments the calling thread's shard by one.
38
    #[inline]
39
800001
    pub fn increment(&self) {
40
800001
        self.add(1);
41
800001
    }
42

            
43
    /// Returns the sum of all per-thread shards.
44
    ///
45
    /// This iterates every thread's shard, so it is `O(threads)`; call it for
46
    /// reporting, not in a hot loop. The result is only a snapshot when no other
47
    /// thread is incrementing concurrently.
48
    #[must_use]
49
4
    pub fn get(&self) -> u64 {
50
10
        self.shards.iter().map(|shard| shard.load(Ordering::Relaxed)).sum()
51
4
    }
52

            
53
    /// Resets every shard to zero.
54
    ///
55
    /// Requires `&mut self`, which statically guarantees no concurrent
56
    /// increments are in flight.
57
1
    pub fn reset(&mut self) {
58
1
        for shard in self.shards.iter_mut() {
59
1
            shard.store(0, Ordering::Relaxed);
60
1
        }
61
1
    }
62

            
63
    /// Returns the calling thread's shard, creating it on first use.
64
    #[inline]
65
800003
    fn shard(&self) -> &AtomicU64 {
66
800003
        self.shards.get_or(|| CachePadded::new(AtomicU64::new(0)))
67
800003
    }
68
}
69

            
70
#[cfg(test)]
71
mod tests {
72
    use std::sync::Arc;
73
    use std::thread;
74

            
75
    use crate::ShardedCounter;
76

            
77
    #[test]
78
1
    fn test_single_thread() {
79
1
        let counter = ShardedCounter::new();
80
1
        assert_eq!(counter.get(), 0);
81

            
82
1
        counter.increment();
83
1
        counter.add(41);
84
1
        assert_eq!(counter.get(), 42);
85
1
    }
86

            
87
    #[test]
88
1
    fn test_reset() {
89
1
        let mut counter = ShardedCounter::new();
90
1
        counter.add(10);
91
1
        counter.reset();
92
1
        assert_eq!(counter.get(), 0);
93
1
    }
94

            
95
    #[test]
96
    #[cfg_attr(miri, ignore)] // Test is too slow under miri
97
1
    fn test_concurrent_increments() {
98
        const THREADS: u64 = 8;
99
        const PER_THREAD: u64 = 100_000;
100

            
101
1
        let counter = Arc::new(ShardedCounter::new());
102
1
        let handles: Vec<_> = (0..THREADS)
103
8
            .map(|_| {
104
8
                let counter = Arc::clone(&counter);
105
8
                thread::spawn(move || {
106
800000
                    for _ in 0..PER_THREAD {
107
800000
                        counter.increment();
108
800000
                    }
109
8
                })
110
8
            })
111
1
            .collect();
112

            
113
8
        for handle in handles {
114
8
            handle.join().unwrap();
115
8
        }
116

            
117
        // No increments are lost: each thread owns its own shard.
118
1
        assert_eq!(counter.get(), THREADS * PER_THREAD);
119
1
    }
120
}