fix: fix AtomicWaker implementation

Author: wyfo

.github/workflows/ci.yml

@@ -11,13 +11,17 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
-      - uses: actions-rs/toolchain@v1
-        with:
-          toolchain: nightly
-          override: true
-          components: rustfmt, clippy, miri
-      - run: cargo +nightly fmt --check
-      - run: cargo clippy --no-default-features -- -D warnings
-      - run: cargo clippy --all-features --tests -- -D warnings
-      - run: cargo test --all-features
-      - run: cargo +nightly miri test --all-features
+      - name: rustfmt
+        run: cargo fmt -- --config "unstable_features=true,imports_granularity=Crate,group_imports=StdExternalCrate,format_code_in_doc_comments=true"
+      - name: clippy no-default-feature
+        run: cargo clippy --no-default-features -- -D warnings
+      - name: clippy all-features
+        run: cargo clippy --all-features --tests -- -D warnings
+      - name: test
+        run: cargo test --all-features
+      - name: miri
+        run: cargo +nightly miri test --all-features --all-seeds
+      - name: loom
+        run: cargo test --release --lib
+        env:
+          RUSTFLAGS: "--cfg loom"

Cargo.toml

@@ -36,9 +36,12 @@ tokio-test = "0.4"
 tokio = { version = "1", features = ["macros", "rt-multi-thread", "test-util", "time"] }
 
 [target.'cfg(loom)'.dev-dependencies]
-loom = { version = "0.6.0", features = ["futures"] }
+loom = { version = "0.7", features = ["futures"] }
 
-# See tokio Cargo.toml
+[lints.rust]
+unexpected_cfgs = { level = "warn", check-cfg = ['cfg(loom)'] }
+
+# see https://users.rust-lang.org/t/how-to-document-optional-features-in-api-docs/64577/3
 [package.metadata.docs.rs]
 all-features = true
 rustdoc-args = [

src/loom.rs

@@ -1,8 +1,10 @@
 #[cfg(not(all(loom, test)))]
 mod without_loom {
+    #[cfg(not(feature = "std"))]
+    pub(crate) use core::sync;
+    pub(crate) use core::{cell, hint};
     #[cfg(feature = "std")]
-    pub(crate) use std::thread;
-    pub(crate) use std::{cell, sync};
+    pub(crate) use std::{sync, thread};
 
     pub(crate) const SPIN_LIMIT: usize = 64;
     pub(crate) const BACKOFF_LIMIT: usize = 6;
@@ -11,6 +13,10 @@ mod without_loom {
     pub(crate) struct LoomUnsafeCell<T>(cell::UnsafeCell<T>);
 
     impl<T> LoomUnsafeCell<T> {
+        pub(crate) const fn new(data: T) -> Self {
+            Self(cell::UnsafeCell::new(data))
+        }
+
         pub(crate) fn with<R>(&self, f: impl FnOnce(*const T) -> R) -> R {
             f(self.0.get())
         }
@@ -28,7 +34,7 @@ pub(crate) use without_loom::*;
 mod with_loom {
     #[cfg(feature = "std")]
     pub(crate) use loom::thread;
-    pub(crate) use loom::{cell, sync};
+    pub(crate) use loom::{cell, hint, sync};
 
     pub(crate) const SPIN_LIMIT: usize = 1;
     pub(crate) const BACKOFF_LIMIT: usize = 1;

src/queue.rs

@@ -6,6 +6,7 @@ use crate::{
     buffer::{Buffer, BufferSlice, Drain, InsertIntoBuffer, Resize},
     error::{TryDequeueError, TryEnqueueError},
     loom::{
+        hint,
         sync::atomic::{AtomicUsize, Ordering},
         LoomUnsafeCell, BACKOFF_LIMIT, SPIN_LIMIT,
     },

src/synchronized.rs

@@ -9,7 +9,7 @@ use std::{
 use crate::{
     buffer::{Buffer, BufferSlice, Drain, InsertIntoBuffer},
     error::{DequeueError, EnqueueError, TryDequeueError, TryEnqueueError},
-    loom::{thread, SPIN_LIMIT},
+    loom::{hint, thread, SPIN_LIMIT},
     notify::Notify,
     synchronized::{atomic_waker::AtomicWaker, waker_list::WakerList},
     Queue,
@@ -429,7 +429,7 @@ where
             }
             res => return Ok(res),
         };
-        std::hint::spin_loop();
+        hint::spin_loop();
     }
     queue.notify().enqueuers.register(cx);
     match queue.try_enqueue(value) {
@@ -450,7 +450,7 @@ where
             Err(TryDequeueError::Empty | TryDequeueError::Pending) => {}
             res => return Some(res),
         }
-        std::hint::spin_loop();
+        hint::spin_loop();
     }
     queue.notify().dequeuer.register(cx);
     match queue.try_dequeue() {

src/synchronized/atomic_waker.rs

@@ -1,103 +1,159 @@
-use std::{mem::MaybeUninit, task};
+use std::task;
 
 use crate::{
     loom::{
-        cell::Cell,
-        sync::atomic::{AtomicU8, Ordering},
-        thread,
+        hint,
+        sync::atomic::{AtomicUsize, Ordering},
+        thread, LoomUnsafeCell,
     },
     synchronized::waker::Waker,
 };
 
-const EMPTY: u8 = 0b000;
-const REGISTERING: u8 = 0b001;
-const REGISTERED: u8 = 0b010;
-const WAKING_FLAG: u8 = 0b100;
-
+// I have to reimplement AtomicWaker because it doesn't use standard `Waker`,
+// see https://internals.rust-lang.org/t/thread-park-waker-a-waker-calling-thread-unpark/19114
 pub(super) struct AtomicWaker {
-    waker: Cell<MaybeUninit<Waker>>,
-    state: AtomicU8,
+    state: AtomicUsize,
+    waker: LoomUnsafeCell<Option<Waker>>,
 }
 
-impl Default for AtomicWaker {
-    fn default() -> Self {
+const WAITING: usize = 0;
+const REGISTERING: usize = 0b01;
+const WAKING: usize = 0b10;
+
+impl AtomicWaker {
+    pub fn new() -> Self {
         Self {
-            waker: Cell::new(MaybeUninit::uninit()),
-            state: AtomicU8::new(0),
+            state: AtomicUsize::new(WAITING),
+            waker: LoomUnsafeCell::new(None),
         }
     }
-}
-
-// SAFETY: `AtomicWaker` waker `Cell` access is synchronized using state
-// (see `AtomicWaker::register`/`AtomicWaker::wake`)
-unsafe impl Send for AtomicWaker {}
-
-// SAFETY: `AtomicWaker` waker `Cell` access is synchronized using state
-// (see `AtomicWaker::register`/`AtomicWaker::wake`)
-unsafe impl Sync for AtomicWaker {}
 
-impl AtomicWaker {
     #[inline]
     pub(super) fn register(&self, cx: Option<&task::Context>) {
-        let mut state = self.state.load(Ordering::Relaxed);
-        loop {
-            if state != EMPTY && state != REGISTERED {
+        match self
+            .state
+            .compare_exchange(WAITING, REGISTERING, Ordering::Acquire, Ordering::Acquire)
+            .unwrap_or_else(|x| x)
+        {
+            WAITING => {
+                // SAFETY: see `futures::task::AtomicWaker` implementation
+                unsafe {
+                    self.waker.with_mut(|w| match &mut *w {
+                        Some(old_waker) if old_waker.will_wake(cx) => (),
+                        _ => *w = Some(Waker::new(cx)),
+                    });
+                    let res = self.state.compare_exchange(
+                        REGISTERING,
+                        WAITING,
+                        Ordering::AcqRel,
+                        Ordering::Acquire,
+                    );
+
+                    match res {
+                        Ok(_) => {}
+                        Err(actual) => {
+                            debug_assert_eq!(actual, REGISTERING | WAKING);
+                            let waker = self.waker.with_mut(|w| (*w).take()).unwrap();
+                            self.state.swap(WAITING, Ordering::AcqRel);
+                            waker.wake();
+                        }
+                    }
+                }
+            }
+            WAKING => {
                 match cx {
                     Some(cx) => cx.waker().wake_by_ref(),
                     None => thread::current().unpark(),
                 }
-                return;
+                // SAFETY: see `AtomicWaker` implementation from tokio (it's needed by loom)
+                hint::spin_loop();
             }
-            match self.state.compare_exchange_weak(
-                state,
-                REGISTERING,
-                Ordering::AcqRel,
-                Ordering::Relaxed,
-            ) {
-                Ok(_) => break,
-                Err(s) => state = s,
+            state => {
+                debug_assert!(state == REGISTERING || state == REGISTERING | WAKING);
             }
         }
-        let prev = self.waker.replace(MaybeUninit::new(Waker::new(cx)));
-        if state == REGISTERED {
-            // SAFETY: state was previously `REGISTERED`, so waker was initialized.
-            unsafe { prev.assume_init() };
-        }
-        if let Err(state) = self.state.compare_exchange(
-            REGISTERING,
-            REGISTERED,
-            Ordering::AcqRel,
-            Ordering::Acquire,
-        ) {
-            debug_assert_eq!(state, REGISTERING | WAKING_FLAG);
-            // SAFETY: waker has been initialized a few lines above, and cannot be read
-            // by `Self::wake` because state has not been set to `REGISTERED`.
-            unsafe { self.waker.replace(MaybeUninit::uninit()).assume_init() }.wake();
-            self.state.store(EMPTY, Ordering::Release);
-        }
     }
 
     #[inline]
     pub(super) fn wake(&self) {
-        let mut state = self.state.load(Ordering::Relaxed);
-        loop {
-            if state != REGISTERING && state != REGISTERED {
-                return;
+        match self.state.fetch_or(WAKING, Ordering::AcqRel) {
+            WAITING => {
+                // SAFETY: see `futures::task::AtomicWaker` implementation
+                let waker = unsafe { self.waker.with_mut(|w| (*w).take()) };
+                self.state.fetch_and(!WAKING, Ordering::Release);
+                if let Some(waker) = waker {
+                    waker.wake()
+                }
             }
-            match self.state.compare_exchange_weak(
-                state,
-                state | WAKING_FLAG,
-                Ordering::AcqRel,
-                Ordering::Relaxed,
-            ) {
-                Ok(REGISTERED) => break,
-                Ok(_) => return,
-                Err(s) => state = s,
+            state => {
+                debug_assert!(
+                    state == REGISTERING || state == REGISTERING | WAKING || state == WAKING
+                );
             }
         }
-        // SAFETY: state was `REGISTERED`, so waker has been registered and can no more
-        // be read/modified by `Self::register`
-        unsafe { self.waker.replace(MaybeUninit::uninit()).assume_init() }.wake();
-        self.state.store(EMPTY, Ordering::Release);
+    }
+}
+
+impl Default for AtomicWaker {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+// SAFETY: see `futures::task::AtomicWaker` implementation
+unsafe impl Send for AtomicWaker {}
+
+// SAFETY: see `futures::task::AtomicWaker` implementation
+unsafe impl Sync for AtomicWaker {}
+
+#[cfg(all(test, loom))]
+mod tests {
+    use std::{
+        future::poll_fn,
+        sync::Arc,
+        task::Poll::{Pending, Ready},
+    };
+
+    use loom::{
+        future::block_on,
+        sync::atomic::{AtomicUsize, Ordering},
+        thread,
+    };
+
+    use super::AtomicWaker;
+
+    struct Chan {
+        num: AtomicUsize,
+        task: AtomicWaker,
+    }
+    #[test]
+    fn basic_notification() {
+        const NUM_NOTIFY: usize = 2;
+
+        loom::model(|| {
+            let chan = Arc::new(Chan {
+                num: AtomicUsize::new(0),
+                task: AtomicWaker::default(),
+            });
+
+            for _ in 0..NUM_NOTIFY {
+                let chan = chan.clone();
+
+                thread::spawn(move || {
+                    chan.num.fetch_add(1, Ordering::Relaxed);
+                    chan.task.wake();
+                });
+            }
+
+            block_on(poll_fn(move |cx| {
+                chan.task.register(Some(cx));
+
+                if NUM_NOTIFY == chan.num.load(Ordering::Relaxed) {
+                    return Ready(());
+                }
+
+                Pending
+            }));
+        });
     }
 }

src/synchronized/waker.rs

@@ -9,13 +9,23 @@ pub(super) enum Waker {
 }
 
 impl Waker {
+    #[inline]
     pub(super) fn new(cx: Option<&task::Context>) -> Self {
         match cx {
             Some(cx) => Self::Async(cx.waker().clone()),
             None => Self::Sync(thread::current()),
         }
     }
 
+    #[inline]
+    pub(super) fn will_wake(&self, cx: Option<&task::Context>) -> bool {
+        match (self, cx) {
+            (Self::Async(w), Some(cx)) => w.will_wake(cx.waker()),
+            _ => false,
+        }
+    }
+
+    #[inline]
     pub(super) fn wake(self) {
         match self {
             Self::Async(waker) => waker.wake(),