Implement WakerListLock

WakerListLock is an optimised version of Spinlock<WakerList> which is
more efficient in performance and space.

Author: nbdd0121

src/sync/mutex.rs

@@ -7,7 +7,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use crate::future::Future;
 use crate::task::{Context, Poll};
 
-use super::waker_list::WakerList;
+use super::waker_list::{WakerList, WakerListLock};
 use std::num::NonZeroUsize;
 
 /// Set if the mutex is locked.
@@ -18,7 +18,7 @@ const BLOCKED: usize = 1 << 1;
 
 struct RawMutex {
     state: AtomicUsize,
-    blocked: std::sync::Mutex<WakerList>,
+    blocked: WakerListLock,
 }
 
 unsafe impl Send for RawMutex {}
@@ -30,7 +30,7 @@ impl RawMutex {
     pub fn new() -> RawMutex {
         RawMutex {
             state: AtomicUsize::new(0),
-            blocked: std::sync::Mutex::new(WakerList::new()),
+            blocked: WakerListLock::new(WakerList::new()),
         }
     }
 
@@ -54,7 +54,7 @@ impl RawMutex {
 
     #[cold]
     fn unlock_slow(&self) {
-        let mut blocked = self.blocked.lock().unwrap();
+        let mut blocked = self.blocked.lock();
         blocked.wake_one_weak();
     }
 
@@ -84,7 +84,7 @@ impl<'a> Future for RawLockFuture<'a> {
             self.acquired = true;
             Poll::Ready(())
         } else {
-            let mut blocked = self.mutex.blocked.lock().unwrap();
+            let mut blocked = self.mutex.blocked.lock();
 
             // Register the current task.
             match self.opt_key {
@@ -124,7 +124,7 @@ impl<'a> Future for RawLockFuture<'a> {
 impl Drop for RawLockFuture<'_> {
     fn drop(&mut self) {
         if let Some(key) = self.opt_key {
-            let mut blocked = self.mutex.blocked.lock().unwrap();
+            let mut blocked = self.mutex.blocked.lock();
             let opt_waker = unsafe { blocked.remove(key) };
 
             if opt_waker.is_none() && !self.acquired {

src/sync/waker_list.rs

@@ -1,6 +1,9 @@
 use crate::task::Waker;
 
+use crossbeam_utils::Backoff;
 use std::num::NonZeroUsize;
+use std::ops::{Deref, DerefMut};
+use std::sync::atomic::{AtomicUsize, Ordering};
 
 struct WakerNode {
     /// Previous `WakerNode` in the queue. If this node is the first node, it shall point to the last node.
@@ -19,6 +22,7 @@ unsafe impl Sync for WakerList {}
 
 impl WakerList {
     /// Create a new empty `WakerList`
+    #[inline]
     pub fn new() -> Self {
         Self {
             head: std::ptr::null_mut(),
@@ -128,3 +132,67 @@ impl<'a> Iterator for Iter<'a> {
         Some(unsafe { &mut (*ptr).waker })
     }
 }
+
+/// This is identical to a Spinlock<WakerList>, but is efficient in space and occupies the same
+/// amount of memory as a pointer. Performance-wise it can be better than a spinlock because once
+/// the lock is acquired, modification to the WakerList is on the stack, and thus is not on the
+/// same cache line as the atomic variable itself, relieving some contention.
+pub struct WakerListLock {
+    /// If this value is 1, it represents that it is locked.
+    /// All other values represent a valid `*mut WakerNode`.
+    value: AtomicUsize,
+}
+
+impl WakerListLock {
+    /// Returns a new pointer lock initialized with `value`.
+    #[inline]
+    pub fn new(value: WakerList) -> Self {
+        Self {
+            value: AtomicUsize::new(value.head as usize),
+        }
+    }
+
+    /// Locks the `WakerListLock`.
+    pub fn lock(&self) -> WakerListLockGuard<'_> {
+        let backoff = Backoff::new();
+        loop {
+            let value = self.value.swap(1, Ordering::Acquire);
+            if value != 1 {
+                return WakerListLockGuard {
+                    parent: self,
+                    list: WakerList {
+                        head: value as *mut WakerNode,
+                    },
+                };
+            }
+            backoff.snooze();
+        }
+    }
+}
+
+pub struct WakerListLockGuard<'a> {
+    parent: &'a WakerListLock,
+    list: WakerList,
+}
+
+impl<'a> Drop for WakerListLockGuard<'a> {
+    fn drop(&mut self) {
+        self.parent
+            .value
+            .store(self.list.head as usize, Ordering::Release);
+    }
+}
+
+impl<'a> Deref for WakerListLockGuard<'a> {
+    type Target = WakerList;
+
+    fn deref(&self) -> &WakerList {
+        &self.list
+    }
+}
+
+impl<'a> DerefMut for WakerListLockGuard<'a> {
+    fn deref_mut(&mut self) -> &mut WakerList {
+        &mut self.list
+    }
+}