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 {
48            None
49        } else {
50            Some(state)
51        }
52    }
53
54    /// Returns `true` if the current stamp is equal to `stamp`.
55    ///
56    /// This method should be called after optimistic reads to check whether they are valid. The
57    /// argument `stamp` should correspond to the one returned by method `optimistic_read`.
58    #[inline]
59    pub(super) fn validate_read(&self, stamp: Stamp) -> bool {
60        atomic::fence(Ordering::Acquire);
61        self.state.load(Ordering::Relaxed) == stamp
62    }
63
64    /// Grabs the lock for writing.
65    #[inline]
66    pub(super) fn write(&self) -> SeqLockWriteGuard<'_> {
67        let mut backoff = Backoff::new();
68        loop {
69            let previous = self.state.swap(1, Ordering::Acquire);
70
71            if previous != 1 {
72                atomic::fence(Ordering::Release);
73
74                return SeqLockWriteGuard { lock: self, state: previous };
75            }
76
77            while self.state.load(Ordering::Relaxed) == 1 {
78                backoff.snooze();
79            }
80        }
81    }
82}
83
84/// An RAII guard that releases the lock and increments the stamp when dropped.
85#[must_use]
86pub(super) struct SeqLockWriteGuard<'a> {
87    /// The parent lock.
88    lock: &'a SeqLock,
89
90    /// The stamp before locking.
91    state: Stamp,
92}
93
94impl SeqLockWriteGuard<'_> {
95    /// Releases the lock without incrementing the stamp.
96    #[inline]
97    pub(super) fn abort(self) {
98        // We specifically don't want to call drop(), since that's
99        // what increments the stamp.
100        let this = ManuallyDrop::new(self);
101
102        // Restore the stamp.
103        //
104        // Release ordering for synchronizing with `optimistic_read`.
105        this.lock.state.store(this.state, Ordering::Release);
106    }
107}
108
109impl Drop for SeqLockWriteGuard<'_> {
110    #[inline]
111    fn drop(&mut self) {
112        // Release the lock and increment the stamp.
113        //
114        // Release ordering for synchronizing with `optimistic_read`.
115        self.lock.state.store(self.state.wrapping_add(2), Ordering::Release);
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::SeqLock;
122
123    #[test]
124    fn smoke() {
125        let lock = SeqLock::new();
126        let before = lock.optimistic_read().unwrap();
127        assert!(lock.validate_read(before));
128        {
129            let _guard = lock.write();
130        }
131        assert!(!lock.validate_read(before));
132        let after = lock.optimistic_read().unwrap();
133        assert_ne!(before, after);
134    }
135
136    #[test]
137    fn test_abort() {
138        let lock = SeqLock::new();
139        let before = lock.optimistic_read().unwrap();
140        {
141            let guard = lock.write();
142            guard.abort();
143        }
144        let after = lock.optimistic_read().unwrap();
145        assert_eq!(before, after, "aborted write does not update the stamp");
146    }
147}