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> = ManuallyDrop::new(self);
97
98 // Restore the stamp.
99 //
100 // Release ordering for synchronizing with `optimistic_read`.
101 this.lock.state.store(val:this.state, order: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), order: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}
144