Merge pull request #196 from Berrysoft/atomic-waker

Author: George-Miao

compio-dispatcher/src/lib.rs

@@ -63,7 +63,7 @@ impl Dispatcher {
                             *f.result.lock().unwrap() = Some(std::panic::catch_unwind(|| {
                                 Runtime::current().block_on((f.func)());
                             }));
-                            f.handle.notify().ok();
+                            f.handle.notify();
                         }
                     })
                 }
@@ -98,8 +98,8 @@ impl Dispatcher {
         &self,
         f: Fn,
     ) -> io::Result<DispatcherJoinHandle> {
-        let event = Event::new()?;
-        let handle = event.handle()?;
+        let event = Event::new();
+        let handle = event.handle();
         let join_handle = DispatcherJoinHandle::new(event);
         let closure = DispatcherClosure {
             handle,
@@ -117,8 +117,8 @@ impl Dispatcher {
     pub async fn join(self) -> io::Result<()> {
         drop(self.sender);
         let results = Arc::new(Mutex::new(vec![]));
-        let event = Event::new()?;
-        let handle = event.handle()?;
+        let event = Event::new();
+        let handle = event.handle();
         if let Err(f) = self.pool.dispatch({
             let results = results.clone();
             move || {
@@ -127,12 +127,12 @@ impl Dispatcher {
                     .into_iter()
                     .map(|thread| thread.join())
                     .collect();
-                handle.notify().ok();
+                handle.notify();
             }
         }) {
             std::thread::spawn(f);
         }
-        event.wait().await?;
+        event.wait().await;
         let mut guard = results.lock().unwrap();
         for res in std::mem::take::<Vec<std::thread::Result<()>>>(guard.as_mut()) {
             // The thread should not panic.
@@ -223,7 +223,7 @@ impl DispatcherJoinHandle {
 
     /// Wait for the task to complete.
     pub async fn join(self) -> io::Result<std::thread::Result<()>> {
-        self.event.wait().await?;
+        self.event.wait().await;
         Ok(self
             .result
             .lock()

compio-net/src/resolve/windows.rs

@@ -47,7 +47,7 @@ impl AsyncResolver {
         let overlapped_ptr = lpoverlapped.cast::<GAIOverlapped>().cast_mut();
         if let Some(overlapped) = overlapped_ptr.as_mut() {
             if let Some(handle) = overlapped.handle.take() {
-                handle.notify().ok();
+                handle.notify();
             }
         }
     }
@@ -151,14 +151,14 @@ pub async fn resolve_sock_addrs(
     hints.ai_socktype = SOCK_STREAM;
     hints.ai_protocol = IPPROTO_TCP;
 
-    let event = Event::new()?;
-    let handle = event.handle()?;
+    let event = Event::new();
+    let handle = event.handle();
     match unsafe { resolver.call(&hints, handle) } {
         Poll::Ready(res) => {
             res?;
         }
         Poll::Pending => {
-            event.wait().await?;
+            event.wait().await;
         }
     }
 

compio-runtime/src/event/eventfd.rs

@@ -1,21 +1,111 @@
 //! Asynchronous events.
-//!
-//! Only for waking up the driver.
-
-cfg_if::cfg_if! {
-    if #[cfg(windows)] {
-        mod iocp;
-        pub use iocp::*;
-    } else if #[cfg(any(
-        target_os = "android",
-        target_os = "freebsd",
-        target_os = "illumos",
-        target_os = "linux",
-    ))] {
-        mod eventfd;
-        pub use eventfd::*;
-    } else if #[cfg(unix)] {
-        mod pipe;
-        pub use pipe::*;
+
+use std::{
+    pin::Pin,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    task::{Context, Poll},
+};
+
+use futures_util::{task::AtomicWaker, Future};
+
+#[derive(Debug)]
+struct Inner {
+    waker: AtomicWaker,
+    set: AtomicBool,
+}
+
+#[derive(Debug, Clone)]
+struct Flag(Arc<Inner>);
+
+impl Flag {
+    pub fn new() -> Self {
+        Self(Arc::new(Inner {
+            waker: AtomicWaker::new(),
+            set: AtomicBool::new(false),
+        }))
+    }
+
+    pub fn notify(&self) {
+        self.0.set.store(true, Ordering::Relaxed);
+        self.0.waker.wake();
+    }
+
+    pub fn notified(&self) -> bool {
+        self.0.set.load(Ordering::Relaxed)
+    }
+}
+
+impl Future for Flag {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+        // quick check to avoid registration if already done.
+        if self.0.set.load(Ordering::Relaxed) {
+            return Poll::Ready(());
+        }
+
+        self.0.waker.register(cx.waker());
+
+        // Need to check condition **after** `register` to avoid a race
+        // condition that would result in lost notifications.
+        if self.0.set.load(Ordering::Relaxed) {
+            Poll::Ready(())
+        } else {
+            Poll::Pending
+        }
+    }
+}
+
+/// An event that won't wake until [`EventHandle::notify`] is called
+/// successfully.
+#[derive(Debug)]
+pub struct Event {
+    flag: Flag,
+}
+
+impl Default for Event {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Event {
+    /// Create [`Event`].
+    pub fn new() -> Self {
+        Self { flag: Flag::new() }
+    }
+
+    /// Get a notify handle.
+    pub fn handle(&self) -> EventHandle {
+        EventHandle::new(self.flag.clone())
+    }
+
+    /// Get if the event has been notified.
+    pub fn notified(&self) -> bool {
+        self.flag.notified()
+    }
+
+    /// Wait for [`EventHandle::notify`] called.
+    pub async fn wait(self) {
+        self.flag.await
+    }
+}
+
+/// A wake up handle to [`Event`].
+pub struct EventHandle {
+    flag: Flag,
+}
+
+impl EventHandle {
+    fn new(flag: Flag) -> Self {
+        Self { flag }
+    }
+
+    /// Notify the event.
+    pub fn notify(self) {
+        self.flag.notify()
     }
 }