wasi-io: Reimplement wasi-io/poll using a Pollable trait

crates/test-programs/src/bin/preview2_pollable_correct.rs

@@ -0,0 +1,12 @@
+use test_programs::wasi::cli::stdin;
+
+fn main() {
+    let stdin = stdin::get_stdin();
+    let p1 = stdin.subscribe();
+    let p2 = stdin.subscribe();
+
+    // Should work:
+    // - Exactly the same pollable passed in multiple times.
+    // - Distinct pollables for the same backing resource (stdin in this case).
+    test_programs::wasi::io::poll::poll(&[&p1, &p2, &p1, &p2]);
+}

crates/test-programs/src/bin/preview2_pollable_traps.rs

@@ -0,0 +1,4 @@
+fn main() {
+    // Polling an empty list should trap:
+    test_programs::wasi::io::poll::poll(&[]);
+}

crates/wasi-http/src/types_impl.rs

@@ -644,7 +644,7 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFutureTrailers for T {
         &mut self,
         index: Resource<HostFutureTrailers>,
     ) -> wasmtime::Result<Resource<Pollable>> {
-        wasmtime_wasi::preview2::subscribe(self.table(), index)
+        wasmtime_wasi::preview2::subscribe(self.table(), &index)
     }
 
     fn get(
@@ -852,7 +852,7 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFutureIncomingResponse f
         &mut self,
         id: Resource<HostFutureIncomingResponse>,
     ) -> wasmtime::Result<Resource<Pollable>> {
-        wasmtime_wasi::preview2::subscribe(self.table(), id)
+        wasmtime_wasi::preview2::subscribe(self.table(), &id)
     }
 }
 

crates/wasi-http/wit/deps/io/poll.wit

@@ -27,8 +27,9 @@ interface poll {
     /// The result `list<u32>` contains one or more indices of handles in the
     /// argument list that is ready for I/O.
     ///
-    /// If the list contains more elements than can be indexed with a `u32`
-    /// value, this function traps.
+    /// This function traps if either:
+    /// - the list is empty, or:
+    /// - the list contains more elements than can be indexed with a `u32` value.
     ///
     /// A timeout can be implemented by adding a pollable from the
     /// wasi-clocks API to the list.

crates/wasi/Cargo.toml

@@ -39,6 +39,7 @@ io-lifetimes = { workspace = true, optional = true }
 fs-set-times = { workspace = true, optional = true }
 bitflags = { workspace = true, optional = true }
 async-trait = { workspace = true, optional = true }
+smallvec = { workspace = true }
 system-interface = { workspace = true, optional = true}
 futures = { workspace = true, optional = true }
 

crates/wasi/src/preview2/host/clocks.rs

@@ -1,10 +1,8 @@
-#![allow(unused_variables)]
-
 use crate::preview2::bindings::{
     clocks::monotonic_clock::{self, Duration as WasiDuration, Instant},
     clocks::wall_clock::{self, Datetime},
 };
-use crate::preview2::poll::{subscribe, Subscribe};
+use crate::preview2::poll;
 use crate::preview2::{Pollable, WasiView};
 use cap_std::time::SystemTime;
 use std::time::Duration;
@@ -42,23 +40,16 @@ impl<T: WasiView> wall_clock::Host for T {
     }
 }
 
-fn subscribe_to_duration(
-    table: &mut wasmtime::component::ResourceTable,
-    duration: tokio::time::Duration,
-) -> anyhow::Result<Resource<Pollable>> {
-    let sleep = if duration.is_zero() {
-        table.push(Deadline::Past)?
+fn make_pollable(duration: tokio::time::Duration) -> Pollable {
+    if duration.is_zero() {
+        poll::ready()
     } else if let Some(deadline) = tokio::time::Instant::now().checked_add(duration) {
-        // NB: this resource created here is not actually exposed to wasm, it's
-        // only an internal implementation detail used to match the signature
-        // expected by `subscribe`.
-        table.push(Deadline::Instant(deadline))?
+        poll::once(async move { tokio::time::sleep_until(deadline).await })
     } else {
         // If the user specifies a time so far in the future we can't
         // represent it, wait forever rather than trap.
-        table.push(Deadline::Never)?
-    };
-    subscribe(table, sleep)
+        poll::pending()
+    }
 }
 
 impl<T: WasiView> monotonic_clock::Host for T {
@@ -77,27 +68,12 @@ impl<T: WasiView> monotonic_clock::Host for T {
         } else {
             Duration::from_nanos(0)
         };
-        subscribe_to_duration(&mut self.table(), duration)
+        let pollable = make_pollable(duration);
+        Ok(self.table().push(pollable)?)
     }
 
     fn subscribe_duration(&mut self, duration: WasiDuration) -> anyhow::Result<Resource<Pollable>> {
-        subscribe_to_duration(&mut self.table(), Duration::from_nanos(duration))
-    }
-}
-
-enum Deadline {
-    Past,
-    Instant(tokio::time::Instant),
-    Never,
-}
-
-#[async_trait::async_trait]
-impl Subscribe for Deadline {
-    async fn ready(&mut self) {
-        match self {
-            Deadline::Past => {}
-            Deadline::Instant(instant) => tokio::time::sleep_until(*instant).await,
-            Deadline::Never => std::future::pending().await,
-        }
+        let pollable = make_pollable(Duration::from_nanos(duration));
+        Ok(self.table().push(pollable)?)
     }
 }

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

@@ -49,7 +49,7 @@ impl<T: WasiView> streams::HostOutputStream for T {
     }
 
     fn subscribe(&mut self, stream: Resource<OutputStream>) -> anyhow::Result<Resource<Pollable>> {
-        subscribe(self.table(), stream)
+        subscribe(self.table(), &stream)
     }
 
     async fn blocking_write_and_flush(
@@ -220,7 +220,7 @@ impl<T: WasiView> streams::HostInputStream for T {
     }
 
     fn subscribe(&mut self, stream: Resource<InputStream>) -> anyhow::Result<Resource<Pollable>> {
-        crate::preview2::poll::subscribe(self.table(), stream)
+        crate::preview2::poll::subscribe(self.table(), &stream)
     }
 }
 
@@ -230,9 +230,8 @@ pub mod sync {
             self as async_streams, Host as AsyncHost, HostInputStream as AsyncHostInputStream,
             HostOutputStream as AsyncHostOutputStream,
         },
-        bindings::sync_io::io::poll::Pollable,
         bindings::sync_io::io::streams::{self, InputStream, OutputStream},
-        in_tokio, StreamError, StreamResult, WasiView,
+        in_tokio, Pollable, StreamError, StreamResult, WasiView,
     };
     use wasmtime::component::Resource;
 

crates/wasi/src/preview2/host/tcp.rs

@@ -560,7 +560,7 @@ impl<T: WasiView> crate::preview2::host::tcp::tcp::HostTcpSocket for T {
     }
 
     fn subscribe(&mut self, this: Resource<tcp::TcpSocket>) -> anyhow::Result<Resource<Pollable>> {
-        crate::preview2::poll::subscribe(self.table(), this)
+        crate::preview2::poll::subscribe(self.table(), &this)
     }
 
     fn shutdown(

crates/wasi/src/preview2/host/udp.rs

@@ -98,12 +98,14 @@ impl<T: WasiView> udp::HostUdpSocket for T {
     )> {
         let table = self.table();
 
-        let has_active_streams = table
-            .iter_children(&this)?
-            .any(|c| c.is::<IncomingDatagramStream>() || c.is::<OutgoingDatagramStream>());
+        for child_result in table.iter_children(&this)? {
+            let Ok(child) = child_result else {
+                return Err(SocketError::trap(anyhow!("UDP stream taken.")));
+            };
 
-        if has_active_streams {
-            return Err(SocketError::trap(anyhow!("UDP streams not dropped yet")));
+            if child.is::<IncomingDatagramStream>() || child.is::<OutgoingDatagramStream>() {
+                return Err(SocketError::trap(anyhow!("UDP streams not dropped yet")));
+            }
         }
 
         let socket = table.get_mut(&this)?;
@@ -283,7 +285,7 @@ impl<T: WasiView> udp::HostUdpSocket for T {
     }
 
     fn subscribe(&mut self, this: Resource<udp::UdpSocket>) -> anyhow::Result<Resource<Pollable>> {
-        crate::preview2::poll::subscribe(self.table(), this)
+        crate::preview2::poll::subscribe(self.table(), &this)
     }
 
     fn drop(&mut self, this: Resource<udp::UdpSocket>) -> Result<(), anyhow::Error> {
@@ -363,7 +365,7 @@ impl<T: WasiView> udp::HostIncomingDatagramStream for T {
         &mut self,
         this: Resource<udp::IncomingDatagramStream>,
     ) -> anyhow::Result<Resource<Pollable>> {
-        crate::preview2::poll::subscribe(self.table(), this)
+        crate::preview2::poll::subscribe(self.table(), &this)
     }
 
     fn drop(&mut self, this: Resource<udp::IncomingDatagramStream>) -> Result<(), anyhow::Error> {
@@ -497,7 +499,7 @@ impl<T: WasiView> udp::HostOutgoingDatagramStream for T {
         &mut self,
         this: Resource<udp::OutgoingDatagramStream>,
     ) -> anyhow::Result<Resource<Pollable>> {
-        crate::preview2::poll::subscribe(self.table(), this)
+        crate::preview2::poll::subscribe(self.table(), &this)
     }
 
     fn drop(&mut self, this: Resource<udp::OutgoingDatagramStream>) -> Result<(), anyhow::Error> {

crates/wasi/src/preview2/ip_name_lookup.rs

@@ -69,7 +69,7 @@ impl<T: WasiView> HostResolveAddressStream for T {
         &mut self,
         resource: Resource<ResolveAddressStream>,
     ) -> Result<Resource<Pollable>> {
-        subscribe(self.table(), resource)
+        subscribe(self.table(), &resource)
     }
 
     fn drop(&mut self, resource: Resource<ResolveAddressStream>) -> Result<()> {

crates/wasi/src/preview2/mod.rs

@@ -28,7 +28,7 @@ mod host;
 mod ip_name_lookup;
 mod network;
 pub mod pipe;
-mod poll;
+pub mod poll;
 #[cfg(feature = "preview1-on-preview2")]
 pub mod preview0;
 #[cfg(feature = "preview1-on-preview2")]
@@ -45,7 +45,7 @@ pub use self::ctx::{WasiCtx, WasiCtxBuilder, WasiView};
 pub use self::error::{I32Exit, TrappableError};
 pub use self::filesystem::{DirPerms, FilePerms, FsError, FsResult};
 pub use self::network::{Network, SocketError, SocketResult};
-pub use self::poll::{subscribe, ClosureFuture, MakeFuture, Pollable, PollableFuture, Subscribe};
+pub use self::poll::{subscribe, Pollable, Subscribe};
 pub use self::random::{thread_rng, Deterministic};
 pub use self::stdio::{
     stderr, stdin, stdout, IsATTY, Stderr, Stdin, StdinStream, Stdout, StdoutStream,
@@ -214,7 +214,7 @@ impl<T> From<tokio::task::JoinHandle<T>> for AbortOnDropJoinHandle<T> {
 impl<T> Future for AbortOnDropJoinHandle<T> {
     type Output = T;
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match Pin::new(&mut self.as_mut().0).poll(cx) {
+        match Future::poll(Pin::new(&mut self.as_mut().0), cx) {
             Poll::Pending => Poll::Pending,
             Poll::Ready(r) => Poll::Ready(r.expect("child task panicked")),
         }