Cleanup

lib/runtime/src/platform/macos_aarch64/driver.rs

@@ -1,12 +1,37 @@
 //! Public runtime driver primitives for macOS.
 
-use std::cell::Cell;
+use std::cell::{Cell, RefCell};
+use std::collections::HashMap;
 use std::io;
 use std::os::fd::RawFd;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::Duration;
 
+use crate::op::completion::CompletionHandle;
+
+type FdCompletion = CompletionHandle<io::Result<()>>;
+
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+pub(crate) struct FdReadinessToken(u64);
+
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+pub(crate) enum FdInterest {
+    Readable,
+    Writable,
+}
+
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+struct FdKey {
+    fd: RawFd,
+    interest: FdInterest,
+}
+
+struct FdWaiter {
+    token: FdReadinessToken,
+    completion: FdCompletion,
+}
+
 #[derive(Clone)]
 struct NotifierInner {
     write_fd: RawFd,
@@ -73,6 +98,8 @@ pub struct Driver {
     timer_deadline: Cell<Option<Duration>>,
     pending_wakes: Cell<u64>,
     pending_timers: Cell<u64>,
+    next_fd_token: Cell<u64>,
+    fd_waiters: RefCell<HashMap<FdKey, Vec<FdWaiter>>>,
 }
 
 /// Creates a new driver and its paired [`ThreadNotifier`].
@@ -130,6 +157,8 @@ pub fn create_driver() -> io::Result<(Driver, ThreadNotifier)> {
         timer_deadline: Cell::new(None),
         pending_wakes: Cell::new(0),
         pending_timers: Cell::new(0),
+        next_fd_token: Cell::new(1),
+        fd_waiters: RefCell::new(HashMap::new()),
     };
 
     let notifier = ThreadNotifier {
@@ -205,6 +234,51 @@ impl Driver {
         }
     }
 
+    pub(crate) fn register_fd_readiness(
+        &self,
+        fd: RawFd,
+        interest: FdInterest,
+        completion: FdCompletion,
+    ) -> io::Result<FdReadinessToken> {
+        let key = FdKey { fd, interest };
+        let should_register = !self.fd_waiters.borrow().contains_key(&key);
+        if should_register {
+            self.update_fd_interest(key, libc::EV_ADD | libc::EV_ENABLE)?;
+        }
+
+        let token = self.allocate_fd_token();
+        self.fd_waiters
+            .borrow_mut()
+            .entry(key)
+            .or_default()
+            .push(FdWaiter { token, completion });
+        Ok(token)
+    }
+
+    pub(crate) fn cancel_fd_readiness(&self, token: FdReadinessToken) {
+        let mut empty_key = None;
+        {
+            let mut waiters = self.fd_waiters.borrow_mut();
+            for (key, entries) in waiters.iter_mut() {
+                if let Some(index) = entries.iter().position(|entry| entry.token == token) {
+                    entries.swap_remove(index);
+                    if entries.is_empty() {
+                        empty_key = Some(*key);
+                    }
+                    break;
+                }
+            }
+
+            if let Some(key) = empty_key {
+                waiters.remove(&key);
+            }
+        }
+
+        if let Some(key) = empty_key {
+            let _ = self.update_fd_interest(key, libc::EV_DELETE);
+        }
+    }
+
     fn process(&self, timeout: Option<Duration>) -> io::Result<Option<ReadyEvents>> {
         let mut ready = ReadyEvents::default();
 
@@ -241,6 +315,8 @@ impl Driver {
                     let wakes = drain_wake_pipe(self.wake_read_fd)?;
                     self.pending_wakes
                         .set(self.pending_wakes.get().saturating_add(wakes));
+                } else if let Some(interest) = interest_from_filter(event.filter) {
+                    self.complete_fd_waiters(event.ident as RawFd, interest, event);
                 }
             }
         }
@@ -257,6 +333,59 @@ impl Driver {
 
         if saw_any { Ok(Some(ready)) } else { Ok(None) }
     }
+
+    fn allocate_fd_token(&self) -> FdReadinessToken {
+        let token = self.next_fd_token.get();
+        self.next_fd_token.set(
+            token
+                .checked_add(1)
+                .expect("fd readiness token space exhausted"),
+        );
+        FdReadinessToken(token)
+    }
+
+    fn update_fd_interest(&self, key: FdKey, flags: u16) -> io::Result<()> {
+        let event = libc::kevent {
+            ident: key.fd as usize,
+            filter: filter_for_interest(key.interest),
+            flags,
+            fflags: 0,
+            data: 0,
+            udata: std::ptr::null_mut(),
+        };
+        let submitted = unsafe {
+            libc::kevent(
+                self.kqueue_fd,
+                &event,
+                1,
+                std::ptr::null_mut(),
+                0,
+                std::ptr::null(),
+            )
+        };
+        if submitted < 0 {
+            Err(io::Error::last_os_error())
+        } else {
+            Ok(())
+        }
+    }
+
+    fn complete_fd_waiters(&self, fd: RawFd, interest: FdInterest, event: &libc::kevent) {
+        let key = FdKey { fd, interest };
+        let waiters = self.fd_waiters.borrow_mut().remove(&key);
+        let Some(waiters) = waiters else {
+            return;
+        };
+
+        let _ = self.update_fd_interest(key, libc::EV_DELETE);
+        let result = fd_event_result(event);
+        for waiter in waiters {
+            waiter.completion.complete(match &result {
+                Ok(()) => Ok(()),
+                Err(error) => Err(io::Error::new(error.kind(), error.to_string())),
+            });
+        }
+    }
 }
 
 impl Drop for Driver {
@@ -290,6 +419,31 @@ fn cvt(value: libc::c_int) -> io::Result<libc::c_int> {
     }
 }
 
+fn filter_for_interest(interest: FdInterest) -> i16 {
+    match interest {
+        FdInterest::Readable => libc::EVFILT_READ,
+        FdInterest::Writable => libc::EVFILT_WRITE,
+    }
+}
+
+fn interest_from_filter(filter: i16) -> Option<FdInterest> {
+    if filter == libc::EVFILT_READ {
+        Some(FdInterest::Readable)
+    } else if filter == libc::EVFILT_WRITE {
+        Some(FdInterest::Writable)
+    } else {
+        None
+    }
+}
+
+fn fd_event_result(event: &libc::kevent) -> io::Result<()> {
+    if event.flags & libc::EV_ERROR != 0 && event.data != 0 {
+        Err(io::Error::from_raw_os_error(event.data as i32))
+    } else {
+        Ok(())
+    }
+}
+
 fn set_nonblocking(fd: RawFd) -> io::Result<()> {
     let flags = cvt(unsafe { libc::fcntl(fd, libc::F_GETFL) })?;
     cvt(unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) })?;