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

@@ -14,7 +14,7 @@ use std::str::FromStr;
 use wasmtime::component::{Resource, ResourceTable};
 use wasmtime_wasi::preview2::{
     bindings::io::streams::{InputStream, OutputStream},
-    Pollable,
+    PollableResource,
 };
 
 impl<T: WasiHttpView> crate::bindings::http::types::Host for T {
@@ -643,8 +643,8 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFutureTrailers for T {
     fn subscribe(
         &mut self,
         index: Resource<HostFutureTrailers>,
-    ) -> wasmtime::Result<Resource<Pollable>> {
-        wasmtime_wasi::preview2::subscribe(self.table(), index)
+    ) -> wasmtime::Result<Resource<PollableResource>> {
+        wasmtime_wasi::preview2::subscribe(self.table(), &index)
     }
 
     fn get(
@@ -851,8 +851,8 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFutureIncomingResponse f
     fn subscribe(
         &mut self,
         id: Resource<HostFutureIncomingResponse>,
-    ) -> wasmtime::Result<Resource<Pollable>> {
-        wasmtime_wasi::preview2::subscribe(self.table(), id)
+    ) -> wasmtime::Result<Resource<PollableResource>> {
+        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,11 +1,9 @@
-#![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::{Pollable, WasiView};
+use crate::preview2::poll;
+use crate::preview2::{Pollable, PollableResource, WasiView};
 use cap_std::time::SystemTime;
 use std::time::Duration;
 use wasmtime::component::Resource;
@@ -42,23 +40,18 @@ 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) -> Box<dyn Pollable> {
+    if duration.is_zero() {
+        Box::new(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))?
+        Box::new(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)
+        Box::new(poll::pending())
+    }
 }
 
 impl<T: WasiView> monotonic_clock::Host for T {
@@ -70,34 +63,22 @@ impl<T: WasiView> monotonic_clock::Host for T {
         Ok(self.ctx().monotonic_clock.resolution())
     }
 
-    fn subscribe_instant(&mut self, when: Instant) -> anyhow::Result<Resource<Pollable>> {
+    fn subscribe_instant(&mut self, when: Instant) -> anyhow::Result<Resource<PollableResource>> {
         let clock_now = self.ctx().monotonic_clock.now();
         let duration = if when > clock_now {
             Duration::from_nanos(when - clock_now)
         } else {
             Duration::from_nanos(0)
         };
-        subscribe_to_duration(&mut self.table(), duration)
-    }
-
-    fn subscribe_duration(&mut self, duration: WasiDuration) -> anyhow::Result<Resource<Pollable>> {
-        subscribe_to_duration(&mut self.table(), Duration::from_nanos(duration))
+        let pollable = make_pollable(duration);
+        Ok(self.table().push(PollableResource::new(pollable))?)
     }
-}
-
-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,
-        }
+    fn subscribe_duration(
+        &mut self,
+        duration: WasiDuration,
+    ) -> anyhow::Result<Resource<PollableResource>> {
+        let pollable = make_pollable(Duration::from_nanos(duration));
+        Ok(self.table().push(PollableResource::new(pollable))?)
     }
 }

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

@@ -2,7 +2,7 @@ use crate::preview2::{
     bindings::io::error,
     bindings::io::streams::{self, InputStream, OutputStream},
     poll::subscribe,
-    Pollable, StreamError, StreamResult, WasiView,
+    PollableResource, StreamError, StreamResult, WasiView,
 };
 use wasmtime::component::Resource;
 
@@ -48,8 +48,11 @@ impl<T: WasiView> streams::HostOutputStream for T {
         Ok(())
     }
 
-    fn subscribe(&mut self, stream: Resource<OutputStream>) -> anyhow::Result<Resource<Pollable>> {
-        subscribe(self.table(), stream)
+    fn subscribe(
+        &mut self,
+        stream: Resource<OutputStream>,
+    ) -> anyhow::Result<Resource<PollableResource>> {
+        subscribe(self.table(), &stream)
     }
 
     async fn blocking_write_and_flush(
@@ -219,8 +222,11 @@ impl<T: WasiView> streams::HostInputStream for T {
         self.skip(stream, len).await
     }
 
-    fn subscribe(&mut self, stream: Resource<InputStream>) -> anyhow::Result<Resource<Pollable>> {
-        crate::preview2::poll::subscribe(self.table(), stream)
+    fn subscribe(
+        &mut self,
+        stream: Resource<InputStream>,
+    ) -> anyhow::Result<Resource<PollableResource>> {
+        crate::preview2::poll::subscribe(self.table(), &stream)
     }
 }
 
@@ -230,9 +236,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, PollableResource, StreamError, StreamResult, WasiView,
     };
     use wasmtime::component::Resource;
 
@@ -290,7 +295,7 @@ pub mod sync {
         fn subscribe(
             &mut self,
             stream: Resource<OutputStream>,
-        ) -> anyhow::Result<Resource<Pollable>> {
+        ) -> anyhow::Result<Resource<PollableResource>> {
             Ok(AsyncHostOutputStream::subscribe(self, stream)?)
         }
 
@@ -359,7 +364,7 @@ pub mod sync {
         fn subscribe(
             &mut self,
             stream: Resource<InputStream>,
-        ) -> anyhow::Result<Resource<Pollable>> {
+        ) -> anyhow::Result<Resource<PollableResource>> {
             AsyncHostInputStream::subscribe(self, stream)
         }
     }

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

@@ -9,7 +9,7 @@ use crate::preview2::{
     },
     network::SocketAddressFamily,
 };
-use crate::preview2::{Pollable, SocketResult, WasiView};
+use crate::preview2::{PollableResource, SocketResult, WasiView};
 use cap_net_ext::Blocking;
 use io_lifetimes::AsSocketlike;
 use rustix::io::Errno;
@@ -559,8 +559,11 @@ impl<T: WasiView> crate::preview2::host::tcp::tcp::HostTcpSocket for T {
         Ok(())
     }
 
-    fn subscribe(&mut self, this: Resource<tcp::TcpSocket>) -> anyhow::Result<Resource<Pollable>> {
-        crate::preview2::poll::subscribe(self.table(), this)
+    fn subscribe(
+        &mut self,
+        this: Resource<tcp::TcpSocket>,
+    ) -> anyhow::Result<Resource<PollableResource>> {
+        crate::preview2::poll::subscribe(self.table(), &this)
     }
 
     fn shutdown(

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

@@ -8,7 +8,7 @@ use crate::preview2::{
     udp::{IncomingDatagramStream, OutgoingDatagramStream, SendState, UdpState},
     Subscribe,
 };
-use crate::preview2::{Pollable, SocketError, SocketResult, WasiView};
+use crate::preview2::{PollableResource, SocketError, SocketResult, WasiView};
 use anyhow::anyhow;
 use async_trait::async_trait;
 use io_lifetimes::AsSocketlike;
@@ -282,8 +282,11 @@ impl<T: WasiView> udp::HostUdpSocket for T {
         Ok(())
     }
 
-    fn subscribe(&mut self, this: Resource<udp::UdpSocket>) -> anyhow::Result<Resource<Pollable>> {
-        crate::preview2::poll::subscribe(self.table(), this)
+    fn subscribe(
+        &mut self,
+        this: Resource<udp::UdpSocket>,
+    ) -> anyhow::Result<Resource<PollableResource>> {
+        crate::preview2::poll::subscribe(self.table(), &this)
     }
 
     fn drop(&mut self, this: Resource<udp::UdpSocket>) -> Result<(), anyhow::Error> {
@@ -362,8 +365,8 @@ impl<T: WasiView> udp::HostIncomingDatagramStream for T {
     fn subscribe(
         &mut self,
         this: Resource<udp::IncomingDatagramStream>,
-    ) -> anyhow::Result<Resource<Pollable>> {
-        crate::preview2::poll::subscribe(self.table(), this)
+    ) -> anyhow::Result<Resource<PollableResource>> {
+        crate::preview2::poll::subscribe(self.table(), &this)
     }
 
     fn drop(&mut self, this: Resource<udp::IncomingDatagramStream>) -> Result<(), anyhow::Error> {
@@ -496,8 +499,8 @@ impl<T: WasiView> udp::HostOutgoingDatagramStream for T {
     fn subscribe(
         &mut self,
         this: Resource<udp::OutgoingDatagramStream>,
-    ) -> anyhow::Result<Resource<Pollable>> {
-        crate::preview2::poll::subscribe(self.table(), this)
+    ) -> anyhow::Result<Resource<PollableResource>> {
+        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

@@ -1,7 +1,7 @@
 use crate::preview2::bindings::sockets::ip_name_lookup::{Host, HostResolveAddressStream};
 use crate::preview2::bindings::sockets::network::{ErrorCode, IpAddress, Network};
 use crate::preview2::host::network::util;
-use crate::preview2::poll::{subscribe, Pollable, Subscribe};
+use crate::preview2::poll::{subscribe, PollableResource, Subscribe};
 use crate::preview2::{spawn_blocking, AbortOnDropJoinHandle, SocketError, WasiView};
 use anyhow::Result;
 use std::mem;
@@ -68,8 +68,8 @@ impl<T: WasiView> HostResolveAddressStream for T {
     fn subscribe(
         &mut self,
         resource: Resource<ResolveAddressStream>,
-    ) -> Result<Resource<Pollable>> {
-        subscribe(self.table(), resource)
+    ) -> Result<Resource<PollableResource>> {
+        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, PollableResource, Subscribe};
 pub use self::random::{thread_rng, Deterministic};
 pub use self::stdio::{
     stderr, stdin, stdout, IsATTY, Stderr, Stdin, StdinStream, Stdout, StdoutStream,
@@ -171,7 +171,7 @@ pub mod bindings {
             "wasi:io/streams/input-stream": super::stream::InputStream,
             "wasi:io/streams/output-stream": super::stream::OutputStream,
             "wasi:io/error/error": super::stream::Error,
-            "wasi:io/poll/pollable": super::poll::Pollable,
+            "wasi:io/poll/pollable": super::poll::PollableResource,
             "wasi:cli/terminal-input/terminal-input": super::stdio::TerminalInput,
             "wasi:cli/terminal-output/terminal-output": super::stdio::TerminalOutput,
         },