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> = ManuallyDrop::new(self);
101
102 // Restore the stamp.
103 //
104 // Release ordering for synchronizing with `optimistic_read`.
105 this.lock.state.store(val:this.state, order: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), order: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}
148