don't trap on idle in Instance::run_concurrent (#11756)

Previously, `Instance::run_concurrent` returned `Trap::AsyncDeadlock` when all
guest tasks and background host tasks had completed, and yet the future
parameter it was passed still hadn't resolved.  The theory was that this
indicated a mistake on the host embedder's part, but it turns out there are
scenarios where this is actually what the embedder wanted.

For example, consider a host embedder that implements a pool of worker tasks,
each of which runs a loop inside async closure passed to
`Instance::run_concurrent`.  In this case, each worker accepts jobs (which
involve calling guest functions) from a multiple-producer, multiple-consumer job
queue, adding them to a `futures::stream::FuturesUnordered` so they can be run
concurrently.  When all the jobs accepted by a given worker have finished, there
may be a lull during which no new jobs are yet available.  In that case, the
worker _could_ break out of the loop, resolve the future, allow
`Instance::run_concurrent` to finish, and wait until the next job arrives before
calling `Instance::run_concurrent` again, but that's more awkward (i.e. nested
loops, complicated control flow) than just a single loop inside
`Instance::run_concurrent` that goes idle now and then.

In short, the closure passed to `Instance::run_concurrent` might experience
delays between when a set of guest tasks have completed and when the next set
are ready to start, and that's not necessarily a bug.

Internally, I've added a new `run_concurrent_trap_on_idle` function, which
provides the original, trapping behavior, and I'm using it to implement
`[Typed]Func::call_async`, in which case it _is_ an error if the event loop goes
idle without the future resolving.  If this turns out to be useful as part of
the public API, we can change the `pub(super)` to `pub`.

Signed-off-by: Joel Dice <[email protected]>

Author: dicej

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

@@ -1,12 +1,14 @@
 use super::util::test_run;
 use crate::scenario::util::{config, make_component};
 use anyhow::Result;
+use component_async_tests::util;
 use component_async_tests::{Ctx, sleep};
 use std::future;
 use std::pin::pin;
 use std::sync::{Arc, Mutex};
 use std::task::Poll;
-use wasmtime::component::{Instance, Linker, ResourceTable};
+use std::time::Duration;
+use wasmtime::component::{Accessor, Instance, Linker, ResourceTable};
 use wasmtime::{AsContextMut, Engine, Store, StoreContextMut};
 use wasmtime_wasi::WasiCtxBuilder;
 
@@ -70,41 +72,65 @@ async fn test_sleep_post_return(components: &[&str]) -> Result<()> {
     );
 
     let instance = linker.instantiate_async(&mut store, &component).await?;
+    let guest = sleep_post_return::SleepPostReturnCallee::new(&mut store, &instance)?;
+
+    async fn run_with(
+        accessor: &Accessor<Ctx>,
+        guest: &sleep_post_return::SleepPostReturnCallee,
+    ) -> Result<()> {
+        // This function should return immediately, then sleep the specified
+        // number of milliseconds after returning, and then finally exit.
+        let exit = guest
+            .local_local_sleep_post_return()
+            .call_run(accessor, 100)
+            .await?
+            .1;
+        // The function has returned, now we wait for it (and any subtasks
+        // it may have spawned) to exit.
+        exit.block(accessor).await;
+        anyhow::Ok(())
+    }
 
-    async fn run(mut store: StoreContextMut<'_, Ctx>, instance: Instance) -> Result<()> {
-        let guest = sleep_post_return::SleepPostReturnCallee::new(&mut store, &instance)?;
+    async fn run(
+        store: StoreContextMut<'_, Ctx>,
+        instance: Instance,
+        guest: &sleep_post_return::SleepPostReturnCallee,
+    ) -> Result<()> {
         instance
             .run_concurrent(store, async |accessor| {
-                // This function should return immediately, then sleep the specified
-                // number of milliseconds after returning, and then finally exit.
-                let exit = guest
-                    .local_local_sleep_post_return()
-                    .call_run(accessor, 100)
-                    .await?
-                    .1;
-                // The function has returned, now we wait for it (and any subtasks
-                // it may have spawned) to exit.
-                exit.block(accessor).await;
+                run_with(accessor, guest).await?;
+
+                // Go idle for a bit before doing it again.  This tests that
+                // `Instance::run_concurrent` is okay with having no outstanding
+                // guest or host tasks to poll for a while, trusting that we'll
+                // resolve the future independently, with or without giving it
+                // more work to do.
+                util::sleep(Duration::from_millis(100)).await;
+
+                run_with(accessor, guest).await?;
+
                 anyhow::Ok(())
             })
             .await?
     }
 
-    run(store.as_context_mut(), instance).await?;
+    run(store.as_context_mut(), instance, &guest).await?;
     // At this point, all subtasks should have exited, meaning no waitables,
     // tasks, or other concurrent state should remain present in the instance.
     instance.assert_concurrent_state_empty(&mut store);
 
     // Do it again, but this time cancel the event loop before it exits:
     assert!(
-        future::poll_fn(|cx| Poll::Ready(pin!(run(store.as_context_mut(), instance)).poll(cx)))
-            .await
-            .is_pending()
+        future::poll_fn(|cx| Poll::Ready(
+            pin!(run(store.as_context_mut(), instance, &guest)).poll(cx)
+        ))
+        .await
+        .is_pending()
     );
 
     // Assuming the event loop is cancel-safe, this should complete without
     // errors or panics:
-    run(store.as_context_mut(), instance).await?;
+    run(store.as_context_mut(), instance, &guest).await?;
 
     Ok(())
 }

crates/wasmtime/src/runtime/component/concurrent.rs

@@ -1034,9 +1034,8 @@ impl Instance {
     /// Run the specified closure `fun` to completion as part of this instance's
     /// event loop.
     ///
-    /// Like [`Self::run`], this will run `fun` as part of this instance's event
-    /// loop until it yields a result _or_ there are no more tasks to run.
-    /// Unlike [`Self::run`], `fun` is provided an [`Accessor`], which provides
+    /// This will run `fun` as part of this instance's event loop until it
+    /// yields a result.  `fun` is provided an [`Accessor`], which provides
     /// controlled access to the `Store` and its data.
     ///
     /// This function can be used to invoke [`Func::call_concurrent`] for
@@ -1076,9 +1075,32 @@ impl Instance {
     /// # }
     /// ```
     pub async fn run_concurrent<T, R>(
+        self,
+        store: impl AsContextMut<Data = T>,
+        fun: impl AsyncFnOnce(&Accessor<T>) -> R,
+    ) -> Result<R>
+    where
+        T: Send + 'static,
+    {
+        self.do_run_concurrent(store, fun, false).await
+    }
+
+    pub(super) async fn run_concurrent_trap_on_idle<T, R>(
+        self,
+        store: impl AsContextMut<Data = T>,
+        fun: impl AsyncFnOnce(&Accessor<T>) -> R,
+    ) -> Result<R>
+    where
+        T: Send + 'static,
+    {
+        self.do_run_concurrent(store, fun, true).await
+    }
+
+    async fn do_run_concurrent<T, R>(
         self,
         mut store: impl AsContextMut<Data = T>,
         fun: impl AsyncFnOnce(&Accessor<T>) -> R,
+        trap_on_idle: bool,
     ) -> Result<R>
     where
         T: Send + 'static,
@@ -1112,7 +1134,7 @@ impl Instance {
         // SAFETY: We never move `dropper` nor its `value` field.
         let future = unsafe { Pin::new_unchecked(dropper.value.deref_mut()) };
 
-        self.poll_until(dropper.store.as_context_mut(), future)
+        self.poll_until(dropper.store.as_context_mut(), future, trap_on_idle)
             .await
     }
 
@@ -1161,13 +1183,12 @@ impl Instance {
 
     /// Run this instance's event loop.
     ///
-    /// The returned future will resolve when either the specified future
-    /// completes (in which case we return its result) or no further progress
-    /// can be made (in which case we trap with `Trap::AsyncDeadlock`).
+    /// The returned future will resolve when the specified future completes.
     async fn poll_until<T, R>(
         self,
         mut store: StoreContextMut<'_, T>,
         mut future: Pin<&mut impl Future<Output = R>>,
+        trap_on_idle: bool,
     ) -> Result<R>
     where
         T: Send + 'static,
@@ -1258,35 +1279,17 @@ impl Instance {
                                     // futures in `ConcurrentState::futures`,
                                     // there are no remaining work items, _and_
                                     // the future we were passed as an argument
-                                    // still hasn't completed, meaning we're
-                                    // stuck, so we return an error.  The
-                                    // underlying assumption is that `future`
-                                    // depends on this component instance making
-                                    // such progress, and thus there's no point
-                                    // in continuing to poll it given we've run
-                                    // out of work to do.
-                                    //
-                                    // Note that we'd also reach this point if
-                                    // the host embedder passed e.g. a
-                                    // `std::future::Pending` to
-                                    // `Instance::run_concurrent`, in which case
-                                    // we'd return a "deadlock" error even when
-                                    // any and all tasks have completed
-                                    // normally.  However, that's not how
-                                    // `Instance::run_concurrent` is intended
-                                    // (and documented) to be used, so it seems
-                                    // reasonable to lump that case in with
-                                    // "real" deadlocks.
-                                    //
-                                    // TODO: Once we've added host APIs for
-                                    // cancelling in-progress tasks, we can
-                                    // return some other, non-error value here,
-                                    // treating it as "normal" and giving the
-                                    // host embedder a chance to intervene by
-                                    // cancelling one or more tasks and/or
-                                    // starting new tasks capable of waking the
-                                    // existing ones.
-                                    Poll::Ready(Err(anyhow!(crate::Trap::AsyncDeadlock)))
+                                    // still hasn't completed.
+                                    if trap_on_idle {
+                                        // `trap_on_idle` is true, so we exit
+                                        // immediately.
+                                        Poll::Ready(Err(anyhow!(crate::Trap::AsyncDeadlock)))
+                                    } else {
+                                        // `trap_on_idle` is false, so we assume
+                                        // that future will wake up and give us
+                                        // more work to do when it's ready to.
+                                        Poll::Pending
+                                    }
                                 }
                             }
                             // There is at least one pending future in

crates/wasmtime/src/runtime/component/func.rs

@@ -281,7 +281,7 @@ impl Func {
         #[cfg(feature = "component-model-async")]
         {
             self.instance
-                .run_concurrent(&mut store, async |store| {
+                .run_concurrent_trap_on_idle(&mut store, async |store| {
                     self.call_concurrent_dynamic(store, params, results, false)
                         .await
                         .map(drop)

crates/wasmtime/src/runtime/component/func/typed.rs

@@ -234,7 +234,7 @@ where
             let result = concurrent::queue_call(wrapper.store.as_context_mut(), prepared)?;
             self.func
                 .instance
-                .run_concurrent(wrapper.store.as_context_mut(), async |_| {
+                .run_concurrent_trap_on_idle(wrapper.store.as_context_mut(), async |_| {
                     Ok(result.await?.0)
                 })
                 .await?