Rework how pipes and event FDs interact with the epoll system

Author: Arshia001

Cargo.lock

@@ -11,6 +11,7 @@ rust-version.workspace = true
 
 [dependencies]
 wasmer-package.workspace = true
+virtual-mio = { path = "../virtual-io", version = "0.601.0-rc.2", default-features = false }
 dashmap.workspace = true
 derive_more.workspace = true
 dunce = "1.0.4"
@@ -40,7 +41,7 @@ serde = { workspace = true, default-features = false, features = [
 	"derive",
 ], optional = true }
 getrandom.workspace = true
-web-time = { version = "1.1", optional = true}
+web-time = { version = "1.1", optional = true }
 
 [dev-dependencies]
 pretty_assertions.workspace = true
@@ -64,10 +65,7 @@ host-fs = [
 webc-fs = ["webc", "anyhow"]
 static-fs = ["webc", "anyhow"]
 enable-serde = ["typetag", "serde"]
-js = [
-	"dep:web-time",
-	"getrandom/js",
-]
+js = ["dep:web-time", "getrandom/js"]
 # Enables memory tracking/limiting functionality for the in-memory filesystem.
 tracking = []
 futures = []

lib/virtual-fs/Cargo.toml

@@ -1,16 +1,20 @@
 use bytes::{Buf, Bytes};
-use std::io::{self, Read, Seek, SeekFrom};
 use std::pin::Pin;
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::task::Context;
 use std::task::Poll;
 use std::{io::IoSlice, sync::MutexGuard};
+use std::{
+    io::{self, Read, Seek, SeekFrom},
+    sync::Weak,
+};
 use tokio::sync::mpsc;
 use tokio::{
     io::{AsyncRead, AsyncSeek, AsyncWrite},
     sync::mpsc::error::TryRecvError,
 };
+use virtual_mio::{InterestHandler, InterestType};
 
 use crate::{ArcFile, FsError, VirtualFile};
 
@@ -31,6 +35,7 @@ pub struct Pipe {
 pub struct PipeTx {
     /// Sends bytes down the pipe
     tx: Option<mpsc::UnboundedSender<Vec<u8>>>,
+    rx_end: Weak<Mutex<PipeReceiver>>,
 }
 
 #[derive(Debug, Clone)]
@@ -124,6 +129,22 @@ impl PipeRx {
             rx.buffer.replace(Bytes::from(data));
         }
     }
+
+    pub fn set_interest_handler(&mut self, interest_handler: Box<dyn InterestHandler>) {
+        let Some(ref rx) = self.rx else {
+            return;
+        };
+        let mut rx = rx.lock().unwrap();
+        rx.interest_handler.replace(interest_handler);
+    }
+
+    pub fn remove_interest_handler(&mut self) -> Option<Box<dyn InterestHandler>> {
+        let Some(ref rx) = self.rx else {
+            return None;
+        };
+        let mut rx = rx.lock().unwrap();
+        rx.interest_handler.take()
+    }
 }
 
 #[derive(Debug)]
@@ -132,20 +153,24 @@ struct PipeReceiver {
     // actual receiver, we can't make use of an mpmc channel
     chan: mpsc::UnboundedReceiver<Vec<u8>>,
     buffer: Option<Bytes>,
+    interest_handler: Option<Box<dyn InterestHandler>>,
 }
 
 impl Pipe {
     pub fn new() -> Self {
         let (tx, rx) = mpsc::unbounded_channel();
 
+        let recv = Arc::new(Mutex::new(PipeReceiver {
+            chan: rx,
+            buffer: None,
+            interest_handler: None,
+        }));
         Pipe {
-            send: PipeTx { tx: Some(tx) },
-            recv: PipeRx {
-                rx: Some(Arc::new(Mutex::new(PipeReceiver {
-                    chan: rx,
-                    buffer: None,
-                }))),
+            send: PipeTx {
+                tx: Some(tx),
+                rx_end: Arc::downgrade(&recv),
             },
+            recv: PipeRx { rx: Some(recv) },
         }
     }
 
@@ -174,23 +199,19 @@ impl Pipe {
         self.send.close();
         self.recv.close();
     }
-}
 
-impl Default for Pipe {
-    fn default() -> Self {
-        Self::new()
+    pub fn set_interest_handler(&mut self, interest_handler: Box<dyn InterestHandler>) {
+        self.recv.set_interest_handler(interest_handler);
     }
-}
 
-impl From<Pipe> for PipeTx {
-    fn from(val: Pipe) -> Self {
-        val.send
+    pub fn remove_interest_handler(&mut self) -> Option<Box<dyn InterestHandler>> {
+        self.recv.remove_interest_handler()
     }
 }
 
-impl From<Pipe> for PipeRx {
-    fn from(val: Pipe) -> Self {
-        val.recv
+impl Default for Pipe {
+    fn default() -> Self {
+        Self::new()
     }
 }
 
@@ -213,6 +234,15 @@ impl PipeTx {
             Poll::Ready(Ok(8192))
         }
     }
+
+    fn push_readable_interest_to_rx_end(&self) {
+        if let Some(rx_end) = self.rx_end.upgrade() {
+            let mut guard = rx_end.lock().unwrap();
+            if let Some(interest_handler) = guard.interest_handler.as_mut() {
+                interest_handler.push_interest(InterestType::Readable);
+            }
+        }
+    }
 }
 
 impl Seek for Pipe {
@@ -302,6 +332,7 @@ impl std::io::Write for PipeTx {
 
         tx.send(buf.to_vec())
             .map_err(|_| Into::<std::io::Error>::into(std::io::ErrorKind::BrokenPipe))?;
+        self.push_readable_interest_to_rx_end();
         Ok(buf.len())
     }
 
@@ -387,7 +418,10 @@ impl AsyncWrite for PipeTx {
         };
 
         match tx.send(buf.to_vec()) {
-            Ok(()) => Poll::Ready(Ok(buf.len())),
+            Ok(()) => {
+                self.push_readable_interest_to_rx_end();
+                Poll::Ready(Ok(buf.len()))
+            }
             Err(_) => Poll::Ready(Err(Into::<std::io::Error>::into(
                 std::io::ErrorKind::BrokenPipe,
             ))),

lib/virtual-fs/src/pipe.rs

@@ -2,23 +2,19 @@ use std::{
     borrow::Cow,
     collections::{HashMap, HashSet},
     path::PathBuf,
-    pin::Pin,
     sync::{atomic::AtomicU64, Arc, RwLock, RwLockReadGuard, RwLockWriteGuard},
-    task::Context,
 };
 
-use futures::Future;
 use serde_derive::{Deserialize, Serialize};
 use std::sync::Mutex as StdMutex;
 use tokio::sync::{watch, Mutex as AsyncMutex};
 use virtual_fs::{Pipe, PipeRx, PipeTx, VirtualFile};
 use wasmer_wasix_types::wasi::{EpollType, Fd as WasiFd, Fdflags, Fdflagsext, Filestat, Rights};
 
-use crate::{net::socket::InodeSocket, syscalls::EpollJoinWaker};
+use crate::net::socket::InodeSocket;
 
 use super::{
-    InodeGuard, InodeValFilePollGuard, InodeValFilePollGuardJoin, InodeValFilePollGuardMode,
-    InodeWeakGuard, NotificationInner,
+    InodeGuard, InodeValFilePollGuard, InodeValFilePollGuardMode, InodeWeakGuard, NotificationInner,
 };
 
 #[derive(Debug, Clone)]
@@ -111,48 +107,32 @@ pub struct EpollInterest {
 
 /// Guard the cleans up the selector registrations
 #[derive(Debug)]
-pub enum EpollJoinGuard {
-    Join {
-        join_guard: InodeValFilePollGuardJoin,
-        epoll_waker: Arc<EpollJoinWaker>,
-    },
-    Handler {
-        fd_guard: InodeValFilePollGuard,
-    },
+pub struct EpollJoinGuard {
+    pub(crate) fd_guard: InodeValFilePollGuard,
 }
 impl Drop for EpollJoinGuard {
     fn drop(&mut self) {
-        if let Self::Handler { fd_guard, .. } = self {
-            if let InodeValFilePollGuardMode::Socket { inner } = &mut fd_guard.mode {
+        match &self.fd_guard.mode {
+            InodeValFilePollGuardMode::File(_) => {
+                // Intentionally ignored, epoll doesn't work with files
+            }
+            InodeValFilePollGuardMode::Socket { inner } => {
                 let mut inner = inner.protected.write().unwrap();
                 inner.remove_handler();
             }
-        }
-    }
-}
-impl EpollJoinGuard {
-    pub fn is_spent(&self) -> bool {
-        match self {
-            Self::Join { join_guard, .. } => join_guard.is_spent(),
-            Self::Handler { .. } => false,
-        }
-    }
-    pub fn renew(&mut self) {
-        if let Self::Join {
-            join_guard,
-            epoll_waker,
-        } = self
-        {
-            let fd = join_guard.fd();
-            join_guard.reset();
-
-            let waker = epoll_waker.as_waker();
-            let mut cx = Context::from_waker(&waker);
-            if Pin::new(join_guard).poll(&mut cx).is_ready() {
-                tracing::trace!(fd, "join renew already woken");
-                waker.wake();
-            } else {
-                tracing::trace!(fd, "join waker reinstalled");
+            InodeValFilePollGuardMode::EventNotifications(inner) => {
+                inner.remove_interest_handler();
+            }
+            InodeValFilePollGuardMode::DuplexPipe { pipe } => {
+                let mut inner = pipe.write().unwrap();
+                inner.remove_interest_handler();
+            }
+            InodeValFilePollGuardMode::PipeRx { rx } => {
+                let mut inner = rx.write().unwrap();
+                inner.remove_interest_handler();
+            }
+            InodeValFilePollGuardMode::PipeTx { .. } => {
+                // Intentionally ignored, the sending end of a pipe can't have an interest handler
             }
         }
     }

lib/wasix/src/fs/fd.rs

@@ -134,7 +134,6 @@ pub struct InodeValFilePollGuardJoin {
     fd: u32,
     peb: PollEventSet,
     subscription: Subscription,
-    spent: bool,
 }
 
 impl InodeValFilePollGuardJoin {
@@ -144,7 +143,6 @@ impl InodeValFilePollGuardJoin {
             fd: guard.fd,
             peb: guard.peb,
             subscription: guard.subscription,
-            spent: false,
         }
     }
     pub(crate) fn fd(&self) -> u32 {
@@ -153,22 +151,6 @@ impl InodeValFilePollGuardJoin {
     pub(crate) fn peb(&self) -> PollEventSet {
         self.peb
     }
-    pub fn is_spent(&self) -> bool {
-        self.spent
-    }
-    pub fn reset(&mut self) {
-        match &self.mode {
-            InodeValFilePollGuardMode::File(_) => {}
-            InodeValFilePollGuardMode::EventNotifications(inner) => {
-                inner.reset();
-            }
-            InodeValFilePollGuardMode::Socket { .. }
-            | InodeValFilePollGuardMode::PipeRx { .. }
-            | InodeValFilePollGuardMode::PipeTx { .. }
-            | InodeValFilePollGuardMode::DuplexPipe { .. } => {}
-        }
-        self.spent = false;
-    }
 }
 
 pub const POLL_GUARD_MAX_RET: usize = 4;
@@ -403,7 +385,6 @@ impl Future for InodeValFilePollGuardJoin {
             };
         }
         if !ret.is_empty() {
-            self.spent = true;
             return Poll::Ready(ret);
         }
         Poll::Pending

lib/wasix/src/fs/inode_guard.rs

@@ -4,6 +4,8 @@ use std::{
     task::{Poll, Waker},
 };
 
+use virtual_mio::{InterestHandler, InterestType};
+
 #[derive(Debug)]
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
 struct NotificationState {
@@ -17,6 +19,9 @@ struct NotificationState {
     /// All the registered wakers
     #[cfg_attr(feature = "enable-serde", serde(skip))]
     wakers: VecDeque<Waker>,
+    /// InterestHandler for use with epoll
+    #[cfg_attr(feature = "enable-serde", serde(skip))]
+    interest_handler: Option<Box<dyn InterestHandler>>,
 }
 
 impl NotificationState {
@@ -31,6 +36,9 @@ impl NotificationState {
         while let Some(waker) = self.wakers.pop_front() {
             waker.wake();
         }
+        if let Some(handler) = self.interest_handler.as_mut() {
+            handler.push_interest(InterestType::Readable);
+        }
     }
 
     fn inc(&mut self, val: u64) {
@@ -70,6 +78,7 @@ impl NotificationInner {
                 last_poll: u64::MAX,
                 is_semaphore,
                 wakers: Default::default(),
+                interest_handler: None,
             }),
         }
     }
@@ -111,4 +120,14 @@ impl NotificationInner {
         let mut state = self.state.lock().unwrap();
         state.last_poll = u64::MAX;
     }
+
+    pub fn set_interest_handler(&self, handler: Box<dyn InterestHandler>) {
+        let mut state = self.state.lock().unwrap();
+        state.interest_handler.replace(handler);
+    }
+
+    pub fn remove_interest_handler(&self) -> Option<Box<dyn InterestHandler>> {
+        let mut state = self.state.lock().unwrap();
+        state.interest_handler.take()
+    }
 }