impl WakerCell to replace ArrayQueue(size=1)

for faster mpsc/spsc

crossfire_bounded_100_async_n_1
mpsc/1 20.298 Melem/s +60%
mpsc/2 17.767 Melem/s +67%

Author: frostyplanet

src/collections.rs

@@ -1,6 +1,58 @@
 use parking_lot::Mutex;
 use std::collections::VecDeque;
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::ptr;
+use std::sync::{
+    atomic::{AtomicBool, AtomicPtr, Ordering},
+    Arc, Weak,
+};
+
+pub struct WeakCell<T> {
+    ptr: AtomicPtr<T>,
+}
+
+unsafe impl<T> Send for WeakCell<T> {}
+unsafe impl<T> Sync for WeakCell<T> {}
+
+impl<T> WeakCell<T> {
+    #[inline(always)]
+    pub fn new() -> Self {
+        Self { ptr: AtomicPtr::new(ptr::null_mut()) }
+    }
+
+    #[inline(always)]
+    pub fn exists(&self) -> bool {
+        self.ptr.load(Ordering::Acquire) == ptr::null_mut()
+    }
+
+    #[inline(always)]
+    pub fn pop(&self) -> Option<Arc<T>> {
+        if self.ptr.load(Ordering::Acquire) == ptr::null_mut() {
+            return None;
+        }
+        let ptr = self.ptr.swap(ptr::null_mut(), Ordering::SeqCst);
+        if ptr != ptr::null_mut() {
+            return unsafe { Weak::from_raw(ptr) }.upgrade();
+        } else {
+            None
+        }
+    }
+
+    pub fn clear(&self) {
+        let ptr = self.ptr.swap(ptr::null_mut(), Ordering::SeqCst);
+        if ptr != ptr::null_mut() {
+            // Convert into Weak and drop
+            let _ = unsafe { Weak::from_raw(ptr) };
+        }
+    }
+
+    #[inline(always)]
+    pub fn put(&self, item: Weak<T>) {
+        let old_ptr = self.ptr.swap(item.into_raw() as *mut T, Ordering::SeqCst);
+        if old_ptr != ptr::null_mut() {
+            let _ = unsafe { Weak::from_raw(old_ptr) };
+        }
+    }
+}
 
 pub struct LockedQueue<T> {
     empty: AtomicBool,

src/locked_waker.rs

@@ -1,3 +1,4 @@
+use crate::collections::WeakCell;
 use std::fmt;
 use std::sync::{
     atomic::{AtomicBool, Ordering},
@@ -21,9 +22,7 @@ struct LockedWakerInner {
     seq: u64,
 }
 
-pub struct LockedWakerRef {
-    w: Weak<LockedWakerInner>,
-}
+pub struct LockedWakerRef(Weak<LockedWakerInner>);
 
 impl fmt::Debug for LockedWakerRef {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -61,7 +60,7 @@ impl LockedWaker {
 
     #[inline(always)]
     pub(crate) fn weak(&self) -> LockedWakerRef {
-        LockedWakerRef { w: Arc::downgrade(&self.0) }
+        LockedWakerRef(Arc::downgrade(&self.0))
     }
 
     #[inline(always)]
@@ -84,7 +83,7 @@ impl LockedWaker {
 impl LockedWakerRef {
     #[inline(always)]
     pub(crate) fn wake(&self) -> bool {
-        if let Some(_self) = self.w.upgrade() {
+        if let Some(_self) = self.0.upgrade() {
             return LockedWaker(_self).wake();
         } else {
             return false;
@@ -93,7 +92,7 @@ impl LockedWakerRef {
 
     /// return true to stop; return false to continue the search.
     pub(crate) fn try_to_clear(&self, seq: u64) -> bool {
-        if let Some(w) = self.w.upgrade() {
+        if let Some(w) = self.0.upgrade() {
             let waker = LockedWaker(w);
             let _seq = waker.get_seq();
             if _seq == seq {
@@ -111,17 +110,56 @@ impl LockedWakerRef {
         return false;
     }
 
+    #[allow(dead_code)]
     #[inline(always)]
     pub(crate) fn check_eq(&self, other: LockedWakerRef) -> bool {
-        if self.w.ptr_eq(&other.w) {
+        if self.0.ptr_eq(&other.0) {
             return true;
         }
         other.wake();
         false
     }
 }
 
-#[test]
-fn test_waker() {
-    println!("waker size {}", std::mem::size_of::<LockedWakerRef>());
+pub struct WakerCell(WeakCell<LockedWakerInner>);
+
+impl WakerCell {
+    #[inline(always)]
+    pub fn new() -> Self {
+        Self(WeakCell::new())
+    }
+
+    #[inline(always)]
+    pub fn wake(&self) -> bool {
+        if let Some(waker) = self.0.pop() {
+            return LockedWaker(waker).wake();
+        }
+        false
+    }
+
+    #[inline(always)]
+    pub fn clear(&self) {
+        self.0.clear();
+    }
+
+    #[inline(always)]
+    pub fn put(&self, waker: LockedWakerRef) {
+        self.0.put(waker.0);
+    }
+
+    #[inline(always)]
+    pub fn exists(&self) -> bool {
+        self.0.exists()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn test_waker() {
+        println!("waker size {}", std::mem::size_of::<LockedWakerRef>());
+        println!("arc size {}", std::mem::size_of::<Arc<WakerCell>>());
+        println!("arc size {}", std::mem::size_of::<Weak<WakerCell>>());
+    }
 }

src/waker_registry.rs

@@ -1,6 +1,5 @@
 use crate::collections::LockedQueue;
 use crate::locked_waker::*;
-use crossbeam::queue::ArrayQueue;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::task::Context;
 
@@ -64,13 +63,13 @@ impl RegistryTrait for RegistryDummy {
 }
 
 pub struct RegistrySingle {
-    queue: ArrayQueue<LockedWakerRef>,
+    cell: WakerCell,
 }
 
 impl RegistrySingle {
     #[inline(always)]
     pub fn new() -> Registry {
-        Registry::Single(Self { queue: ArrayQueue::new(1) })
+        Registry::Single(Self { cell: WakerCell::new() })
     }
 }
 
@@ -79,42 +78,25 @@ impl RegistryTrait for RegistrySingle {
     #[inline(always)]
     fn reg_async(&self, ctx: &mut Context) -> LockedWaker {
         let waker = LockedWaker::new(ctx, 0);
-        match self.queue.push(waker.weak()) {
-            Ok(_) => {}
-            Err(_weak) => {
-                if let Some(old_waker) = self.queue.pop() {
-                    _weak.check_eq(old_waker);
-                    self.queue.push(_weak).expect("Do not misuse mpsc as mpmc");
-                } else {
-                    self.queue.push(_weak).expect("Do not misuse mpsc as mpmc");
-                }
-            }
-        }
+        self.cell.put(waker.weak());
         waker
     }
 
     #[inline(always)]
-    fn cancel_waker(&self, waker: LockedWaker) {
+    fn cancel_waker(&self, _waker: LockedWaker) {
         // Got to be it, because only one single thread.
-        if let Some(waker_ref) = self.queue.pop() {
-            // protect miss-use of multi thread
-            waker.weak().check_eq(waker_ref);
-        }
+        self.cell.clear();
     }
 
     #[inline(always)]
     fn clear_wakers(&self, _seq: u64) {
         // Got to be it, because only one single thread.
-        if let Some(_waker) = self.queue.pop() {
-            _waker.wake();
-        }
+        self.cell.clear();
     }
 
     #[inline(always)]
     fn fire(&self) {
-        if let Some(waker) = self.queue.pop() {
-            waker.wake();
-        }
+        self.cell.wake();
     }
 
     #[inline(always)]
@@ -125,7 +107,11 @@ impl RegistryTrait for RegistrySingle {
     /// return waker queue size
     #[inline(always)]
     fn get_size(&self) -> usize {
-        self.queue.len()
+        if self.cell.exists() {
+            1
+        } else {
+            0
+        }
     }
 }