portable_atomic/imp/fallback/
seq_lock.rs

1// SPDX-License-Identifier: Apache-2.0 OR MIT
2
3// Adapted from https://github.com/crossbeam-rs/crossbeam/blob/crossbeam-utils-0.8.7/crossbeam-utils/src/atomic/seq_lock.rs.
4
5use core::{
6    mem::ManuallyDrop,
7    sync::atomic::{self, Ordering},
8};
9
10use super::utils::Backoff;
11
12// See mod.rs for details.
13#[cfg(any(target_pointer_width = "16", target_pointer_width = "32"))]
14pub(super) use core::sync::atomic::AtomicU64 as AtomicStamp;
15#[cfg(not(any(target_pointer_width = "16", target_pointer_width = "32")))]
16pub(super) use core::sync::atomic::AtomicUsize as AtomicStamp;
17#[cfg(not(any(target_pointer_width = "16", target_pointer_width = "32")))]
18pub(super) type Stamp = usize;
19#[cfg(any(target_pointer_width = "16", target_pointer_width = "32"))]
20pub(super) type Stamp = u64;
21
22// See mod.rs for details.
23pub(super) type AtomicChunk = AtomicStamp;
24pub(super) type Chunk = Stamp;
25
26/// A simple stamped lock.
27pub(super) struct SeqLock {
28    /// The current state of the lock.
29    ///
30    /// All bits except the least significant one hold the current stamp. When locked, the state
31    /// equals 1 and doesn't contain a valid stamp.
32    state: AtomicStamp,
33}
34
35impl SeqLock {
36    #[inline]
37    pub(super) const fn new() -> Self {
38        Self { state: AtomicStamp::new(0) }
39    }
40
41    /// If not locked, returns the current stamp.
42    ///
43    /// This method should be called before optimistic reads.
44    #[inline]
45    pub(super) fn optimistic_read(&self) -> Option<Stamp> {
46        let state = self.state.load(Ordering::Acquire);
47        if state == 1 { None } else { Some(state) }
48    }
49
50    /// Returns `true` if the current stamp is equal to `stamp`.
51    ///
52    /// This method should be called after optimistic reads to check whether they are valid. The
53    /// argument `stamp` should correspond to the one returned by method `optimistic_read`.
54    #[inline]
55    pub(super) fn validate_read(&self, stamp: Stamp) -> bool {
56        atomic::fence(Ordering::Acquire);
57        self.state.load(Ordering::Relaxed) == stamp
58    }
59
60    /// Grabs the lock for writing.
61    #[inline]
62    pub(super) fn write(&self) -> SeqLockWriteGuard<'_> {
63        let mut backoff = Backoff::new();
64        loop {
65            let previous = self.state.swap(1, Ordering::Acquire);
66
67            if previous != 1 {
68                atomic::fence(Ordering::Release);
69
70                return SeqLockWriteGuard { lock: self, state: previous };
71            }
72
73            while self.state.load(Ordering::Relaxed) == 1 {
74                backoff.snooze();
75            }
76        }
77    }
78}
79
80/// An RAII guard that releases the lock and increments the stamp when dropped.
81#[must_use]
82pub(super) struct SeqLockWriteGuard<'a> {
83    /// The parent lock.
84    lock: &'a SeqLock,
85
86    /// The stamp before locking.
87    state: Stamp,
88}
89
90impl SeqLockWriteGuard<'_> {
91    /// Releases the lock without incrementing the stamp.
92    #[inline]
93    pub(super) fn abort(self) {
94        // We specifically don't want to call drop(), since that's
95        // what increments the stamp.
96        let this = ManuallyDrop::new(self);
97
98        // Restore the stamp.
99        //
100        // Release ordering for synchronizing with `optimistic_read`.
101        this.lock.state.store(this.state, Ordering::Release);
102    }
103}
104
105impl Drop for SeqLockWriteGuard<'_> {
106    #[inline]
107    fn drop(&mut self) {
108        // Release the lock and increment the stamp.
109        //
110        // Release ordering for synchronizing with `optimistic_read`.
111        self.lock.state.store(self.state.wrapping_add(2), Ordering::Release);
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::SeqLock;
118
119    #[test]
120    fn smoke() {
121        let lock = SeqLock::new();
122        let before = lock.optimistic_read().unwrap();
123        assert!(lock.validate_read(before));
124        {
125            let _guard = lock.write();
126        }
127        assert!(!lock.validate_read(before));
128        let after = lock.optimistic_read().unwrap();
129        assert_ne!(before, after);
130    }
131
132    #[test]
133    fn test_abort() {
134        let lock = SeqLock::new();
135        let before = lock.optimistic_read().unwrap();
136        {
137            let guard = lock.write();
138            guard.abort();
139        }
140        let after = lock.optimistic_read().unwrap();
141        assert_eq!(before, after, "aborted write does not update the stamp");
142    }
143}