Rename join_background_tasks to flush_output

This builds on the prior commit to redefine the `join_background_tasks`
as purely a "flush" operation that flushes all output streams. This no
longer attempts to gracefully join output tasks one-by-one since that
should no longer be necessary and everything will get cleaned up during
drop when `abort` calls are invoked.

Author: alexcrichton

crates/wasi/src/preview2/ctx.rs

@@ -2,14 +2,11 @@ use super::clocks::host::{monotonic_clock, wall_clock};
 use crate::preview2::{
     clocks::{self, HostMonotonicClock, HostWallClock},
     filesystem::{Dir, TableFsExt},
-    pipe,
-    poll::TablePollableExt,
-    random, stdio,
+    pipe, random, stdio,
     stream::{HostInputStream, HostOutputStream, TableStreamExt},
     DirPerms, FilePerms, Table,
 };
 use cap_rand::{Rng, RngCore, SeedableRng};
-use std::future::Future;
 use std::mem;
 
 pub struct WasiCtxBuilder {
@@ -263,8 +260,7 @@ pub struct WasiCtx {
 }
 
 impl WasiCtx {
-    /// Wait for all background tasks to join (complete) gracefully, after flushing any
-    /// buffered output.
+    /// Flush all buffered output and wait for it to reach the destination.
     ///
     /// NOTE: This function should be used when [`WasiCtx`] is used in an async embedding
     /// (i.e. with [`crate::preview2::command::add_to_linker`]). Use its counterpart
@@ -274,92 +270,40 @@ impl WasiCtx {
     /// In order to implement non-blocking streams, we often often need to offload async
     /// operations to background `tokio::task`s. These tasks are aborted when the resources
     /// in the `Table` referencing them are dropped. In some cases, this abort may occur before
-    /// buffered output has been flushed. Use this function to wait for all background tasks to
-    /// join gracefully.
+    /// buffered output has been flushed. Use this function to wait for all
+    /// written data to reach its destination gracefully, even if wasm didn't
+    /// explicitly wait for this.
     ///
     /// In some embeddings, a misbehaving client might cause this graceful exit to await for an
     /// unbounded amount of time, so we recommend bounding this with a timeout or other mechanism.
-    pub fn join_background_tasks<'a>(&mut self, table: &'a mut Table) -> impl Future<Output = ()> {
-        use std::pin::Pin;
-        use std::task::{Context, Poll};
+    pub async fn flush_output(&mut self, table: &mut Table) {
         let keys = table.keys().cloned().collect::<Vec<u32>>();
-        let mut set = Vec::new();
-        // we can't remove an stream from the table if it has any child pollables,
-        // so first delete all pollables from the table.
-        for k in keys.iter() {
-            let _ = table.delete_host_pollable(*k);
-        }
         for k in keys {
             match table.delete_output_stream(k) {
                 Ok(mut ostream) => {
-                    // async block takes ownership of the ostream and flushes it
-                    let f = async move { ostream.join_background_tasks().await };
-                    set.push(Box::pin(f) as _)
-                }
-                _ => {}
-            }
-            match table.delete_input_stream(k) {
-                Ok(mut istream) => {
-                    // async block takes ownership of the istream and flushes it
-                    let f = async move { istream.join_background_tasks().await };
-                    set.push(Box::pin(f) as _)
+                    let _ = ostream.flush().await;
                 }
                 _ => {}
             }
         }
-        // poll futures until all are ready.
-        // We can't write this as an `async fn` because we want to eagerly poll on each possible
-        // join, rather than sequentially awaiting on them.
-        struct JoinAll(Vec<Pin<Box<dyn Future<Output = ()> + Send>>>);
-        impl Future for JoinAll {
-            type Output = ();
-            fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
-                // Iterate through the set, polling each future and removing it from the set if it
-                // is ready:
-                self.as_mut()
-                    .0
-                    .retain_mut(|fut| match fut.as_mut().poll(cx) {
-                        Poll::Ready(_) => false,
-                        _ => true,
-                    });
-                // Ready if set is empty:
-                if self.as_mut().0.is_empty() {
-                    Poll::Ready(())
-                } else {
-                    Poll::Pending
-                }
-            }
-        }
-
-        JoinAll(set)
     }
 
-    /// Wait for all background tasks to join (complete) gracefully, after flushing any
-    /// buffered output.
-    ///
-    /// NOTE: This function should be used when [`WasiCtx`] is used in an synchronous embedding
-    /// (i.e. with [`crate::preview2::command::sync::add_to_linker`]). Use its counterpart
-    /// `join_background_tasks` in an async embedding (i.e. with
-    /// [`crate::preview2::command::add_to_linker`].
+    /// Same as [`WasiCtx::flush_output`] except suitable for synchronous
+    /// embeddings.
     ///
-    /// In order to implement non-blocking streams, we often often need to offload async
-    /// operations to background `tokio::task`s. These tasks are aborted when the resources
-    /// in the `Table` referencing them are dropped. In some cases, this abort may occur before
-    /// buffered output has been flushed. Use this function to wait for all background tasks to
-    /// join gracefully.
-    ///
-    /// In some embeddings, a misbehaving client might cause this graceful exit to await for an
-    /// unbounded amount of time, so we recommend providing a timeout for this method.
+    /// This will block the current thread up to the `timeout` specified, or
+    /// forever if `None` is specified, until all wasm output has been flushed
+    /// out.
     pub fn sync_join_background_tasks<'a>(
         &mut self,
         table: &'a mut Table,
         timeout: Option<std::time::Duration>,
     ) {
         crate::preview2::in_tokio(async move {
             if let Some(timeout) = timeout {
-                let _ = tokio::time::timeout(timeout, self.join_background_tasks(table)).await;
+                let _ = tokio::time::timeout(timeout, self.flush_output(table)).await;
             } else {
-                self.join_background_tasks(table).await
+                self.flush_output(table).await
             }
         })
     }

crates/wasi/src/preview2/pipe.rs

@@ -87,6 +87,10 @@ impl HostOutputStream for MemoryOutputPipe {
         // This stream is always ready for writing.
         Ok(())
     }
+
+    async fn flush(&mut self) -> Result<(), Error> {
+        Ok(())
+    }
 }
 
 /// TODO
@@ -105,9 +109,7 @@ pub struct AsyncReadStream {
     state: StreamState,
     buffer: Option<Result<Bytes, std::io::Error>>,
     receiver: tokio::sync::mpsc::Receiver<Result<(Bytes, StreamState), std::io::Error>>,
-    // the join handle for the background task is Some until join_background_tasks, after which
-    // further use of the AsyncReadStream is not allowed.
-    join_handle: Option<tokio::task::JoinHandle<()>>,
+    join_handle: tokio::task::JoinHandle<()>,
 }
 
 impl AsyncReadStream {
@@ -136,17 +138,13 @@ impl AsyncReadStream {
             state: StreamState::Open,
             buffer: None,
             receiver,
-            join_handle: Some(join_handle),
+            join_handle,
         }
     }
     // stdio implementation uses this to determine if the backing tokio runtime has been shutdown and
     // restarted:
     pub(crate) fn is_finished(&self) -> bool {
-        assert!(
-            self.join_handle.is_some(),
-            "illegal use of AsyncReadStream after join_background_tasks"
-        );
-        self.join_handle.as_ref().unwrap().is_finished()
+        self.join_handle.is_finished()
     }
 }
 
@@ -155,18 +153,14 @@ impl AsyncReadStream {
 // on reader.read_buf's await it could hold the reader open indefinitely.
 impl Drop for AsyncReadStream {
     fn drop(&mut self) {
-        self.join_handle.take().map(|h| h.abort());
+        self.join_handle.abort();
     }
 }
 
 #[async_trait::async_trait]
 impl HostInputStream for AsyncReadStream {
     fn read(&mut self, size: usize) -> Result<(Bytes, StreamState), Error> {
         use tokio::sync::mpsc::error::TryRecvError;
-        assert!(
-            self.join_handle.is_some(),
-            "illegal use of AsyncReadStream after join_background_tasks"
-        );
 
         match self.buffer.take() {
             Some(Ok(mut bytes)) => {
@@ -209,11 +203,6 @@ impl HostInputStream for AsyncReadStream {
     }
 
     async fn ready(&mut self) -> Result<(), Error> {
-        assert!(
-            self.join_handle.is_some(),
-            "illegal use of AsyncReadStream after join_background_tasks"
-        );
-
         if self.buffer.is_some() || self.state == StreamState::Closed {
             return Ok(());
         }
@@ -233,9 +222,6 @@ impl HostInputStream for AsyncReadStream {
         }
         Ok(())
     }
-    async fn join_background_tasks(&mut self) {
-        self.join_handle.take().map(|h| h.abort());
-    }
 }
 
 #[derive(Debug)]
@@ -348,40 +334,6 @@ impl AsyncWriteStream {
     fn has_pending_op(&self) -> bool {
         matches!(self.state, Some(WriteState::Pending))
     }
-
-    async fn flush(&mut self) -> anyhow::Result<()> {
-        // NB: This method needs to be "cancel safe" where it can be cancelled
-        // at any `.await` point but the flush operation still needs to be able
-        // to be restarted successfully.
-
-        // First wait for any pending operation to complete to have the ability
-        // to send another message.
-        self.ready().await?;
-
-        // Queue up a flush operation in our background task, and if it's
-        // already gone then that's ok as flushing has completed anyway.
-        //
-        // Note that this may end up returning a queued error from a previous
-        // write or flush.
-        if !self.send(WriteMesage::Flush)? {
-            return Ok(());
-        }
-
-        // Wait again for the flush to fully complete before considering this
-        // request to flush as fully complete.
-        self.ready().await?;
-
-        // Extract the error, if any, that occurred.
-        match mem::replace(&mut self.state, Some(WriteState::Ready)) {
-            Some(WriteState::Err(e)) => Err(e.into()),
-            Some(WriteState::Pending) => unreachable!(),
-            Some(WriteState::Ready) => Ok(()),
-            None => {
-                self.state = None;
-                Ok(())
-            }
-        }
-    }
 }
 
 // Make sure the background task does not outlive the AsyncWriteStream handle.
@@ -429,8 +381,38 @@ impl HostOutputStream for AsyncWriteStream {
         Ok(())
     }
 
-    async fn join_background_tasks(&mut self) {
-        let _ = self.flush().await;
+    async fn flush(&mut self) -> anyhow::Result<()> {
+        // NB: This method needs to be "cancel safe" where it can be cancelled
+        // at any `.await` point but the flush operation still needs to be able
+        // to be restarted successfully.
+
+        // First wait for any pending operation to complete to have the ability
+        // to send another message.
+        self.ready().await?;
+
+        // Queue up a flush operation in our background task, and if it's
+        // already gone then that's ok as flushing has completed anyway.
+        //
+        // Note that this may end up returning a queued error from a previous
+        // write or flush.
+        if !self.send(WriteMesage::Flush)? {
+            return Ok(());
+        }
+
+        // Wait again for the flush to fully complete before considering this
+        // request to flush as fully complete.
+        self.ready().await?;
+
+        // Extract the error, if any, that occurred.
+        match mem::replace(&mut self.state, Some(WriteState::Ready)) {
+            Some(WriteState::Err(e)) => Err(e.into()),
+            Some(WriteState::Pending) => unreachable!(),
+            Some(WriteState::Ready) => Ok(()),
+            None => {
+                self.state = None;
+                Ok(())
+            }
+        }
     }
 }
 
@@ -446,6 +428,10 @@ impl HostOutputStream for SinkOutputStream {
     async fn ready(&mut self) -> Result<(), Error> {
         Ok(())
     }
+
+    async fn flush(&mut self) -> Result<(), Error> {
+        Ok(())
+    }
 }
 
 /// A stream that is ready immediately, but will always report that it's closed.
@@ -474,6 +460,10 @@ impl HostOutputStream for ClosedOutputStream {
     async fn ready(&mut self) -> Result<(), Error> {
         Ok(())
     }
+
+    async fn flush(&mut self) -> Result<(), Error> {
+        Ok(())
+    }
 }
 
 #[cfg(test)]

crates/wasi/src/preview2/stream.rs

@@ -43,11 +43,6 @@ pub trait HostInputStream: Send + Sync {
     /// Check for read readiness: this method blocks until the stream is ready
     /// for reading.
     async fn ready(&mut self) -> Result<(), Error>;
-
-    /// Terminate all background tasks. Exposed only to the host, not accessible from WebAssembly.
-    /// Must cancel background tasks even if dropped before completion.
-    /// No other methods may be used after calling this method.
-    async fn join_background_tasks(&mut self) {}
 }
 
 /// Host trait for implementing the `wasi:io/streams.output-stream` resource:
@@ -94,10 +89,11 @@ pub trait HostOutputStream: Send + Sync {
     /// ready for writing.
     async fn ready(&mut self) -> Result<(), Error>;
 
-    /// Flush any output which has been buffered, and terminate all background tasks. Exposed only
-    /// to the host, not accessible from WebAssembly. Must cancel background tasks even if dropped
-    /// before completion. No other methods may be used after calling this method.
-    async fn join_background_tasks(&mut self) {}
+    /// Flush any output which has been buffered.
+    ///
+    /// This will attempt to flush buffers and wait for all bytes to reach the
+    /// "end", whatever the "end" means in the context of this stream.
+    async fn flush(&mut self) -> Result<(), Error>;
 }
 
 pub(crate) enum InternalInputStream {
@@ -288,6 +284,9 @@ mod test {
             async fn ready(&mut self) -> Result<(), Error> {
                 unimplemented!();
             }
+            async fn flush(&mut self) -> Result<(), Error> {
+                Ok(())
+            }
         }
 
         let dummy = DummyOutputStream;