added try_read and try_write

Author: rj-calvin

futures-util/src/lock/backoff.rs

@@ -0,0 +1,32 @@
+use std::cell::Cell;
+use std::sync::atomic::spin_loop_hint;
+use std::thread::yield_now;
+
+// Taken from crossbeam_utils
+const SPIN_LIMIT: u32 = 6;
+const YIELD_LIMIT: u32 = 10;
+
+pub(crate) struct Backoff {
+    step: Cell<u32>,
+}
+
+impl Backoff {
+    pub(crate) fn new() -> Backoff {
+        Backoff {
+            step: Cell::new(0u32),
+        }
+    }
+
+    pub(crate) fn snooze(&self) {
+        if self.step.get() <= SPIN_LIMIT {
+            for _ in 0..1 << self.step.get() {
+                spin_loop_hint();
+            }
+        } else {
+            yield_now();
+        }
+        if self.step.get() <= YIELD_LIMIT {
+            self.step.set(self.step.get() + 1);
+        }
+    }
+}

futures-util/src/lock/mod.rs

@@ -16,6 +16,8 @@ pub use self::rwlock::{
 };
 #[cfg(feature = "std")]
 pub(crate) mod waiter;
+#[cfg(feature = "std")]
+pub(crate) mod backoff;
 
 #[cfg(any(feature = "bilock", feature = "sink", feature = "io"))]
 #[cfg_attr(docsrs, doc(cfg(feature = "bilock")))]

futures-util/src/lock/rwlock.rs

@@ -5,8 +5,18 @@ use std::cell::UnsafeCell;
 use std::fmt;
 use std::ops::{Deref, DerefMut};
 use std::pin::Pin;
+use std::sync::RwLock as StdRwLock;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
+#[allow(clippy::identity_op)]
+const PHASE: usize = 1 << 0;
+const ONE_WRITER: usize = 1 << 1;
+const ONE_READER: usize = 1 << 2;
+const WRITE_BITS: usize = ONE_WRITER | PHASE;
+
+// Sentinel for when no slot in the `Slab` has been dedicated to this object.
+const WAIT_KEY_NONE: usize = usize::max_value();
+
 struct State {
     ins: AtomicUsize,
     out: AtomicUsize,
@@ -33,16 +43,26 @@ impl AtomicState {
         self.read.out.fetch_add(ONE_READER, Ordering::SeqCst)
     }
 
+    #[inline]
+    fn waiting_writers(&self) -> usize {
+        self.write.ins.load(Ordering::Relaxed)
+    }
+
     #[inline]
     fn reserve_reader(&self) -> usize {
         self.read.ins.fetch_add(ONE_READER, Ordering::SeqCst) & WRITE_BITS
     }
 
     #[inline]
-    fn reserve_writer(&self, phase: usize) -> usize {
+    fn reserve_writer(&self, ticket: usize) -> usize {
         self.read
             .ins
-            .fetch_add(ONE_WRITER | phase, Ordering::SeqCst)
+            .fetch_add(ONE_WRITER | (ticket & PHASE), Ordering::SeqCst)
+    }
+
+    #[inline]
+    fn reserve_transient_writer(&self) -> usize {
+        self.read.ins.fetch_add(PHASE, Ordering::SeqCst)
     }
 
     #[inline]
@@ -71,6 +91,8 @@ pub struct RwLock<T: ?Sized> {
     atomic: AtomicState,
     readers: WaiterSet,
     writers: WaiterSet,
+    block_read_tickets: StdRwLock<()>,
+    block_write_tickets: StdRwLock<()>,
     value: UnsafeCell<T>,
 }
 
@@ -80,15 +102,6 @@ impl<T: ?Sized> fmt::Debug for RwLock<T> {
     }
 }
 
-#[allow(clippy::identity_op)]
-const PHASE: usize = 1 << 0;
-const ONE_WRITER: usize = 1 << 1;
-const ONE_READER: usize = 1 << 2;
-const WRITE_BITS: usize = ONE_WRITER | PHASE;
-
-// Sentinel for when no slot in the `Slab` has been dedicated to this object.
-const WAIT_KEY_NONE: usize = usize::max_value();
-
 impl<T> RwLock<T> {
     /// Creates a new futures-aware read-write lock.
     pub fn new(t: T) -> RwLock<T> {
@@ -105,6 +118,8 @@ impl<T> RwLock<T> {
             },
             readers: WaiterSet::new(),
             writers: WaiterSet::new(),
+            block_read_tickets: StdRwLock::new(()),
+            block_write_tickets: StdRwLock::new(()),
             value: UnsafeCell::new(t),
         }
     }
@@ -149,6 +164,54 @@ impl<T: ?Sized> RwLock<T> {
         }
     }
 
+    /// Attempt to acquire a read access lock synchronously.
+    pub fn try_read(&self) -> Option<RwLockReadGuard<'_, T>> {
+        let lock = self.block_read_tickets.write().unwrap();
+        if self.atomic.phase() == 0 {
+            self.atomic.reserve_reader();
+            drop(lock);
+            self.writers.notify_all();
+            Some(RwLockReadGuard { rwlock: self })
+        } else {
+            drop(lock);
+            self.writers.notify_all();
+            None
+        }
+    }
+
+    /// Attempt to acquire a write access lock synchronously.
+    pub fn try_write(&self) -> Option<RwLockWriteGuard<'_, T>> {
+        let read_lock = self.block_read_tickets.write().unwrap();
+        if self.atomic.phase() == 0 {
+            let write_lock = self.block_write_tickets.write().unwrap();
+            if self.atomic.waiting_writers() == self.atomic.finished_writers()
+                && self.atomic.reserve_transient_writer() == self.atomic.finished_readers()
+            {
+                self.atomic.insert_writer();
+                drop(write_lock);
+                drop(read_lock);
+                self.writers.notify_all();
+                Some(RwLockWriteGuard { rwlock: self })
+            } else if self.atomic.phase() != 0 {
+                self.atomic.clear_phase();
+                drop(write_lock);
+                drop(read_lock);
+                self.writers.notify_all();
+                self.readers.notify_all();
+                None
+            } else {
+                drop(write_lock);
+                drop(read_lock);
+                self.writers.notify_all();
+                None
+            }
+        } else {
+            drop(read_lock);
+            self.writers.notify_all();
+            None
+        }
+    }
+
     /// Returns a mutable reference to the underlying data.
     ///
     /// Since this call borrows the lock mutably, no actual locking needs to
@@ -304,11 +367,13 @@ impl<'a, T: ?Sized> Future for RwLockWriteFuture<'a, T> {
 
         match self.ticket {
             None => {
+                let _write_lock = rwlock.block_write_tickets.read().unwrap();
                 let ticket = rwlock.atomic.insert_writer();
                 self.ticket = Some(Ticket::Write(ticket));
                 if ticket == rwlock.atomic.finished_writers() {
                     // Note that the WRITE_BITS are always cleared at this point.
-                    let ticket = rwlock.atomic.reserve_writer(ticket & PHASE);
+                    let _read_lock = rwlock.block_read_tickets.read().unwrap();
+                    let ticket = rwlock.atomic.reserve_writer(ticket);
                     self.ticket = Some(Ticket::Read(ticket));
                     if ticket == rwlock.atomic.finished_readers() {
                         self.rwlock = None;
@@ -325,7 +390,8 @@ impl<'a, T: ?Sized> Future for RwLockWriteFuture<'a, T> {
             Some(Ticket::Write(ticket)) => {
                 if ticket == rwlock.atomic.finished_writers() {
                     // Note that the WRITE_BITS are always cleared at this point.
-                    let ticket = rwlock.atomic.reserve_writer(ticket & PHASE);
+                    let _read_lock = rwlock.block_read_tickets.read().unwrap();
+                    let ticket = rwlock.atomic.reserve_writer(ticket);
                     self.ticket = Some(Ticket::Read(ticket));
                     if ticket == rwlock.atomic.finished_readers() {
                         rwlock.writers.remove(self.wait_key);

futures-util/src/lock/waiter.rs

@@ -1,14 +1,10 @@
+use crate::lock::backoff::Backoff;
 use slab::Slab;
-use std::cell::{Cell, UnsafeCell};
+use std::cell::UnsafeCell;
 use std::mem;
 use std::ops::{Deref, DerefMut};
-use std::sync::atomic::{spin_loop_hint, AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::task::Waker;
-use std::thread::yield_now;
-
-// Taken from crossbeam_utils
-const SPIN_LIMIT: u32 = 6;
-const YIELD_LIMIT: u32 = 10;
 
 #[allow(clippy::identity_op)]
 const IS_LOCKED: usize = 1 << 0;
@@ -142,30 +138,20 @@ impl WaiterSet {
         notified
     }
 
-    fn lock(&self) -> Lock<'_> {
-        let step = Cell::new(0u32);
+    fn lock(&self) -> WaiterSetGuard<'_> {
+        let backoff = Backoff::new();
         while self.state.fetch_or(IS_LOCKED, Ordering::Acquire) & IS_LOCKED != 0 {
-            if step.get() <= SPIN_LIMIT {
-                for _ in 0..1 << step.get() {
-                    spin_loop_hint();
-                }
-            } else {
-                yield_now();
-            }
-            if step.get() <= YIELD_LIMIT {
-                step.set(step.get() + 1);
-            }
+            backoff.snooze();
         }
-
-        Lock { waiter_set: self }
+        WaiterSetGuard { waiter_set: self }
     }
 }
 
-struct Lock<'a> {
+struct WaiterSetGuard<'a> {
     waiter_set: &'a WaiterSet,
 }
 
-impl Drop for Lock<'_> {
+impl Drop for WaiterSetGuard<'_> {
     fn drop(&mut self) {
         let mut state = 0;
 
@@ -181,15 +167,15 @@ impl Drop for Lock<'_> {
     }
 }
 
-impl Deref for Lock<'_> {
+impl Deref for WaiterSetGuard<'_> {
     type Target = Inner;
 
     fn deref(&self) -> &Inner {
         unsafe { &*self.waiter_set.inner.get() }
     }
 }
 
-impl DerefMut for Lock<'_> {
+impl DerefMut for WaiterSetGuard<'_> {
     fn deref_mut(&mut self) -> &mut Inner {
         unsafe { &mut *self.waiter_set.inner.get() }
     }