refactor(driver): make the flag more meaningful

Author: Berrysoft

compio-driver/src/sys/driver/iocp/mod.rs

@@ -186,14 +186,16 @@ impl Driver {
             entry.notify();
             has_entry = true;
         }
+        if self.notify.reset() {
+            has_entry = true;
+        }
 
         if !has_entry {
             for e in self.notify.port.poll(timeout)? {
-                self.notify.set_awake(true);
                 if let Some(e) = Self::create_entry(notify, &mut self.waits, e) {
+                    self.notifier.set_awake();
                     e.notify()
                 }
-                self.notify.set_awake(false);
             }
         }
 
@@ -231,8 +233,12 @@ impl Notify {
         }
     }
 
-    fn set_awake(&self, awake: bool) {
-        self.awake.set(awake);
+    fn set_awake(&self) {
+        self.awake.set();
+    }
+
+    fn reset(&self) -> bool {
+        self.awake.reset()
     }
 }
 

compio-driver/src/sys/driver/iour/mod.rs

@@ -149,12 +149,12 @@ impl Driver {
     }
 
     // Auto means that it choose to wait or not automatically.
-    fn submit_auto(&mut self, timeout: Option<Duration>) -> io::Result<()> {
+    fn submit_auto(&mut self, timeout: Option<Duration>, need_wait: bool) -> io::Result<()> {
         instrument!(compio_log::Level::TRACE, "submit_auto", ?timeout);
 
         // when taskrun is true, there are completed cqes wait to handle, no need to
         // block the submit
-        let want_sqe = if self.inner.submission().taskrun() {
+        let want_sqe = if !need_wait || self.inner.submission().taskrun() {
             0
         } else {
             1
@@ -196,14 +196,10 @@ impl Driver {
         has_entry
     }
 
-    fn poll_entries(&mut self, hot_path: bool) -> bool {
+    fn poll_entries(&mut self) -> bool {
         let mut cqueue = self.inner.completion();
         cqueue.sync();
         let has_entry = !cqueue.is_empty();
-        // TODO: likely hint
-        if hot_path {
-            self.notifier.set_awake(true);
-        }
         for entry in cqueue {
             match entry.user_data() {
                 Self::CANCEL => {}
@@ -235,9 +231,6 @@ impl Driver {
                 }
             }
         }
-        if hot_path {
-            self.notifier.set_awake(false);
-        }
         has_entry
     }
 
@@ -289,7 +282,7 @@ impl Driver {
                 }
                 Err(_) => {
                     drop(squeue);
-                    match self.submit_auto(Some(Duration::ZERO)) {
+                    match self.submit_auto(Some(Duration::ZERO), true) {
                         Ok(()) => {}
                         Err(e)
                             if matches!(
@@ -303,7 +296,7 @@ impl Driver {
                     // event indefinitely.
                     //
                     // Anyway it is not a hot path, so we can afford an extra `write` syscall here.
-                    self.poll_entries(false);
+                    self.poll_entries();
                 }
             }
         }
@@ -372,6 +365,8 @@ impl Driver {
 
         trace!("start polling");
 
+        let need_wait = !self.notifier.reset();
+
         if self.need_push_notifier {
             #[allow(clippy::useless_conversion)]
             self.push_raw(
@@ -384,8 +379,10 @@ impl Driver {
             self.need_push_notifier = false;
         }
 
-        self.submit_auto(timeout)?;
-        self.poll_entries(true);
+        self.submit_auto(timeout, need_wait)?;
+
+        self.notifier.set_awake();
+        self.poll_entries();
 
         Ok(())
     }

compio-driver/src/sys/driver/iour/notify.rs

@@ -30,8 +30,12 @@ impl Notifier {
         Ok(())
     }
 
-    pub fn set_awake(&self, awake: bool) {
-        self.notify.set_awake(awake);
+    pub fn set_awake(&self) {
+        self.notify.set_awake();
+    }
+
+    pub fn reset(&self) -> bool {
+        self.notify.reset()
     }
 
     pub fn waker(&self) -> Waker {
@@ -66,8 +70,12 @@ impl Notify {
         }
     }
 
-    pub fn set_awake(&self, awake: bool) {
-        self.awake.set(awake);
+    pub fn set_awake(&self) {
+        self.awake.set();
+    }
+
+    pub fn reset(&self) -> bool {
+        self.awake.reset()
     }
 }
 

compio-driver/src/sys/driver/mod.rs

@@ -1,4 +1,4 @@
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::atomic::{AtomicU8, Ordering};
 
 cfg_if::cfg_if! {
     if #[cfg(windows)] {
@@ -24,19 +24,33 @@ cfg_if::cfg_if! {
 crate::assert_not_impl!(Driver, Send);
 crate::assert_not_impl!(Driver, Sync);
 
+const IDLE: u8 = 0b00;
+const NOTIFIED: u8 = 0b01;
+const AWAKE: u8 = 0b10;
+
 #[derive(Debug)]
-struct AwakeFlag(AtomicBool);
+struct AwakeFlag(AtomicU8);
 
 impl AwakeFlag {
     pub fn new() -> Self {
-        Self(AtomicBool::new(false))
+        Self(AtomicU8::new(IDLE))
+    }
+
+    /// Set the awake flag. It is true before the driver sleeps, and false after
+    /// it wakes up.
+    pub fn set(&self) {
+        self.0.fetch_or(AWAKE, Ordering::SeqCst);
     }
 
-    pub fn set(&self, awake: bool) {
-        self.0.store(awake, Ordering::SeqCst);
+    /// Reset the flags. Returns true if it was notified.
+    pub fn reset(&self) -> bool {
+        (self.0.swap(IDLE, Ordering::SeqCst) & NOTIFIED) != 0
     }
 
+    /// Set the notified flag. Returns true if the awake flag is set or the
+    /// notified flag is set. If the awake flag is not set, the driver needs
+    /// to be notified through a syscall.
     pub fn wake(&self) -> bool {
-        self.0.swap(true, Ordering::SeqCst)
+        self.0.fetch_or(NOTIFIED, Ordering::SeqCst) != 0
     }
 }

compio-driver/src/sys/driver/poll/mod.rs

@@ -172,10 +172,8 @@ impl Driver {
         F: FnOnce(&mut Self, &mut Events) -> R,
     {
         let mut events = std::mem::take(&mut self.events);
-        self.notify.set_awake(true);
         let res = f(self, &mut events);
         self.events = events;
-        self.notify.set_awake(false);
         res
     }
 
@@ -440,20 +438,22 @@ impl Driver {
 
     pub fn poll(&mut self, mut timeout: Option<Duration>) -> io::Result<()> {
         instrument!(compio_log::Level::TRACE, "poll", ?timeout);
+        let timeout_is_some = timeout.is_some();
         let has_completed = !self.completed_rx.is_empty();
-        if has_completed {
+        let need_wait = !self.notify.reset();
+        if !need_wait || has_completed {
             timeout = Some(Duration::ZERO);
         }
         // We need to poll the poller first to make sure it handles the internal notify
         // event (if any).
         self.events.clear();
         self.notify.poll.wait(&mut self.events, timeout)?;
-        self.notify.set_awake(false);
+        self.notify.set_awake();
         if self.events.is_empty() {
             if self.poll_completed() {
                 return Ok(());
             }
-            if timeout.is_some() {
+            if timeout_is_some {
                 return Err(io::Error::from_raw_os_error(libc::ETIMEDOUT));
             }
         } else if has_completed {
@@ -551,8 +551,12 @@ impl Notify {
         }
     }
 
-    fn set_awake(&self, awake: bool) {
-        self.awake.set(awake);
+    fn set_awake(&self) {
+        self.awake.set();
+    }
+
+    fn reset(&self) -> bool {
+        self.awake.reset()
     }
 }