wasip3: Attempt cancellation in file ops (#13163)

This commit updates the file stream implementations to at least attempt
cancellation of blocking work instead of completely ignoring it. This
isn't guaranteed to work but at least enables exploring some possible
state spaces in wasi-libc for streams interacting with the host.

Author: alexcrichton

crates/wasi/src/p3/filesystem/host.rs

@@ -167,9 +167,7 @@ impl<D> StreamProducer<D> for ReadStreamProducer {
         cx: &mut Context<'_>,
         store: StoreContextMut<'a, D>,
         mut dst: Destination<'a, Self::Item, Self::Buffer>,
-        // Intentionally ignore this as in blocking mode everything is always
-        // ready and otherwise spawned blocking work can't be cancelled.
-        _finish: bool,
+        finish: bool,
     ) -> Poll<wasmtime::Result<StreamResult>> {
         if let Some(file) = self.file.as_blocking_file() {
             // Once a blocking file, always a blocking file, so assert as such.
@@ -213,22 +211,37 @@ impl<D> StreamProducer<D> for ReadStreamProducer {
         // Await the completion of the read task. Note that this is not a
         // cancellable await point because we can't cancel the other task, so
         // the `finish` parameter is ignored.
-        let res = ready!(Pin::new(task).poll(cx)).expect("I/O task should not panic");
+        let result = match Pin::new(&mut *task).poll(cx) {
+            // If cancellation is requested, then flag that to Tokio. Note that
+            // this still waits for the actual completion of the spawned task,
+            // which won't actually happen if it's already executing.
+            Poll::Pending if finish => {
+                task.abort();
+                ready!(Pin::new(task).poll(cx))
+            }
+            other => ready!(other),
+        };
         self.task = None;
-        match res {
-            Ok(buf) if buf.is_empty() => {
+        match result {
+            Ok(Ok(buf)) if buf.is_empty() => {
                 self.close(Ok(()));
                 Poll::Ready(Ok(StreamResult::Dropped))
             }
-            Ok(buf) => {
+            Ok(Ok(buf)) => {
                 let n = buf.len();
                 dst.set_buffer(Cursor::new(buf));
                 Poll::Ready(Ok(self.complete_read(n)))
             }
-            Err(err) => {
+            Ok(Err(err)) => {
                 self.close(Err(err.into()));
                 Poll::Ready(Ok(StreamResult::Dropped))
             }
+            Err(err) => {
+                if err.is_cancelled() {
+                    return Poll::Ready(Ok(StreamResult::Cancelled));
+                }
+                panic!("I/O task should not panic: {err}")
+            }
         }
     }
 }
@@ -434,9 +447,7 @@ impl<D> StreamConsumer<D> for WriteStreamConsumer {
         cx: &mut Context<'_>,
         store: StoreContextMut<D>,
         src: Source<Self::Item>,
-        // Intentionally ignore this as in blocking mode everything is always
-        // ready and otherwise spawned blocking work can't be cancelled.
-        _finish: bool,
+        finish: bool,
     ) -> Poll<wasmtime::Result<StreamResult>> {
         let mut src = src.as_direct(store);
         if let Some(file) = self.file.as_blocking_file() {
@@ -462,19 +473,34 @@ impl<D> StreamConsumer<D> for WriteStreamConsumer {
             let location = me.location;
             spawn_blocking(move || location.write(&file, &buf).map(|n| (buf, n)))
         });
-        let res = ready!(Pin::new(task).poll(cx)).expect("I/O task should not panic");
+        let result = match Pin::new(&mut *task).poll(cx) {
+            // If cancellation is requested, then flag that to Tokio. Note that
+            // this still waits for the actual completion of the spawned task,
+            // which won't actually happen if it's already executing.
+            Poll::Pending if finish => {
+                task.abort();
+                ready!(Pin::new(task).poll(cx))
+            }
+            other => ready!(other),
+        };
         self.task = None;
-        match res {
-            Ok((buf, n)) => {
+        match result {
+            Ok(Ok((buf, n))) => {
                 src.mark_read(n);
                 self.buffer = buf;
                 self.buffer.clear();
                 Poll::Ready(Ok(self.complete_write(n)))
             }
-            Err(err) => {
+            Ok(Err(err)) => {
                 self.close(Err(err.into()));
                 Poll::Ready(Ok(StreamResult::Dropped))
             }
+            Err(err) => {
+                if err.is_cancelled() {
+                    return Poll::Ready(Ok(StreamResult::Cancelled));
+                }
+                panic!("I/O task should not panic: {err}")
+            }
         }
     }
 }