1use core::mem;
2use core::sync::atomic::{self, AtomicUsize, Ordering};
3
4use crate::Backoff;
5
6/// A simple stamped lock.
7pub(crate) struct SeqLock {
8 /// The current state of the lock.
9 ///
10 /// All bits except the least significant one hold the current stamp. When locked, the state
11 /// equals 1 and doesn't contain a valid stamp.
12 state: AtomicUsize,
13}
14
15impl SeqLock {
16 pub(crate) const fn new() -> Self {
17 Self {
18 state: AtomicUsize::new(0),
19 }
20 }
21
22 /// If not locked, returns the current stamp.
23 ///
24 /// This method should be called before optimistic reads.
25 #[inline]
26 pub(crate) fn optimistic_read(&self) -> Option<usize> {
27 let state = self.state.load(Ordering::Acquire);
28 if state == 1 {
29 None
30 } else {
31 Some(state)
32 }
33 }
34
35 /// Returns `true` if the current stamp is equal to `stamp`.
36 ///
37 /// This method should be called after optimistic reads to check whether they are valid. The
38 /// argument `stamp` should correspond to the one returned by method `optimistic_read`.
39 #[inline]
40 pub(crate) fn validate_read(&self, stamp: usize) -> bool {
41 atomic::fence(Ordering::Acquire);
42 self.state.load(Ordering::Relaxed) == stamp
43 }
44
45 /// Grabs the lock for writing.
46 #[inline]
47 pub(crate) fn write(&'static self) -> SeqLockWriteGuard {
48 let backoff = Backoff::new();
49 loop {
50 let previous = self.state.swap(1, Ordering::Acquire);
51
52 if previous != 1 {
53 atomic::fence(Ordering::Release);
54
55 return SeqLockWriteGuard {
56 lock: self,
57 state: previous,
58 };
59 }
60
61 backoff.snooze();
62 }
63 }
64}
65
66/// An RAII guard that releases the lock and increments the stamp when dropped.
67pub(crate) struct SeqLockWriteGuard {
68 /// The parent lock.
69 lock: &'static SeqLock,
70
71 /// The stamp before locking.
72 state: usize,
73}
74
75impl SeqLockWriteGuard {
76 /// Releases the lock without incrementing the stamp.
77 #[inline]
78 pub(crate) fn abort(self) {
79 self.lock.state.store(self.state, order:Ordering::Release);
80
81 // We specifically don't want to call drop(), since that's
82 // what increments the stamp.
83 mem::forget(self);
84 }
85}
86
87impl Drop for SeqLockWriteGuard {
88 #[inline]
89 fn drop(&mut self) {
90 // Release the lock and increment the stamp.
91 self.lock
92 .state
93 .store(self.state.wrapping_add(2), order:Ordering::Release);
94 }
95}
96
97#[cfg(test)]
98mod tests {
99 use super::SeqLock;
100
101 #[test]
102 fn test_abort() {
103 static LK: SeqLock = SeqLock::new();
104 let before = LK.optimistic_read().unwrap();
105 {
106 let guard = LK.write();
107 guard.abort();
108 }
109 let after = LK.optimistic_read().unwrap();
110 assert_eq!(before, after, "aborted write does not update the stamp");
111 }
112}
113