[37.0.0] Backport WASIp3-related PRs (#11624)

* make `Destination::as_direct` work for both host and guest readers (#11612)

In order to reduce code duplication (and code paths to test) in `wasmtime-wasi`
and custom host embeddings, I've made `Destination::as_direct` (formerly known
as `as_direct_destination`) work for host readers as well as guest ones.  In the
process, I noticed and fixed a couple of related issues:

- I had forgotten to implement or test host reader support in `DirectSource` 🤦
- The code to support host-to-host pipes failed to account for partial reads

I've also simplified the `StreamConsumer` and `StreamProducer` APIs slightly by
having them take their `Source` and `Destination` parameters by value rather
than by reference, respectively.

Note that, per WebAssembly/component-model#561, I've
tweaked the documentation for `StreamProducer` to indicate that implementations
might reasonably opt to "pretend" they're ready without buffering any items when
handling zero-length reads given that buffering has its own hazards.  Likewise,
I've updated the `wasi-filesystem` and `wasi-cli` implementations to "pretend"
instead of buffering.

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

* wasip3: Refine 0-length behavior in cli/sockets (#11614)

Add some comments for issues and `if` blocks in a few places. I'd like
to add some tests but I ran into #11611 so I'll defer tests to later.

* wasip3: Refactor the readdir stream iterators (#11615)

* wasip3: Refactor `BlockingDirectoryStreamProducer`

* Generalize this into a `FallibleStreamProducer` structure
* Don't read the entire iterator on the first call to `poll_produce`
* Do a blocking read of `dir.entries()` in the original function call to
  avoid handling state in the iterator itself.

* wasip3: Refactor `NonblockingDirectoryStreamProducer`

* Start the reading task before iteration starts to move the spawn out
  of the `poll_*` method.
* Rely on fusing behavior of mpsc/tasks to avoid extra state structure.
* Specifically handle 0-length reads.

Mostly try to refactor the state representation to be more struct-like
rather than enum like which is a little easier to follow.

* wasip3: Port `preview1_fd_readdir` to WASIp3

Have at least one test looking at the readdir behavior.

* Fix `finish` handling in `FallibleIteratorProducer`

* Fix a typo

---------

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

Author: alexcrichton

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

@@ -53,7 +53,7 @@ impl<D, T: Send + Sync + Lower + 'static, S: Stream<Item = T> + Send + 'static>
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         _: StoreContextMut<D>,
-        destination: &'a mut Destination<'a, Self::Item, Self::Buffer>,
+        mut destination: Destination<'a, Self::Item, Self::Buffer>,
         finish: bool,
     ) -> Poll<Result<StreamResult>> {
         // SAFETY: This is a standard pin-projection, and we never move
@@ -94,7 +94,7 @@ impl<D, T: Lift + 'static, S: Sink<T, Error: std::error::Error + Send + Sync> +
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         store: StoreContextMut<D>,
-        source: &mut Source<Self::Item>,
+        mut source: Source<Self::Item>,
         finish: bool,
     ) -> Poll<Result<StreamResult>> {
         // SAFETY: This is a standard pin-projection, and we never move

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

@@ -6,18 +6,107 @@ use {
         util::{OneshotConsumer, OneshotProducer, PipeConsumer, PipeProducer},
     },
     futures::{
-        SinkExt, StreamExt,
+        Sink, SinkExt, Stream, StreamExt,
         channel::{mpsc, oneshot},
         future,
     },
-    std::sync::{Arc, Mutex},
+    std::{
+        pin::Pin,
+        sync::{Arc, Mutex},
+        task::{Context, Poll},
+    },
     wasmtime::{
-        Engine, Store,
-        component::{FutureReader, Linker, ResourceTable, StreamReader},
+        Engine, Store, StoreContextMut,
+        component::{
+            Destination, FutureReader, Linker, ResourceTable, Source, StreamConsumer,
+            StreamProducer, StreamReader, StreamResult,
+        },
     },
     wasmtime_wasi::WasiCtxBuilder,
 };
 
+pub struct DirectPipeProducer<S>(S);
+
+impl<D, S: Stream<Item = u8> + Send + 'static> StreamProducer<D> for DirectPipeProducer<S> {
+    type Item = u8;
+    type Buffer = Option<u8>;
+
+    fn poll_produce<'a>(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        store: StoreContextMut<D>,
+        destination: 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)) => {
+                let mut destination = destination.as_direct(store, 1);
+                destination.remaining()[0] = item;
+                destination.mark_written(1);
+                Poll::Ready(Ok(StreamResult::Completed))
+            }
+            Poll::Ready(None) => Poll::Ready(Ok(StreamResult::Dropped)),
+        }
+    }
+}
+
+pub struct DirectPipeConsumer<S>(S);
+
+impl<D, S: Sink<u8, Error: std::error::Error + Send + Sync> + Send + 'static> StreamConsumer<D>
+    for DirectPipeConsumer<S>
+{
+    type Item = u8;
+
+    fn poll_consume(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        store: StoreContextMut<D>,
+        source: 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 mut source = source.as_direct(store);
+                        let item = source.remaining()[0];
+                        source.mark_read(1);
+                        sink.start_send(item)?;
+                        Poll::Ready(Ok(StreamResult::Completed))
+                    }
+                }
+            }
+        }
+    }
+}
+
 #[tokio::test]
 pub async fn async_closed_streams() -> Result<()> {
     let engine = Engine::new(&config())?;
@@ -49,35 +138,45 @@ pub async fn async_closed_streams() -> Result<()> {
     let value = 42_u8;
 
     // First, test stream host->host
-    {
-        let (mut input_tx, input_rx) = mpsc::channel(1);
-        let (output_tx, mut output_rx) = mpsc::channel(1);
-        StreamReader::new(instance, &mut store, PipeProducer::new(input_rx))
-            .pipe(&mut store, PipeConsumer::new(output_tx));
-
-        instance
-            .run_concurrent(&mut store, async |_| {
-                let (a, b) = future::join(
-                    async {
-                        for &value in &values {
-                            input_tx.send(value).await?;
-                        }
-                        drop(input_tx);
-                        anyhow::Ok(())
-                    },
-                    async {
-                        for &value in &values {
-                            assert_eq!(Some(value), output_rx.next().await);
-                        }
-                        assert!(output_rx.next().await.is_none());
-                        Ok(())
-                    },
-                )
-                .await;
-
-                a.and(b)
-            })
-            .await??;
+    for direct_producer in [true, false] {
+        for direct_consumer in [true, false] {
+            let (mut input_tx, input_rx) = mpsc::channel(1);
+            let (output_tx, mut output_rx) = mpsc::channel(1);
+            let reader = if direct_producer {
+                StreamReader::new(instance, &mut store, DirectPipeProducer(input_rx))
+            } else {
+                StreamReader::new(instance, &mut store, PipeProducer::new(input_rx))
+            };
+            if direct_consumer {
+                reader.pipe(&mut store, DirectPipeConsumer(output_tx));
+            } else {
+                reader.pipe(&mut store, PipeConsumer::new(output_tx));
+            }
+
+            instance
+                .run_concurrent(&mut store, async |_| {
+                    let (a, b) = future::join(
+                        async {
+                            for &value in &values {
+                                input_tx.send(value).await?;
+                            }
+                            drop(input_tx);
+                            anyhow::Ok(())
+                        },
+                        async {
+                            for &value in &values {
+                                assert_eq!(Some(value), output_rx.next().await);
+                            }
+                            assert!(output_rx.next().await.is_none());
+                            Ok(())
+                        },
+                    )
+                    .await;
+
+                    a.and(b)
+                })
+                .await??;
+        }
     }
 
     // Next, test futures host->host

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

@@ -42,7 +42,7 @@ impl<D> StreamProducer<D> for ReadinessProducer {
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         mut store: StoreContextMut<'a, D>,
-        destination: &'a mut Destination<'a, Self::Item, Self::Buffer>,
+        destination: Destination<'a, Self::Item, Self::Buffer>,
         finish: bool,
     ) -> Poll<Result<StreamResult>> {
         let me = self.get_mut();
@@ -62,7 +62,7 @@ impl<D> StreamProducer<D> for ReadinessProducer {
                     Poll::Ready(Ok(StreamResult::Completed))
                 } else {
                     assert_eq!(capacity, Some(me.buffer.len()));
-                    let mut destination = destination.as_direct_destination(store).unwrap();
+                    let mut destination = destination.as_direct(store, me.buffer.len());
                     destination.remaining().copy_from_slice(&me.buffer);
                     destination.mark_written(me.buffer.len());
 
@@ -85,7 +85,7 @@ impl<D> StreamConsumer<D> for ReadinessConsumer {
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         mut store: StoreContextMut<D>,
-        source: &mut Source<Self::Item>,
+        source: Source<Self::Item>,
         finish: bool,
     ) -> Poll<Result<StreamResult>> {
         let me = self.get_mut();
@@ -105,7 +105,7 @@ impl<D> StreamConsumer<D> for ReadinessConsumer {
                     Poll::Ready(Ok(StreamResult::Completed))
                 } else {
                     assert_eq!(available, me.expected.len());
-                    let mut source = source.as_direct_source(store);
+                    let mut source = source.as_direct(store);
                     assert_eq!(&me.expected, source.remaining());
                     source.mark_read(me.expected.len());
 

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

@@ -0,0 +1,95 @@
+use test_programs::p3::wasi;
+use test_programs::p3::wasi::filesystem::types::{
+    Descriptor, DescriptorFlags, DescriptorType, DirectoryEntry, OpenFlags, PathFlags,
+};
+
+struct Component;
+
+test_programs::p3::export!(Component);
+
+impl test_programs::p3::exports::wasi::cli::run::Guest for Component {
+    async fn run() -> Result<(), ()> {
+        let preopens = wasi::filesystem::preopens::get_directories();
+        let (dir, _) = &preopens[0];
+
+        test_readdir(dir).await;
+        test_readdir_lots(dir).await;
+        Ok(())
+    }
+}
+
+fn main() {
+    unreachable!()
+}
+
+async fn read_dir(dir: &Descriptor) -> Vec<DirectoryEntry> {
+    let (dirs, result) = dir.read_directory().await;
+    let mut dirs = dirs.collect().await;
+    result.await.unwrap();
+    dirs.sort_by_key(|d| d.name.clone());
+    dirs
+}
+
+async fn assert_empty_dir(dir: &Descriptor) {
+    let dirs = read_dir(dir).await;
+    assert_eq!(dirs.len(), 0);
+}
+
+async fn test_readdir(dir: &Descriptor) {
+    // Check the behavior in an empty directory
+    assert_empty_dir(dir).await;
+
+    dir.open_at(
+        PathFlags::empty(),
+        "file".to_string(),
+        OpenFlags::CREATE,
+        DescriptorFlags::READ | DescriptorFlags::WRITE,
+    )
+    .await
+    .unwrap();
+
+    dir.create_directory_at("nested".to_string()).await.unwrap();
+    let nested = dir
+        .open_at(
+            PathFlags::empty(),
+            "nested".to_string(),
+            OpenFlags::DIRECTORY,
+            DescriptorFlags::empty(),
+        )
+        .await
+        .unwrap();
+
+    let entries = read_dir(dir).await;
+    assert_eq!(entries.len(), 2);
+    assert_eq!(entries[0].name, "file");
+    assert_eq!(entries[0].type_, DescriptorType::RegularFile);
+    assert_eq!(entries[1].name, "nested");
+    assert_eq!(entries[1].type_, DescriptorType::Directory);
+
+    assert_empty_dir(&nested).await;
+    drop(nested);
+
+    dir.unlink_file_at("file".to_string()).await.unwrap();
+    dir.remove_directory_at("nested".to_string()).await.unwrap();
+}
+
+async fn test_readdir_lots(dir: &Descriptor) {
+    for count in 0..1000 {
+        dir.open_at(
+            PathFlags::empty(),
+            format!("file.{count}"),
+            OpenFlags::CREATE,
+            DescriptorFlags::READ | DescriptorFlags::WRITE,
+        )
+        .await
+        .expect("failed to create file");
+    }
+
+    assert_eq!(read_dir(dir).await.len(), 1000);
+
+    for count in 0..1000 {
+        dir.unlink_file_at(format!("file.{count}"))
+            .await
+            .expect("removing a file");
+    }
+}