Revamp component model stream/future host API (again)

This changes the host APIs for dealing with futures and streams from a
"rendezvous"-style API to a callback-oriented one.

Previously you would create e.g. a `StreamReader`/`StreamWriter` pair and call
their `read` and `write` methods, respectively, and those methods would return
`Future`s that resolved when the operation was matched with a corresponding
`write` or `read` operation on the other end.

With the new API, you instead provide a `StreamProducer` trait implementation
whe creating the stream, whose `produce` method will be called as soon as a read
happens, giving the implementation a chance to respond immediately without
making the reader wait for a rendezvous.  Likewise, you can match the read end
of a stream to a `StreamConsumer` to respond immediately to writes.  This model
should reduce scheduling overhead and make it easier to e.g. pipe items to/from
`AsyncWrite`/`AsyncRead` or `Sink`/`Stream` implementations without needing to
explicitly spawn background tasks.  In addition, the new API provides direct
access to guest read and write buffers for `stream<u8>` operations, enabling
zero-copy operations.

Other changes:

- I've removed the `HostTaskOutput`; we were using it to run extra code with
  access to the store after a host task completes, but we can do that more
  elegantly inside the future using `tls::get`.  This also allowed me to
  simplify `Instance::poll_until` a bit.

- I've removed the `watch_{reader,writer}` functionality; it's not needed now
  given that the runtime will automatically dispose of the producer or consumer
  when the other end of the stream or future is closed -- no need for embedder
  code to manage that.

- In order to make `UntypedWriteBuffer` `Send`, I had to wrap its raw pointer
  `buf` field in a `SendSyncPtr`.

- I've removed `{Future,Stream}Writer` entirely and moved
  `Instance::{future,stream}` to `{Future,Stream}Reader::new`, respectively.

- I've added a bounds check to the beginnings of `Instance::guest_read` and
  `Instance::guest_write` so that we need not do it later in
  `Guest{Source,Destination}::remaining`, meaning those functions can be
  infallible.

Note that I haven't updated `wasmtime-wasi` yet to match; that will happen in
one or more follow-up commits.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

Author: dicej

crates/misc/component-async-tests/src/resource_stream.rs

@@ -1,7 +1,7 @@
+use crate::util::MpscProducer;
 use anyhow::Result;
-use wasmtime::component::{
-    Accessor, AccessorTask, GuardedStreamWriter, Resource, StreamReader, StreamWriter,
-};
+use futures::channel::mpsc;
+use wasmtime::component::{Accessor, Resource, StreamReader};
 
 use super::Ctx;
 
@@ -38,29 +38,15 @@ impl bindings::local::local::resource_stream::HostWithStore for Ctx {
         accessor: &Accessor<T, Self>,
         count: u32,
     ) -> wasmtime::Result<StreamReader<Resource<ResourceStreamX>>> {
-        struct Task {
-            tx: StreamWriter<Resource<ResourceStreamX>>,
-
-            count: u32,
-        }
-
-        impl<T> AccessorTask<T, Ctx, Result<()>> for Task {
-            async fn run(self, accessor: &Accessor<T, Ctx>) -> Result<()> {
-                let mut tx = GuardedStreamWriter::new(accessor, self.tx);
-                for _ in 0..self.count {
-                    let item = accessor.with(|mut view| view.get().table.push(ResourceStreamX))?;
-                    tx.write_all(Some(item)).await;
-                }
-                Ok(())
+        accessor.with(|mut access| {
+            let (mut tx, rx) = mpsc::channel(usize::try_from(count).unwrap());
+            for _ in 0..count {
+                tx.try_send(access.get().table.push(ResourceStreamX)?)
+                    .unwrap()
             }
-        }
-
-        let (tx, rx) = accessor.with(|mut view| {
-            let instance = view.instance();
-            instance.stream(&mut view)
-        })?;
-        accessor.spawn(Task { tx, count });
-        Ok(rx)
+            let instance = access.instance();
+            Ok(StreamReader::new(instance, access, MpscProducer::new(rx)))
+        })
     }
 }
 

crates/misc/component-async-tests/src/util.rs

@@ -1,5 +1,13 @@
-use futures::channel::oneshot;
+use anyhow::Result;
+use futures::{
+    SinkExt, StreamExt,
+    channel::{mpsc, oneshot},
+};
 use std::thread;
+use wasmtime::component::{
+    Accessor, Destination, FutureConsumer, FutureProducer, Lift, Lower, Source, StreamConsumer,
+    StreamProducer, StreamState,
+};
 
 pub async fn sleep(duration: std::time::Duration) {
     if cfg!(miri) {
@@ -21,3 +29,111 @@ pub async fn sleep(duration: std::time::Duration) {
         tokio::time::sleep(duration).await;
     }
 }
+
+pub struct MpscProducer<T> {
+    rx: mpsc::Receiver<T>,
+    closed: bool,
+}
+
+impl<T: Send + Sync + 'static> MpscProducer<T> {
+    pub fn new(rx: mpsc::Receiver<T>) -> Self {
+        Self { rx, closed: false }
+    }
+
+    fn state(&self) -> StreamState {
+        if self.closed {
+            StreamState::Closed
+        } else {
+            StreamState::Open
+        }
+    }
+}
+
+impl<D, T: Send + Sync + Lower + 'static> StreamProducer<D, T> for MpscProducer<T> {
+    async fn produce(
+        &mut self,
+        accessor: &Accessor<D>,
+        destination: &mut Destination<T>,
+    ) -> Result<StreamState> {
+        if let Some(item) = self.rx.next().await {
+            let item = destination.write(accessor, Some(item)).await?;
+            assert!(item.is_none());
+        } else {
+            self.closed = true;
+        }
+
+        Ok(self.state())
+    }
+
+    async fn when_ready(&mut self, _: &Accessor<D>) -> Result<StreamState> {
+        Ok(self.state())
+    }
+}
+
+pub struct MpscConsumer<T> {
+    tx: mpsc::Sender<T>,
+}
+
+impl<T> MpscConsumer<T> {
+    pub fn new(tx: mpsc::Sender<T>) -> Self {
+        Self { tx }
+    }
+
+    fn state(&self) -> StreamState {
+        if self.tx.is_closed() {
+            StreamState::Closed
+        } else {
+            StreamState::Open
+        }
+    }
+}
+
+impl<D, T: Lift + 'static> StreamConsumer<D, T> for MpscConsumer<T> {
+    async fn consume(
+        &mut self,
+        accessor: &Accessor<D>,
+        source: &mut Source<'_, T>,
+    ) -> Result<StreamState> {
+        let item = &mut None;
+        accessor.with(|access| source.read(access, item))?;
+        _ = self.tx.send(item.take().unwrap()).await;
+        Ok(self.state())
+    }
+
+    async fn when_ready(&mut self, _: &Accessor<D>) -> Result<StreamState> {
+        Ok(self.state())
+    }
+}
+
+pub struct OneshotProducer<T> {
+    rx: oneshot::Receiver<T>,
+}
+
+impl<T> OneshotProducer<T> {
+    pub fn new(rx: oneshot::Receiver<T>) -> Self {
+        Self { rx }
+    }
+}
+
+impl<D, T: Send + 'static> FutureProducer<D, T> for OneshotProducer<T> {
+    async fn produce(self, _: &Accessor<D>) -> Result<T> {
+        Ok(self.rx.await?)
+    }
+}
+
+pub struct OneshotConsumer<T> {
+    tx: oneshot::Sender<T>,
+}
+
+impl<T> OneshotConsumer<T> {
+    pub fn new(tx: oneshot::Sender<T>) -> Self {
+        Self { tx }
+    }
+}
+
+impl<D, T: Send + 'static> FutureConsumer<D, T> for OneshotConsumer<T> {
+    async fn consume(self, _: &Accessor<D>, value: T) -> Result<()> {
+        _ = self.tx.send(value);
+        Ok(())
+    }
+}

crates/misc/component-async-tests/tests/scenario/round_trip.rs

@@ -237,7 +237,7 @@ pub async fn test_round_trip(
             component_async_tests::round_trip::bindings::RoundTrip::new(&mut store, &instance)?;
 
         if call_style == 0 || !cfg!(miri) {
-            // Now do it again using `Instance::run_concurrent`:
+            // Run the test using `Instance::run_concurrent`:
             instance
                 .run_concurrent(&mut store, {
                     let inputs_and_outputs = inputs_and_outputs