Make waking more robust

This adds a new type, PollingState, that replaces the old is polling
boolean.

The old boolean was inherently racy because of the timing when the
boolean was set, the waker checked the value and the actual polling
began.

PollingState keep tracks of previous calls to wake and lets the polling
thread know if it's been awoken before the call to poll. If this is the
case we'll only poll with a zero timeout to ensure it's quick to return,
but it will submit all submissions and process any already ready
completions.

Author: Thomasdezeeuw

src/io_uring/cq.rs

@@ -61,9 +61,14 @@ impl Completions {
             // If we have no completions we make a system call to wait for
             // completion events.
             log::trace!(timeout:?; "waiting for completion events");
-            shared.is_polling.store(true, Ordering::Release);
+            let timeout = if shared.polling.set_polling(true) {
+                // Got woken up, so polling without a timeout.
+                Some(Duration::ZERO)
+            } else {
+                timeout
+            };
             let result = shared.enter(1, libc::IORING_ENTER_GETEVENTS, timeout);
-            shared.is_polling.store(false, Ordering::Release);
+            shared.polling.set_polling(false);
             result?;
             tail = load_kernel_shared(self.entries_tail);
         }

src/io_uring/mod.rs

@@ -7,11 +7,11 @@ use std::cmp::min;
 use std::mem::{drop as unlock, swap, take};
 use std::os::fd::{AsRawFd, OwnedFd, RawFd};
 use std::sync::Mutex;
-use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::time::Duration;
 use std::{ptr, task};
 
-use crate::{asan, lock, syscall, try_lock};
+use crate::{PollingState, asan, lock, syscall, try_lock};
 
 pub(crate) mod config;
 pub(crate) mod cq;
@@ -80,10 +80,7 @@ pub(crate) struct Shared {
     /// True if only a single thread can submit submissions, i.e. if
     /// `IORING_SETUP_SINGLE_ISSUER` is enabled.
     single_issuer: bool,
-    /// Boolean indicating a thread is [`Ring::poll`]ing.
-    ///
-    /// [`Ring::poll`]: crate::Ring::poll
-    is_polling: AtomicBool,
+    polling: PollingState,
     /// Futures that are waiting for a slot in submissions.
     blocked_futures: Mutex<Vec<task::Waker>>,
     /// File descriptor of the io_uring.
@@ -135,7 +132,7 @@ impl Shared {
             submissions_len: parameters.sq_entries,
             kernel_thread: (parameters.flags & libc::IORING_SETUP_SQPOLL) != 0,
             single_issuer: (parameters.flags & libc::IORING_SETUP_SINGLE_ISSUER) != 0,
-            is_polling: AtomicBool::new(false),
+            polling: PollingState::new(),
             blocked_futures: Mutex::new(Vec::new()),
             rfd,
         })

src/io_uring/sq.rs

@@ -89,7 +89,7 @@ impl Submissions {
 
     pub(crate) fn wake(&self) -> io::Result<()> {
         log::trace!("waking up ring");
-        if !self.shared.is_polling.load(Ordering::Acquire) {
+        if !self.shared.polling.wake() {
             // If we're not polling we don't need to wake up.
             log::trace!("skipping ring message as it's not polling");
             return Ok(());

src/kqueue/config.rs

@@ -3,11 +3,10 @@
 use std::marker::PhantomData;
 use std::os::fd::{AsRawFd, FromRawFd, OwnedFd};
 use std::sync::Mutex;
-use std::sync::atomic::AtomicBool;
 use std::{io, mem, ptr};
 
 use crate::kqueue::{Completions, Shared, Submissions, cq};
-use crate::syscall;
+use crate::{PollingState, syscall};
 
 #[derive(Debug, Clone)]
 pub(crate) struct Config<'r> {
@@ -86,7 +85,7 @@ impl<'r> crate::Config<'r> {
         let shared = Shared {
             max_change_list_size: self.sys.max_change_list_size,
             change_list: Mutex::new(Vec::new()),
-            is_polling: AtomicBool::new(false),
+            polling: PollingState::new(),
             kq,
         };
         let submissions = Submissions::new(shared);

src/kqueue/cq.rs

@@ -1,5 +1,4 @@
 use std::mem::{drop as unlock, take};
-use std::sync::atomic::Ordering;
 use std::time::Duration;
 use std::{cmp, io, ptr};
 
@@ -24,6 +23,13 @@ impl Completions {
     pub(crate) fn poll(&mut self, shared: &Shared, timeout: Option<Duration>) -> io::Result<()> {
         self.events.clear();
 
+        let timeout = if shared.polling.set_polling(true) {
+            // Got woken up, so polling without a timeout.
+            Some(Duration::ZERO)
+        } else {
+            timeout
+        };
+
         let ts = timeout.map(|to| libc::timespec {
             tv_sec: cmp::min(to.as_secs(), libc::time_t::MAX as u64).cast_signed(),
             // `Duration::subsec_nanos` is guaranteed to be less than one
@@ -43,9 +49,8 @@ impl Completions {
         unlock(change_list); // Unlock, to not block others.
 
         log::trace!(submissions = changes.len(), timeout:?; "waiting for events");
-        shared.is_polling.store(true, Ordering::Release);
         shared.kevent(&mut changes, UseEvents::Some(&mut self.events), ts.as_ref());
-        shared.is_polling.store(false, Ordering::Release);
+        shared.polling.set_polling(false);
         shared.reuse_change_list(changes);
         Ok(())
     }

src/kqueue/mod.rs

@@ -9,10 +9,9 @@
 use std::mem::{drop as unlock, swap};
 use std::os::fd::{AsRawFd, OwnedFd};
 use std::sync::Mutex;
-use std::sync::atomic::AtomicBool;
 use std::{fmt, ptr, task};
 
-use crate::{debug_detail, lock, syscall};
+use crate::{PollingState, debug_detail, lock, syscall};
 
 pub(crate) mod config;
 mod cq;
@@ -40,8 +39,7 @@ pub(crate) struct Shared {
     max_change_list_size: u32,
     /// Batched events to register.
     change_list: Mutex<Vec<Event>>,
-    /// Boolean indicating a thread is [`Ring::poll`]ing.
-    is_polling: AtomicBool,
+    polling: PollingState,
     /// kqueue(2) file descriptor.
     kq: OwnedFd,
 }

src/kqueue/sq.rs

@@ -2,7 +2,6 @@ use std::io;
 use std::mem::{self, drop as unlock, take};
 use std::os::fd::{AsRawFd, RawFd};
 use std::sync::Arc;
-use std::sync::atomic::Ordering;
 
 use crate::kqueue::{Event, Shared, UseEvents, cq};
 use crate::lock;
@@ -83,7 +82,7 @@ impl Submissions {
 
     #[allow(clippy::unnecessary_wraps)]
     pub(crate) fn wake(&self) -> io::Result<()> {
-        if !self.shared.is_polling.load(Ordering::Acquire) {
+        if !self.shared.polling.wake() {
             // If we're not polling we don't need to wake up.
             return Ok(());
         }

src/lib.rs

@@ -71,6 +71,8 @@
 )))]
 compile_error!("OS not supported");
 
+use std::fmt;
+use std::sync::atomic::{AtomicU8, Ordering};
 use std::time::Duration;
 
 // This must come before the other modules for the documentation.
@@ -506,3 +508,49 @@ impl<T> OpPollResult<T> for Option<io::Result<T>> {
         None
     }
 }
+
+/// Polling state of a [`Ring`].
+pub(crate) struct PollingState(AtomicU8);
+
+const IS_POLLING: u8 = 0b01;
+#[allow(unused)] // Used in an assert below, but rustc doesn't seem to care.
+const NOT_POLLING: u8 = 0b00;
+const IS_AWOKEN: u8 = 0b10;
+const NOT_AWOKEN: u8 = 0b00;
+
+impl PollingState {
+    pub(crate) const fn new() -> PollingState {
+        PollingState(AtomicU8::new(0))
+    }
+
+    /// Set the state to currently (not) polling.
+    ///
+    /// Returns a boolean indicating if the poll call should continue like
+    /// normal, or should be cut short (e.g. with a zero timeout).
+    #[allow(clippy::cast_lossless)]
+    pub(crate) fn set_polling(&self, is_polling: bool) -> bool {
+        const _BOOL_CAST_CHECK_TRUE: () = assert!(true as u8 == IS_POLLING);
+        const _BOOL_CAST_CHECK_FALSE: () = assert!(false as u8 == NOT_POLLING);
+        let state = self.0.swap(is_polling as u8 | NOT_AWOKEN, Ordering::AcqRel);
+        (state & IS_AWOKEN) != 0
+    }
+
+    /// Set the state to wake the polling thread.
+    ///
+    /// Returns a boolean indicating if the caller should submit an event to
+    /// wake up the polling thread.
+    pub(crate) fn wake(&self) -> bool {
+        let state = self.0.fetch_or(IS_AWOKEN, Ordering::AcqRel);
+        state == (IS_POLLING | NOT_AWOKEN)
+    }
+}
+
+impl fmt::Debug for PollingState {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let state = self.0.load(Ordering::Relaxed);
+        f.debug_struct("PollingState")
+            .field("polling", &((state & IS_POLLING) != 0))
+            .field("awoken", &((state & IS_AWOKEN) != 0))
+            .finish()
+    }
+}