1 use core::mem; 2 use core::sync::atomic::{self, AtomicUsize, Ordering}; 3 4 use crate::Backoff; 5 6 /// A simple stamped lock. 7 pub(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 15 impl SeqLock { new() -> Self16 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] optimistic_read(&self) -> Option<usize>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] validate_read(&self, stamp: usize) -> bool40 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] write(&'static self) -> SeqLockWriteGuard47 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. 67 pub(crate) struct SeqLockWriteGuard { 68 /// The parent lock. 69 lock: &'static SeqLock, 70 71 /// The stamp before locking. 72 state: usize, 73 } 74 75 impl SeqLockWriteGuard { 76 /// Releases the lock without incrementing the stamp. 77 #[inline] abort(self)78 pub(crate) fn abort(self) { 79 self.lock.state.store(self.state, 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 87 impl Drop for SeqLockWriteGuard { 88 #[inline] drop(&mut self)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), Ordering::Release); 94 } 95 } 96 97 #[cfg(test)] 98 mod tests { 99 use super::SeqLock; 100 101 #[test] test_abort()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