Replace Slab-backed waker list with linked list

Signed-off-by: Gary Guo <gary@garyguo.net>

Author: nbdd0121

src/sync/mod.rs

@@ -37,6 +37,7 @@ pub use rwlock::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 mod mutex;
 mod rwlock;
+mod waker_list;
 
 cfg_unstable! {
     pub use barrier::{Barrier, BarrierWaitResult};

src/sync/mutex.rs

@@ -4,10 +4,11 @@ use std::ops::{Deref, DerefMut};
 use std::pin::Pin;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
-use slab::Slab;
-
 use crate::future::Future;
-use crate::task::{Context, Poll, Waker};
+use crate::task::{Context, Poll};
+
+use super::waker_list::WakerList;
+use std::num::NonZeroUsize;
 
 /// Set if the mutex is locked.
 const LOCK: usize = 1;
@@ -17,7 +18,7 @@ const BLOCKED: usize = 1 << 1;
 
 struct RawMutex {
     state: AtomicUsize,
-    blocked: std::sync::Mutex<Slab<Option<Waker>>>,
+    blocked: std::sync::Mutex<WakerList>,
 }
 
 unsafe impl Send for RawMutex {}
@@ -29,7 +30,7 @@ impl RawMutex {
     pub fn new() -> RawMutex {
         RawMutex {
             state: AtomicUsize::new(0),
-            blocked: std::sync::Mutex::new(Slab::new()),
+            blocked: std::sync::Mutex::new(WakerList::new()),
         }
     }
 
@@ -54,13 +55,7 @@ impl RawMutex {
     #[cold]
     fn unlock_slow(&self) {
         let mut blocked = self.blocked.lock().unwrap();
-
-        if let Some((_, opt_waker)) = blocked.iter_mut().next() {
-            // If there is no waker in this entry, that means it was already woken.
-            if let Some(w) = opt_waker.take() {
-                w.wake();
-            }
-        }
+        blocked.wake_one_weak();
     }
 
     /// Unlock this mutex.
@@ -77,7 +72,7 @@ impl RawMutex {
 
 struct RawLockFuture<'a> {
     mutex: &'a RawMutex,
-    opt_key: Option<usize>,
+    opt_key: Option<NonZeroUsize>,
     acquired: bool,
 }
 
@@ -94,21 +89,22 @@ impl<'a> Future for RawLockFuture<'a> {
             // Register the current task.
             match self.opt_key {
                 None => {
+                    if blocked.is_empty() {
+                        self.mutex.state.fetch_or(BLOCKED, Ordering::Relaxed);
+                    }
+
                     // Insert a new entry into the list of blocked tasks.
                     let w = cx.waker().clone();
                     let key = blocked.insert(Some(w));
                     self.opt_key = Some(key);
-
-                    if blocked.len() == 1 {
-                        self.mutex.state.fetch_or(BLOCKED, Ordering::Relaxed);
-                    }
                 }
                 Some(key) => {
                     // There is already an entry in the list of blocked tasks. Just
                     // reset the waker if it was removed.
-                    if blocked[key].is_none() {
+                    let opt_waker = unsafe { blocked.get(key) };
+                    if opt_waker.is_none() {
                         let w = cx.waker().clone();
-                        blocked[key] = Some(w);
+                        *opt_waker = Some(w);
                     }
                 }
             }
@@ -129,16 +125,11 @@ impl Drop for RawLockFuture<'_> {
     fn drop(&mut self) {
         if let Some(key) = self.opt_key {
             let mut blocked = self.mutex.blocked.lock().unwrap();
-            let opt_waker = blocked.remove(key);
+            let opt_waker = unsafe { blocked.remove(key) };
 
             if opt_waker.is_none() && !self.acquired {
                 // We were awoken but didn't acquire the lock. Wake up another task.
-                if let Some((_, opt_waker)) = blocked.iter_mut().next() {
-                    // If there is no waker in this entry, that means it was already woken.
-                    if let Some(w) = opt_waker.take() {
-                        w.wake();
-                    }
-                }
+                blocked.wake_one_weak();
             }
 
             if blocked.is_empty() {

src/sync/waker_list.rs

@@ -0,0 +1,130 @@
+use crate::task::Waker;
+
+use std::num::NonZeroUsize;
+
+struct WakerNode {
+    /// Previous `WakerNode` in the queue. If this node is the first node, it shall point to the last node.
+    prev_in_queue: *mut WakerNode,
+    /// Next `WakerNode` in the queue. If this node is the last node, it shall be null.
+    next_in_queue: *mut WakerNode,
+    waker: Option<Waker>,
+}
+
+pub struct WakerList {
+    head: *mut WakerNode,
+}
+
+unsafe impl Send for WakerList {}
+unsafe impl Sync for WakerList {}
+
+impl WakerList {
+    /// Create a new empty `WakerList`
+    pub fn new() -> Self {
+        Self {
+            head: std::ptr::null_mut(),
+        }
+    }
+
+    /// Insert a waker to the back of the list, and return its key.
+    pub fn insert(&mut self, waker: Option<Waker>) -> NonZeroUsize {
+        let node = Box::into_raw(Box::new(WakerNode {
+            waker,
+            next_in_queue: std::ptr::null_mut(),
+            prev_in_queue: std::ptr::null_mut(),
+        }));
+
+        if self.head.is_null() {
+            unsafe {
+                (*node).prev_in_queue = node;
+            }
+            self.head = node;
+        } else {
+            unsafe {
+                let prev = std::mem::replace(&mut (*self.head).prev_in_queue, node);
+                (*prev).next_in_queue = node;
+                (*node).prev_in_queue = prev;
+            }
+        }
+
+        unsafe { NonZeroUsize::new_unchecked(node as usize) }
+    }
+
+    /// Remove a waker by its key.
+    ///
+    /// # Safety
+    /// This function is unsafe because there is no guarantee that key is the previously returned
+    /// key, and that the key is only removed once.
+    pub unsafe fn remove(&mut self, key: NonZeroUsize) -> Option<Waker> {
+        let node = key.get() as *mut WakerNode;
+        let prev = (*node).prev_in_queue;
+        let next = (*node).next_in_queue;
+
+        // Special treatment on removing first node
+        if self.head == node {
+            self.head = next;
+        } else {
+            std::mem::replace(&mut (*prev).next_in_queue, next);
+        }
+
+        // Special treatment on removing last node
+        if next.is_null() {
+            if !self.head.is_null() {
+                std::mem::replace(&mut (*self.head).prev_in_queue, prev);
+            }
+        } else {
+            std::mem::replace(&mut (*next).prev_in_queue, prev);
+        }
+
+        Box::from_raw(node).waker
+    }
+
+    /// Get a waker by its key.
+    ///
+    /// # Safety
+    /// This function is unsafe because there is no guarantee that key is the previously returned
+    /// key, and that the key is not removed.
+    pub unsafe fn get(&mut self, key: NonZeroUsize) -> &mut Option<Waker> {
+        &mut (*(key.get() as *mut WakerNode)).waker
+    }
+
+    /// Check if this list is empty.
+    pub fn is_empty(&self) -> bool {
+        self.head.is_null()
+    }
+
+    /// Get an iterator over all wakers.
+    pub fn iter_mut(&mut self) -> Iter<'_> {
+        Iter {
+            ptr: self.head,
+            _marker: std::marker::PhantomData,
+        }
+    }
+
+    /// Wake the first waker in the list, and convert it to `None`. This function is named `weak` as
+    /// nothing is performed when the first waker is waken already.
+    pub fn wake_one_weak(&mut self) {
+        if let Some(opt_waker) = self.iter_mut().next() {
+            if let Some(w) = opt_waker.take() {
+                w.wake();
+            }
+        }
+    }
+}
+
+pub struct Iter<'a> {
+    ptr: *mut WakerNode,
+    _marker: std::marker::PhantomData<&'a ()>,
+}
+
+impl<'a> Iterator for Iter<'a> {
+    type Item = &'a mut Option<Waker>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.ptr.is_null() {
+            return None;
+        }
+        let next = unsafe { (*self.ptr).next_in_queue };
+        let ptr = std::mem::replace(&mut self.ptr, next);
+        Some(unsafe { &mut (*ptr).waker })
+    }
+}