vendor/portable-atomic/src/imp/fallback/seq_lock.rs

   1 // Adapted from https://github.com/crossbeam-rs/crossbeam/blob/crossbeam-utils-0.8.7/crossbeam-utils/src/atomic/seq_lock.rs.
   2
   3 use core::{
   4     mem::ManuallyDrop,
   5     sync::atomic::{self, Ordering},
   6 };
   7
   8 use super::utils::Backoff;
   9
  10 // See mod.rs for details.
  11 #[cfg(any(target_pointer_width = "16", target_pointer_width = "32"))]
  12 pub(super) use core::sync::atomic::AtomicU64 as AtomicStamp;
  13 #[cfg(not(any(target_pointer_width = "16", target_pointer_width = "32")))]
  14 pub(super) use core::sync::atomic::AtomicUsize as AtomicStamp;
  15 #[cfg(not(any(target_pointer_width = "16", target_pointer_width = "32")))]
  16 pub(super) type Stamp = usize;
  17 #[cfg(any(target_pointer_width = "16", target_pointer_width = "32"))]
  18 pub(super) type Stamp = u64;
  19
  20 // See mod.rs for details.
  21 pub(super) type AtomicChunk = AtomicStamp;
  22 pub(super) type Chunk = Stamp;
  23
  24 /// A simple stamped lock.
  25 pub(super) struct SeqLock {
  26     /// The current state of the lock.
  27     ///
  28     /// All bits except the least significant one hold the current stamp. When locked, the state
  29     /// equals 1 and doesn't contain a valid stamp.
  30     state: AtomicStamp,
  31 }
  32
  33 impl SeqLock {
  34     #[inline]
  35     pub(super) const fn new() -> Self {
  36         Self { state: AtomicStamp::new(0) }
  37     }
  38
  39     /// If not locked, returns the current stamp.
  40     ///
  41     /// This method should be called before optimistic reads.
  42     #[inline]
  43     pub(super) fn optimistic_read(&self) -> Option<Stamp> {
  44         let state = self.state.load(Ordering::Acquire);
  45         if state == 1 {
  46             None
  47         } else {
  48             Some(state)
  49         }
  50     }
  51
  52     /// Returns `true` if the current stamp is equal to `stamp`.
  53     ///
  54     /// This method should be called after optimistic reads to check whether they are valid. The
  55     /// argument `stamp` should correspond to the one returned by method `optimistic_read`.
  56     #[inline]
  57     pub(super) fn validate_read(&self, stamp: Stamp) -> bool {
  58         atomic::fence(Ordering::Acquire);
  59         self.state.load(Ordering::Relaxed) == stamp
  60     }
  61
  62     /// Grabs the lock for writing.
  63     #[inline]
  64     pub(super) fn write(&self) -> SeqLockWriteGuard<'_> {
  65         let mut backoff = Backoff::new();
  66         loop {
  67             let previous = self.state.swap(1, Ordering::Acquire);
  68
  69             if previous != 1 {
  70                 atomic::fence(Ordering::Release);
  71
  72                 return SeqLockWriteGuard { lock: self, state: previous };
  73             }
  74
  75             while self.state.load(Ordering::Relaxed) == 1 {
  76                 backoff.snooze();
  77             }
  78         }
  79     }
  80 }
  81
  82 /// An RAII guard that releases the lock and increments the stamp when dropped.
  83 #[must_use]
  84 pub(super) struct SeqLockWriteGuard<'a> {
  85     /// The parent lock.
  86     lock: &'a SeqLock,
  87
  88     /// The stamp before locking.
  89     state: Stamp,
  90 }
  91
  92 impl SeqLockWriteGuard<'_> {
  93     /// Releases the lock without incrementing the stamp.
  94     #[inline]
  95     pub(super) fn abort(self) {
  96         // We specifically don't want to call drop(), since that's
  97         // what increments the stamp.
  98         let this = ManuallyDrop::new(self);
  99
 100         // Restore the stamp.
 101         //
 102         // Release ordering for synchronizing with `optimistic_read`.
 103         this.lock.state.store(this.state, Ordering::Release);
 104     }
 105 }
 106
 107 impl Drop for SeqLockWriteGuard<'_> {
 108     #[inline]
 109     fn drop(&mut self) {
 110         // Release the lock and increment the stamp.
 111         //
 112         // Release ordering for synchronizing with `optimistic_read`.
 113         self.lock.state.store(self.state.wrapping_add(2), Ordering::Release);
 114     }
 115 }
 116
 117 #[cfg(test)]
 118 mod tests {
 119     use super::SeqLock;
 120
 121     #[test]
 122     fn smoke() {
 123         let lock = SeqLock::new();
 124         let before = lock.optimistic_read().unwrap();
 125         assert!(lock.validate_read(before));
 126         {
 127             let _guard = lock.write();
 128         }
 129         assert!(!lock.validate_read(before));
 130         let after = lock.optimistic_read().unwrap();
 131         assert_ne!(before, after);
 132     }
 133
 134     #[test]
 135     fn test_abort() {
 136         let lock = SeqLock::new();
 137         let before = lock.optimistic_read().unwrap();
 138         {
 139             let guard = lock.write();
 140             guard.abort();
 141         }
 142         let after = lock.optimistic_read().unwrap();
 143         assert_eq!(before, after, "aborted write does not update the stamp");
 144     }
 145 }