Resource table tracks child relationships

crates/test-programs/command-tests/src/bin/stream_pollable_lifetimes.rs

@@ -0,0 +1,28 @@
+use command_tests::wasi::cli_base::environment;
+use command_tests::wasi::cli_base::stdin;
+use command_tests::wasi::io::streams;
+use command_tests::wasi::poll::poll;
+
+fn main() {
+    let args = environment::get_arguments();
+
+    if args == &["correct"] {
+        let stdin: streams::InputStream = stdin::get_stdin();
+        let stdin_pollable = streams::subscribe_to_input_stream(stdin);
+        let ready = poll::poll_oneoff(&[stdin_pollable]);
+        assert_eq!(ready, &[true]);
+        poll::drop_pollable(stdin_pollable);
+        streams::drop_input_stream(stdin);
+    } else if args == &["trap"] {
+        let stdin: streams::InputStream = stdin::get_stdin();
+        let stdin_pollable = streams::subscribe_to_input_stream(stdin);
+        let ready = poll::poll_oneoff(&[stdin_pollable]);
+        assert_eq!(ready, &[true]);
+        streams::drop_input_stream(stdin);
+        unreachable!(
+            "execution should have trapped in line above when stream dropped before pollable"
+        );
+    } else {
+        panic!("bad value for args: expected `[\"correct\"]` or `[\"trap\"]`, got {args:?}")
+    }
+}

crates/test-programs/command-tests/src/lib.rs

@@ -0,0 +1 @@
+wit_bindgen::generate!("test-command" in "../../wasi/wit");

crates/test-programs/tests/command.rs

@@ -443,3 +443,50 @@ async fn read_only() -> Result<()> {
         .await?
         .map_err(|()| anyhow::anyhow!("command returned with failing exit status"))
 }
+
+#[test_log::test(tokio::test(flavor = "multi_thread"))]
+async fn stream_pollable_lifetimes() -> Result<()> {
+    // Test program has two modes, dispatching based on argument.
+    {
+        // Correct execution: should succeed
+        let mut table = Table::new();
+        let wasi = WasiCtxBuilder::new()
+            .set_args(&["correct"])
+            .set_stdin(MemoryInputPipe::new(" ".into()))
+            .build(&mut table)?;
+
+        let (mut store, command) = instantiate(
+            get_component("stream_pollable_lifetimes"),
+            CommandCtx { table, wasi },
+        )
+        .await?;
+
+        command
+            .call_run(&mut store)
+            .await?
+            .map_err(|()| anyhow::anyhow!("command returned with failing exit status"))?;
+    }
+    {
+        // Incorrect execution: should trap with a TableError::HasChildren
+        let mut table = Table::new();
+        let wasi = WasiCtxBuilder::new()
+            .set_args(&["trap"])
+            .set_stdin(MemoryInputPipe::new(" ".into()))
+            .build(&mut table)?;
+
+        let (mut store, command) = instantiate(
+            get_component("stream_pollable_lifetimes"),
+            CommandCtx { table, wasi },
+        )
+        .await?;
+
+        let trap = command
+            .call_run(&mut store)
+            .await
+            .err()
+            .expect("should trap");
+        use wasmtime_wasi::preview2::TableError;
+        assert!(matches!(trap.downcast_ref(), Some(TableError::HasChildren)));
+    }
+    Ok(())
+}

crates/wasi/src/preview2/mod.rs

@@ -37,7 +37,7 @@ pub use self::filesystem::{DirPerms, FilePerms};
 pub use self::poll::{ClosureFuture, HostPollable, MakeFuture, PollableFuture, TablePollableExt};
 pub use self::random::{thread_rng, Deterministic};
 pub use self::stream::{HostInputStream, HostOutputStream, StreamState, TableStreamExt};
-pub use self::table::{Table, TableError};
+pub use self::table::{OccupiedEntry, Table, TableError};
 pub use cap_fs_ext::SystemTimeSpec;
 pub use cap_rand::RngCore;
 

crates/wasi/src/preview2/poll.rs

@@ -23,18 +23,6 @@ pub enum HostPollable {
     /// Create a Future by calling a fn on another resource in the table. This
     /// indirection means the created Future can use a mut borrow of another
     /// resource in the Table (e.g. a stream)
-    ///
-    /// FIXME: we currently aren't tracking the lifetime of the resource along
-    /// with this entry, which means that this index could be occupied by something
-    /// unrelated by the time we poll it again. This is a crash vector, because
-    /// the [`MakeFuture`] would panic if the type of the index has changed, and
-    /// would yield undefined behavior otherwise. We'll likely fix this by making
-    /// the parent resources of a pollable clean up their pollable entries when
-    /// they are destroyed (e.g. the HostInputStream would track the pollables it
-    /// has created).
-    ///
-    /// WARNING: do not deploy this library to production until the above issue has
-    /// been fixed.
     TableEntry { index: u32, make_future: MakeFuture },
     /// Create a future by calling an owned, static closure. This is used for
     /// pollables which do not share state with another resource in the Table
@@ -50,7 +38,10 @@ pub trait TablePollableExt {
 
 impl TablePollableExt for Table {
     fn push_host_pollable(&mut self, p: HostPollable) -> Result<u32, TableError> {
-        self.push(Box::new(p))
+        match p {
+            HostPollable::TableEntry { index, .. } => self.push_child(Box::new(p), index),
+            HostPollable::Closure { .. } => self.push(Box::new(p)),
+        }
     }
     fn get_host_pollable_mut(&mut self, fd: u32) -> Result<&mut HostPollable, TableError> {
         self.get_mut::<HostPollable>(fd)

crates/wasi/src/preview2/preview1/mod.rs

@@ -527,18 +527,9 @@ impl TryFrom<filesystem::Error> for types::Error {
     }
 }
 
-impl From<TableError> for types::Errno {
-    fn from(err: TableError) -> Self {
-        match err {
-            TableError::Full => types::Errno::Nomem,
-            TableError::NotPresent | TableError::WrongType => types::Errno::Badf,
-        }
-    }
-}
-
 impl From<TableError> for types::Error {
     fn from(err: TableError) -> Self {
-        types::Errno::from(err).into()
+        types::Error::trap(err.into())
     }
 }
 

crates/wasi/src/preview2/preview2/filesystem.rs

@@ -9,17 +9,8 @@ use filesystem::ErrorCode;
 mod sync;
 
 impl From<TableError> for filesystem::Error {
-    fn from(error: TableError) -> filesystem::Error {
-        match error {
-            TableError::Full => filesystem::Error::trap(anyhow::anyhow!(error)),
-            TableError::NotPresent | TableError::WrongType => ErrorCode::BadDescriptor.into(),
-        }
-    }
-}
-
-impl From<tokio::task::JoinError> for filesystem::Error {
-    fn from(error: tokio::task::JoinError) -> Self {
-        Self::trap(anyhow::anyhow!(error))
+    fn from(error: TableError) -> Self {
+        Self::trap(error.into())
     }
 }
 

crates/wasi/src/preview2/preview2/io.rs

@@ -23,13 +23,7 @@ impl From<anyhow::Error> for streams::Error {
 
 impl From<TableError> for streams::Error {
     fn from(error: TableError) -> streams::Error {
-        match error {
-            TableError::Full => streams::Error::trap(anyhow!(error)),
-            TableError::NotPresent | TableError::WrongType => {
-                // wit definition needs to define a badf-equiv variant:
-                StreamError { dummy: 0 }.into()
-            }
-        }
+        streams::Error::trap(anyhow!(error))
     }
 }
 

crates/wasi/src/preview2/stream.rs

@@ -193,7 +193,7 @@ impl TableStreamExt for Table {
         let occ = self.entry(fd)?;
         match occ.get().downcast_ref::<InternalInputStream>() {
             Some(InternalInputStream::Host(_)) => {
-                let (_, any) = occ.remove_entry();
+                let any = occ.remove_entry()?;
                 match *any.downcast().expect("downcast checked above") {
                     InternalInputStream::Host(h) => Ok(h),
                     _ => unreachable!("variant checked above"),
@@ -219,7 +219,7 @@ impl TableStreamExt for Table {
         let occ = self.entry(fd)?;
         match occ.get().downcast_ref::<InternalOutputStream>() {
             Some(InternalOutputStream::Host(_)) => {
-                let (_, any) = occ.remove_entry();
+                let any = occ.remove_entry()?;
                 match *any.downcast().expect("downcast checked above") {
                     InternalOutputStream::Host(h) => Ok(h),
                     _ => unreachable!("variant checked above"),