Merge pull request #215 from Berrysoft/dev/runtime-hook

George-Miao · web-flow · commit e1bcab57cff3 · 2024-03-04T15:31:21.000+08:00
feat(runtime): a runtime hook API design
diff --git a/compio-driver/Cargo.toml b/compio-driver/Cargo.toml
@@ -61,11 +61,13 @@ windows-sys = { workspace = true, features = [
 [target.'cfg(target_os = "linux")'.dependencies]
 io-uring = { version = "0.6.2", optional = true }
 polling = { version = "3.3.0", optional = true }
+os_pipe = { workspace = true, optional = true }
 paste = { workspace = true }
 
 # Other platform dependencies
 [target.'cfg(all(not(target_os = "linux"), unix))'.dependencies]
 polling = "3.3.0"
+os_pipe = { workspace = true }
 
 [target.'cfg(unix)'.dependencies]
 crossbeam-queue = { workspace = true }
@@ -76,7 +78,7 @@ compio-buf = { workspace = true, features = ["arrayvec"] }
 
 [features]
 default = ["io-uring"]
-polling = ["dep:polling"]
+polling = ["dep:polling", "dep:os_pipe"]
 
 io-uring-sqe128 = []
 io-uring-cqe32 = []
diff --git a/compio-driver/src/iour/mod.rs b/compio-driver/src/iour/mod.rs
@@ -23,7 +23,7 @@ cfg_if::cfg_if! {
     }
 }
 use io_uring::{
-    opcode::{AsyncCancel, Read},
+    opcode::{AsyncCancel, PollAdd},
     types::{Fd, SubmitArgs, Timespec},
     IoUring,
 };
@@ -76,7 +76,6 @@ pub(crate) struct Driver {
     inner: IoUring<SEntry, CEntry>,
     squeue: VecDeque<SEntry>,
     notifier: Notifier,
-    notifier_registered: bool,
     pool: AsyncifyPool,
     pool_completed: Arc<SegQueue<Entry>>,
 }
@@ -88,11 +87,20 @@ impl Driver {
     pub fn new(builder: &ProactorBuilder) -> io::Result<Self> {
         instrument!(compio_log::Level::TRACE, "new", ?builder);
         trace!("new iour driver");
+        let mut squeue = VecDeque::with_capacity(builder.capacity as usize);
+        let notifier = Notifier::new()?;
+        #[allow(clippy::useless_conversion)]
+        squeue.push_back(
+            PollAdd::new(Fd(notifier.as_raw_fd()), libc::POLLIN as _)
+                .multi(true)
+                .build()
+                .user_data(Self::NOTIFY)
+                .into(),
+        );
         Ok(Self {
             inner: IoUring::builder().build(builder.capacity)?,
-            squeue: VecDeque::with_capacity(builder.capacity as usize),
-            notifier: Notifier::new()?,
-            notifier_registered: false,
+            squeue,
+            notifier,
             pool: builder.create_or_get_thread_pool(),
             pool_completed: Arc::new(SegQueue::new()),
         })
@@ -177,7 +185,10 @@ impl Driver {
         let completed_entries = cqueue.filter_map(|entry| match entry.user_data() {
             Self::CANCEL => None,
             Self::NOTIFY => {
-                self.notifier_registered = false;
+                const IORING_CQE_F_MORE: u32 = 1 << 1;
+                let flags = entry.flags();
+                debug_assert!(flags & IORING_CQE_F_MORE == IORING_CQE_F_MORE);
+                self.notifier.clear().expect("cannot clear notifier");
                 None
             }
             _ => Some(create_entry(entry)),
@@ -256,19 +267,6 @@ impl Driver {
         mut entries: OutEntries<impl Extend<usize>>,
     ) -> io::Result<()> {
         instrument!(compio_log::Level::TRACE, "poll", ?timeout);
-        if !self.notifier_registered {
-            let fd = self.notifier.as_raw_fd();
-            let dst = self.notifier.dst();
-            #[allow(clippy::useless_conversion)]
-            self.squeue.push_back(
-                Read::new(Fd(fd), dst.as_mut_ptr(), dst.len() as _)
-                    .build()
-                    .user_data(Self::NOTIFY)
-                    .into(),
-            );
-            trace!("registered notifier");
-            self.notifier_registered = true
-        }
         // Anyway we need to submit once, no matter there are entries in squeue.
         trace!("start polling");
         loop {
@@ -321,27 +319,40 @@ fn timespec(duration: std::time::Duration) -> Timespec {
 #[derive(Debug)]
 struct Notifier {
     fd: OwnedFd,
-    read_dst: Box<[u8; 8]>,
 }
 
 impl Notifier {
     /// Create a new notifier.
     fn new() -> io::Result<Self> {
-        let fd = syscall!(libc::eventfd(0, libc::EFD_CLOEXEC))?;
+        let fd = syscall!(libc::eventfd(0, libc::EFD_CLOEXEC | libc::EFD_NONBLOCK))?;
         let fd = unsafe { OwnedFd::from_raw_fd(fd) };
-        Ok(Self {
-            fd,
-            read_dst: Box::new([0; 8]),
-        })
+        Ok(Self { fd })
     }
 
-    fn dst(&mut self) -> &mut [u8] {
-        self.read_dst.as_mut_slice()
+    pub fn clear(&self) -> io::Result<()> {
+        loop {
+            let mut buffer = [0u64];
+            let res = syscall!(libc::read(
+                self.fd.as_raw_fd(),
+                buffer.as_mut_ptr().cast(),
+                std::mem::size_of::<u64>()
+            ));
+            match res {
+                Ok(len) => {
+                    debug_assert_eq!(len, std::mem::size_of::<u64>() as _);
+                    break Ok(());
+                }
+                // Clear the next time:)
+                Err(e) if e.kind() == io::ErrorKind::WouldBlock => break Ok(()),
+                // Just like read_exact
+                Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
+                Err(e) => break Err(e),
+            }
+        }
     }
 
     pub fn handle(&self) -> io::Result<NotifyHandle> {
-        let fd = self.fd.try_clone()?;
-        Ok(NotifyHandle::new(fd))
+        Ok(NotifyHandle::new(self.fd.try_clone()?))
     }
 }
 
diff --git a/compio-driver/src/poll/mod.rs b/compio-driver/src/poll/mod.rs
@@ -3,7 +3,7 @@
 pub use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
 use std::{
     collections::{HashMap, HashSet, VecDeque},
-    io,
+    io::{self, Read, Write},
     num::NonZeroUsize,
     os::fd::BorrowedFd,
     pin::Pin,
@@ -16,7 +16,7 @@ use std::{
 use compio_log::{instrument, trace};
 use crossbeam_queue::SegQueue;
 pub(crate) use libc::{sockaddr_storage, socklen_t};
-use polling::{Event, Events, Poller};
+use polling::{Event, Events, PollMode, Poller};
 use slab::Slab;
 
 use crate::{syscall, AsyncifyPool, Entry, OutEntries, ProactorBuilder};
@@ -151,6 +151,7 @@ pub(crate) struct Driver {
     poll: Arc<Poller>,
     registry: HashMap<RawFd, FdQueue>,
     cancelled: HashSet<usize>,
+    notifier: Notifier,
     pool: AsyncifyPool,
     pool_completed: Arc<SegQueue<Entry>>,
 }
@@ -166,11 +167,21 @@ impl Driver {
             Events::with_capacity(NonZeroUsize::new(entries).unwrap())
         };
 
+        let notifier = Notifier::new()?;
+        let fd = notifier.reader_fd();
+
+        let poll = Arc::new(Poller::new()?);
+        // Attach the reader to poll.
+        unsafe {
+            poll.add_with_mode(fd, Event::new(fd as _, true, false), PollMode::Level)?;
+        }
+
         Ok(Self {
             events,
-            poll: Arc::new(Poller::new()?),
+            poll,
             registry: HashMap::new(),
             cancelled: HashSet::new(),
+            notifier,
             pool: builder.create_or_get_thread_pool(),
             pool_completed: Arc::new(SegQueue::new()),
         })
@@ -279,6 +290,10 @@ impl Driver {
         }
         for event in self.events.iter() {
             let fd = event.key as RawFd;
+            if fd == self.notifier.reader_fd() {
+                self.notifier.clear()?;
+                continue;
+            }
             let queue = self
                 .registry
                 .get_mut(&fd)
@@ -310,7 +325,7 @@ impl Driver {
     }
 
     pub fn handle(&self) -> io::Result<NotifyHandle> {
-        Ok(NotifyHandle::new(self.poll.clone()))
+        self.notifier.handle()
     }
 }
 
@@ -338,18 +353,59 @@ fn entry_cancelled(user_data: usize) -> Entry {
     )
 }
 
+struct Notifier {
+    notify_reader: os_pipe::PipeReader,
+    notify_writer: os_pipe::PipeWriter,
+}
+
+impl Notifier {
+    pub fn new() -> io::Result<Self> {
+        let (notify_reader, notify_writer) = os_pipe::pipe()?;
+
+        // Set the reader as nonblocking.
+        let fd = notify_reader.as_raw_fd();
+        let current_flags = syscall!(libc::fcntl(fd, libc::F_GETFL))?;
+        let flags = current_flags | libc::O_NONBLOCK;
+        if flags != current_flags {
+            syscall!(libc::fcntl(fd, libc::F_SETFL, flags))?;
+        }
+
+        Ok(Self {
+            notify_reader,
+            notify_writer,
+        })
+    }
+
+    pub fn handle(&self) -> io::Result<NotifyHandle> {
+        Ok(NotifyHandle::new(self.notify_writer.try_clone()?))
+    }
+
+    pub fn clear(&self) -> io::Result<()> {
+        let mut buffer = [0u8];
+        match (&self.notify_reader).read_exact(&mut buffer) {
+            Ok(()) => Ok(()),
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => Ok(()),
+            Err(e) => Err(e),
+        }
+    }
+
+    pub fn reader_fd(&self) -> RawFd {
+        self.notify_reader.as_raw_fd()
+    }
+}
+
 /// A notify handle to the inner driver.
 pub struct NotifyHandle {
-    poll: Arc<Poller>,
+    sender: os_pipe::PipeWriter,
 }
 
 impl NotifyHandle {
-    fn new(poll: Arc<Poller>) -> Self {
-        Self { poll }
+    fn new(sender: os_pipe::PipeWriter) -> Self {
+        Self { sender }
     }
 
     /// Notify the inner driver.
     pub fn notify(&self) -> io::Result<()> {
-        self.poll.notify()
+        (&self.sender).write_all(&[1u8])
     }
 }
diff --git a/compio-runtime/Cargo.toml b/compio-runtime/Cargo.toml
@@ -52,6 +52,16 @@ windows-sys = { workspace = true, features = ["Win32_System_IO"] }
 os_pipe = { workspace = true }
 libc = { workspace = true }
 
+[target.'cfg(windows)'.dev-dependencies]
+windows-sys = { workspace = true, features = ["Win32_UI_WindowsAndMessaging"] }
+
+[target.'cfg(target_os = "macos")'.dev-dependencies]
+core-foundation = "0.9.4"
+block2 = "0.4.0"
+
+[target.'cfg(not(any(windows, target_os = "macos")))'.dev-dependencies]
+glib = "0.19"
+
 [features]
 event = ["dep:cfg-if", "compio-buf/arrayvec"]
 time = ["dep:slab"]
diff --git a/compio-runtime/src/runtime/mod.rs b/compio-runtime/src/runtime/mod.rs
diff --git a/compio-runtime/tests/custom_loop.rs b/compio-runtime/tests/custom_loop.rs
