Merge pull request async-rs#5 from Keruspe/optimistic-rw

skip optimistic read/write when there already are wakers registered

Author: Stjepan Glavina

src/lib.rs

@@ -518,10 +518,14 @@ impl<T> Async<T> {
     pub async fn read_with<R>(&self, op: impl FnMut(&T) -> io::Result<R>) -> io::Result<R> {
         let mut op = op;
         loop {
-            match op(self.get_ref()) {
-                Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
-                res => return res,
+            // If there are no blocked readers, attempt the read operation.
+            if !self.source.readers_registered() {
+                match op(self.get_ref()) {
+                    Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
+                    res => return res,
+                }
             }
+            // Wait until the I/O handle becomes readable.
             optimistic(self.readable()).await?;
         }
     }
@@ -554,10 +558,14 @@ impl<T> Async<T> {
     ) -> io::Result<R> {
         let mut op = op;
         loop {
-            match op(self.get_mut()) {
-                Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
-                res => return res,
+            // If there are no blocked readers, attempt the read operation.
+            if !self.source.readers_registered() {
+                match op(self.get_mut()) {
+                    Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
+                    res => return res,
+                }
             }
+            // Wait until the I/O handle becomes readable.
             optimistic(self.readable()).await?;
         }
     }
@@ -588,10 +596,14 @@ impl<T> Async<T> {
     pub async fn write_with<R>(&self, op: impl FnMut(&T) -> io::Result<R>) -> io::Result<R> {
         let mut op = op;
         loop {
-            match op(self.get_ref()) {
-                Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
-                res => return res,
+            // If there are no blocked readers, attempt the write operation.
+            if !self.source.writers_registered() {
+                match op(self.get_ref()) {
+                    Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
+                    res => return res,
+                }
             }
+            // Wait until the I/O handle becomes writable.
             optimistic(self.writable()).await?;
         }
     }
@@ -625,10 +637,14 @@ impl<T> Async<T> {
     ) -> io::Result<R> {
         let mut op = op;
         loop {
-            match op(self.get_mut()) {
-                Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
-                res => return res,
+            // If there are no blocked readers, attempt the write operation.
+            if !self.source.writers_registered() {
+                match op(self.get_mut()) {
+                    Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
+                    res => return res,
+                }
             }
+            // Wait until the I/O handle becomes writable.
             optimistic(self.writable()).await?;
         }
     }

src/reactor.rs

@@ -6,7 +6,7 @@ use std::os::unix::io::RawFd;
 #[cfg(windows)]
 use std::os::windows::io::RawSocket;
 use std::panic;
-use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex, MutexGuard};
 use std::task::{Poll, Waker};
 use std::thread;
@@ -166,6 +166,7 @@ impl Reactor {
                 readers: Vec::new(),
                 writers: Vec::new(),
             }),
+            wakers_registered: AtomicU8::new(0),
         });
         sources.insert(source.clone());
 
@@ -338,12 +339,18 @@ impl ReactorLock<'_> {
                         if ev.readable {
                             w.tick_readable = tick;
                             wakers.append(&mut w.readers);
+                            source
+                                .wakers_registered
+                                .fetch_and(!Source::READERS_REGISTERED, Ordering::SeqCst);
                         }
 
                         // Wake writers if a writability event was emitted.
                         if ev.writable {
                             w.tick_writable = tick;
                             wakers.append(&mut w.writers);
+                            source
+                                .wakers_registered
+                                .fetch_and(!Source::WRITERS_REGISTERED, Ordering::SeqCst);
                         }
 
                         // Re-register if there are still writers or
@@ -408,6 +415,9 @@ pub(crate) struct Source {
 
     /// Tasks interested in events on this source.
     wakers: Mutex<Wakers>,
+
+    /// Whether there are wakers interrested in events on this source.
+    wakers_registered: AtomicU8,
 }
 
 /// Tasks interested in events on a source.
@@ -427,6 +437,9 @@ struct Wakers {
 }
 
 impl Source {
+    const READERS_REGISTERED: u8 = 1 << 0;
+    const WRITERS_REGISTERED: u8 = 1 << 1;
+
     /// Waits until the I/O source is readable.
     pub(crate) async fn readable(&self) -> io::Result<()> {
         let mut ticks = None;
@@ -453,6 +466,8 @@ impl Source {
                         writable: !w.writers.is_empty(),
                     },
                 )?;
+                self.wakers_registered
+                    .fetch_or(Self::READERS_REGISTERED, Ordering::SeqCst);
             }
 
             // Register the current task's waker if not present already.
@@ -473,6 +488,11 @@ impl Source {
         .await
     }
 
+    pub(crate) fn readers_registered(&self) -> bool {
+        self.wakers_registered.load(Ordering::SeqCst) & Self::READERS_REGISTERED
+            == Self::READERS_REGISTERED
+    }
+
     /// Waits until the I/O source is writable.
     pub(crate) async fn writable(&self) -> io::Result<()> {
         let mut ticks = None;
@@ -499,6 +519,8 @@ impl Source {
                         writable: true,
                     },
                 )?;
+                self.wakers_registered
+                    .fetch_or(Self::WRITERS_REGISTERED, Ordering::SeqCst);
             }
 
             // Register the current task's waker if not present already.
@@ -518,4 +540,9 @@ impl Source {
         })
         .await
     }
+
+    pub(crate) fn writers_registered(&self) -> bool {
+        self.wakers_registered.load(Ordering::SeqCst) & Self::WRITERS_REGISTERED
+            == Self::WRITERS_REGISTERED
+    }
 }