Revamp component model stream/future host API (again)

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

@@ -1,7 +1,7 @@
+use crate::util::PipeProducer;
 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, PipeProducer::new(rx)))
+        })
     }
 }
 

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

@@ -1,5 +1,18 @@
-use futures::channel::oneshot;
-use std::thread;
+use anyhow::Result;
+use futures::{Sink, Stream, channel::oneshot};
+use std::{
+    marker::PhantomData,
+    pin::Pin,
+    task::{Context, Poll},
+    thread,
+};
+use wasmtime::{
+    StoreContextMut,
+    component::{
+        Accessor, Destination, FutureConsumer, FutureProducer, Lift, Lower, Source, StreamConsumer,
+        StreamProducer, StreamResult,
+    },
+};
 
 pub async fn sleep(duration: std::time::Duration) {
     if cfg!(miri) {
@@ -21,3 +34,129 @@ pub async fn sleep(duration: std::time::Duration) {
         tokio::time::sleep(duration).await;
     }
 }
+
+pub struct PipeProducer<S>(S);
+
+impl<S> PipeProducer<S> {
+    pub fn new(rx: S) -> Self {
+        Self(rx)
+    }
+}
+
+impl<D, T: Send + Sync + Lower + 'static, S: Stream<Item = T> + Send + 'static> StreamProducer<D>
+    for PipeProducer<S>
+{
+    type Item = T;
+    type Buffer = Option<T>;
+
+    fn poll_produce<'a>(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        _: StoreContextMut<D>,
+        destination: &'a mut Destination<'a, Self::Item, Self::Buffer>,
+        finish: bool,
+    ) -> Poll<Result<StreamResult>> {
+        // SAFETY: This is a standard pin-projection, and we never move
+        // out of `self`.
+        let stream = unsafe { self.map_unchecked_mut(|v| &mut v.0) };
+
+        match stream.poll_next(cx) {
+            Poll::Pending => {
+                if finish {
+                    Poll::Ready(Ok(StreamResult::Cancelled))
+                } else {
+                    Poll::Pending
+                }
+            }
+            Poll::Ready(Some(item)) => {
+                destination.set_buffer(Some(item));
+                Poll::Ready(Ok(StreamResult::Completed))
+            }
+            Poll::Ready(None) => Poll::Ready(Ok(StreamResult::Dropped)),
+        }
+    }
+}
+
+pub struct PipeConsumer<T, S>(S, PhantomData<fn() -> T>);
+
+impl<T, S> PipeConsumer<T, S> {
+    pub fn new(tx: S) -> Self {
+        Self(tx, PhantomData)
+    }
+}
+
+impl<D, T: Lift + 'static, S: Sink<T, Error: std::error::Error + Send + Sync> + Send + 'static>
+    StreamConsumer<D> for PipeConsumer<T, S>
+{
+    type Item = T;
+
+    fn poll_consume(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        store: StoreContextMut<D>,
+        source: &mut Source<Self::Item>,
+        finish: bool,
+    ) -> Poll<Result<StreamResult>> {
+        // SAFETY: This is a standard pin-projection, and we never move
+        // out of `self`.
+        let mut sink = unsafe { self.map_unchecked_mut(|v| &mut v.0) };
+
+        let on_pending = || {
+            if finish {
+                Poll::Ready(Ok(StreamResult::Cancelled))
+            } else {
+                Poll::Pending
+            }
+        };
+
+        match sink.as_mut().poll_flush(cx) {
+            Poll::Pending => on_pending(),
+            Poll::Ready(result) => {
+                result?;
+                match sink.as_mut().poll_ready(cx) {
+                    Poll::Pending => on_pending(),
+                    Poll::Ready(result) => {
+                        result?;
+                        let item = &mut None;
+                        source.read(store, item)?;
+                        sink.start_send(item.take().unwrap())?;
+                        Poll::Ready(Ok(StreamResult::Completed))
+                    }
+                }
+            }
+        }
+    }
+}
+
+pub struct OneshotProducer<T>(oneshot::Receiver<T>);
+
+impl<T> OneshotProducer<T> {
+    pub fn new(rx: oneshot::Receiver<T>) -> Self {
+        Self(rx)
+    }
+}
+
+impl<D, T: Send + 'static> FutureProducer<D> for OneshotProducer<T> {
+    type Item = T;
+
+    async fn produce(self, _: &Accessor<D>) -> Result<T> {
+        Ok(self.0.await?)
+    }
+}
+
+pub struct OneshotConsumer<T>(oneshot::Sender<T>);
+
+impl<T> OneshotConsumer<T> {
+    pub fn new(tx: oneshot::Sender<T>) -> Self {
+        Self(tx)
+    }
+}
+
+impl<D, T: Send + 'static> FutureConsumer<D> for OneshotConsumer<T> {
+    type Item = T;
+
+    async fn consume(self, _: &Accessor<D>, value: T) -> Result<()> {
+        _ = self.0.send(value);
+        Ok(())
+    }
+}

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

@@ -236,7 +236,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