Fix/improve host subtask cancellation (#12640)

This commit refactors some of the internals of `subtask.cancel` with
respect to host subtasks. Notably a few panics and semantic bugs are
fixed here. The main bug was that host subtasks could be aborted but
their completion might have still been queued up which would produce the
result somewhere or assert that the task exists. Cancellation is changed
to use `wait_for_event` to ensure that this completion is executed
before `subtask.cancel` returns. This helps keep host subtasks looking
more similar to guest subtasks in that respect.

Closes #12631
Closes #12632

Co-authored-by: Jelle van den Hooff <jelle@vandenhooff.name>

Author: alexcrichton

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

@@ -756,7 +756,9 @@ pub(crate) fn poll_and_block<R: Send + Sync + 'static>(
         let result = future.await?;
         tls::get(move |store| {
             let state = store.concurrent_state_mut();
-            state.get_mut(task)?.result = Some(Box::new(result) as _);
+            let host_state = &mut state.get_mut(task)?.state;
+            assert!(matches!(host_state, HostTaskState::CalleeStarted));
+            *host_state = HostTaskState::CalleeFinished(Box::new(result));
 
             Waitable::Host(task).set_event(
                 state,
@@ -808,14 +810,11 @@ pub(crate) fn poll_and_block<R: Send + Sync + 'static>(
     }
 
     // Retrieve and return the result.
-    Ok(*store
-        .concurrent_state_mut()
-        .get_mut(task)?
-        .result
-        .take()
-        .unwrap()
-        .downcast()
-        .unwrap())
+    let host_state = &mut store.concurrent_state_mut().get_mut(task)?.state;
+    match mem::replace(host_state, HostTaskState::CalleeDone) {
+        HostTaskState::CalleeFinished(result) => Ok(*result.downcast().unwrap()),
+        _ => panic!("unexpected host task state after completion"),
+    }
 }
 
 /// Execute the specified guest call.
@@ -1550,7 +1549,7 @@ impl StoreOpaque {
     pub fn enter_host_call(&mut self) -> Result<()> {
         let state = self.concurrent_state_mut();
         let caller = state.unwrap_current_guest_thread();
-        let task = state.push(HostTask::new(caller))?;
+        let task = state.push(HostTask::new(caller, HostTaskState::CalleeStarted))?;
         log::trace!("new host task {task:?}");
         self.set_thread(task);
         Ok(())
@@ -2736,12 +2735,13 @@ impl Instance {
     ///
     /// Whether the future returns `Ready` immediately or later, the `lower`
     /// function will be used to lower the result, if any, into the guest caller's
-    /// stack and linear memory unless the task has been cancelled.
+    /// stack and linear memory. The `lower` function is invoked with `None` if
+    /// the future is cancelled.
     pub(crate) fn first_poll<T: 'static, R: Send + 'static>(
         self,
         mut store: StoreContextMut<'_, T>,
         future: impl Future<Output = Result<R>> + Send + 'static,
-        lower: impl FnOnce(StoreContextMut<T>, R) -> Result<()> + Send + 'static,
+        lower: impl FnOnce(StoreContextMut<T>, Option<R>) -> Result<()> + Send + 'static,
     ) -> Result<Option<u32>> {
         let token = StoreToken::new(store.as_context_mut());
         let state = store.0.concurrent_state_mut();
@@ -2751,9 +2751,9 @@ impl Instance {
         // context state for the future.
         let (join_handle, future) = JoinHandle::run(future);
         {
-            let task = state.get_mut(task)?;
-            assert!(task.join_handle.is_none());
-            task.join_handle = Some(join_handle);
+            let state = &mut state.get_mut(task)?.state;
+            assert!(matches!(state, HostTaskState::CalleeStarted));
+            *state = HostTaskState::CalleeRunning(join_handle);
         }
 
         let mut future = Box::pin(future);
@@ -2771,7 +2771,7 @@ impl Instance {
         match poll {
             // It finished immediately; lower the result and delete the task.
             Poll::Ready(Some(result)) => {
-                lower(store.as_context_mut(), result?)?;
+                lower(store.as_context_mut(), Some(result?))?;
                 return Ok(None);
             }
 
@@ -2792,9 +2792,8 @@ impl Instance {
         // the task returned.
         let future = Box::pin(async move {
             let result = match future.await {
-                Some(result) => result?,
-                // Task was cancelled; nothing left to do.
-                None => return Ok(()),
+                Some(result) => Some(result?),
+                None => None,
             };
             let on_complete = move |store: &mut dyn VMStore| {
                 // Restore the `current_thread` to be the host so `lower` knows
@@ -2805,15 +2804,16 @@ impl Instance {
                 assert!(state.current_thread.is_none());
                 store.0.set_thread(task);
 
+                let status = if result.is_some() {
+                    Status::Returned
+                } else {
+                    Status::ReturnCancelled
+                };
+
                 lower(store.as_context_mut(), result)?;
                 let state = store.0.concurrent_state_mut();
-                state.get_mut(task)?.join_handle.take();
-                Waitable::Host(task).set_event(
-                    state,
-                    Some(Event::Subtask {
-                        status: Status::Returned,
-                    }),
-                )?;
+                state.get_mut(task)?.state = HostTaskState::CalleeDone;
+                Waitable::Host(task).set_event(state, Some(Event::Subtask { status }))?;
 
                 // Go back to "no current thread" at the end.
                 store.0.set_thread(CurrentThread::None);
@@ -3059,8 +3059,12 @@ impl Instance {
         let (waitable, expected_caller, delete) = if is_host {
             let id = TableId::<HostTask>::new(rep);
             let task = concurrent_state.get_mut(id)?;
-            if task.join_handle.is_some() {
-                bail!("cannot drop a subtask which has not yet resolved");
+            match &task.state {
+                HostTaskState::CalleeRunning(_) => {
+                    bail!("cannot drop a subtask which has not yet resolved");
+                }
+                HostTaskState::CalleeDone => {}
+                HostTaskState::CalleeStarted | HostTaskState::CalleeFinished(_) => unreachable!(),
             }
             (Waitable::Host(id), task.caller, true)
         } else {
@@ -3537,11 +3541,30 @@ impl Instance {
 
         log::trace!("subtask_cancel {waitable:?} (handle {task_id})");
 
+        let needs_block;
         if let Waitable::Host(host_task) = waitable {
-            if let Some(handle) = concurrent_state.get_mut(host_task)?.join_handle.take() {
-                handle.abort();
-                return Ok(Status::ReturnCancelled as u32);
+            let state = &mut concurrent_state.get_mut(host_task)?.state;
+            match mem::replace(state, HostTaskState::CalleeDone) {
+                // If the callee is still running, signal an abort is requested.
+                // Then fall through to determine what to do next.
+                HostTaskState::CalleeRunning(handle) => handle.abort(),
+
+                // Cancellation was already requested, so fail as the task can't
+                // be cancelled twice.
+                HostTaskState::CalleeDone => {
+                    bail!("`subtask.cancel` called after terminal status delivered");
+                }
+
+                // These states should not be possible for a subtask that's
+                // visible from the guest, so panic here.
+                HostTaskState::CalleeStarted | HostTaskState::CalleeFinished(_) => unreachable!(),
             }
+
+            // Cancelling host tasks always needs to block on them to await the
+            // result of the completion set up in `first_poll`. This'll resolve
+            // the race of `handle.abort()` above to see if it actually
+            // cancelled something or if the future ended up finishing.
+            needs_block = true;
         } else {
             let caller = concurrent_state.unwrap_current_guest_thread();
             let guest_task = TableId::<GuestTask>::new(rep);
@@ -3622,16 +3645,31 @@ impl Instance {
                     }
                 }
 
-                let concurrent_state = store.concurrent_state_mut();
-                let task = concurrent_state.get_mut(guest_task)?;
-                if !task.returned_or_cancelled() {
-                    if async_ {
-                        return Ok(BLOCKED);
-                    } else {
-                        store.wait_for_event(Waitable::Guest(guest_task))?;
-                    }
-                }
+                // Guest tasks need to block if they have not yet returned or
+                // cancelled, even as a result of the event delivery above.
+                needs_block = !store
+                    .concurrent_state_mut()
+                    .get_mut(guest_task)?
+                    .returned_or_cancelled()
+            } else {
+                needs_block = false;
+            }
+        };
+
+        // If we need to block waiting on the terminal status of this subtask
+        // then return immediately in `async` mode, or otherwise wait for the
+        // event to get signaled through the store.
+        if needs_block {
+            if async_ {
+                return Ok(BLOCKED);
             }
+
+            // Wait for this waitable to get signaled with its terminal status
+            // from the completion callback enqueued by `first_poll`. Once
+            // that's done fall through to the sahred
+            store.wait_for_event(waitable)?;
+
+            // .. fall through to determine what event's in store for us.
         }
 
         let event = waitable.take_event(store.concurrent_state_mut())?;
@@ -4147,24 +4185,39 @@ struct HostTask {
     /// borrows to the host, for example.
     call_context: CallContext,
 
-    /// For host tasks which end up doing some asynchronous work (e.g.
-    /// async-lowered and didn't complete on the first poll) this handle is used
-    /// as a signal to cancel the future as it resides in the store's
-    /// `FuturesUnordered`.
-    join_handle: Option<JoinHandle>,
+    state: HostTaskState,
+}
 
-    /// Box<Any> of the result of this host task.
-    result: Option<LiftedResult>,
+enum HostTaskState {
+    /// A host task has been created and it's considered "started".
+    ///
+    /// The host task has yet to enter `first_poll` or `poll_and_block` which
+    /// is where this will get updated further.
+    CalleeStarted,
+
+    /// State used for tasks in `first_poll` meaning that the guest did an async
+    /// lower of a host async function which is blocked. The specified handle is
+    /// linked to the future in the main `FuturesUnordered` of a store which is
+    /// used to cancel it if the guest requests cancellation.
+    CalleeRunning(JoinHandle),
+
+    /// Terminal state used for tasks in `poll_and_block` to store the result of
+    /// their computation. Note that this state is not used for tasks in
+    /// `first_poll`.
+    CalleeFinished(LiftedResult),
+
+    /// Terminal state for host tasks meaning that the task was cancelled or the
+    /// result was taken.
+    CalleeDone,
 }
 
 impl HostTask {
-    fn new(caller: QualifiedThreadId) -> Self {
+    fn new(caller: QualifiedThreadId, state: HostTaskState) -> Self {
         Self {
             common: WaitableCommon::default(),
             call_context: CallContext::default(),
             caller,
-            join_handle: None,
-            result: None,
+            state,
         }
     }
 }

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

@@ -447,7 +447,7 @@ where
                 ptr,
             )?)
         };
-        Self::lower_result_and_exit_call(&mut lower, ty, ret, dst)
+        Self::lower_result_and_exit_call(&mut lower, ty, Some(ret), dst)
     }
 
     /// Implementation of the "async" ABI of the component model.
@@ -499,7 +499,7 @@ where
                 Self::lower_result_and_exit_call(
                     &mut LowerContext::new(store, options, instance),
                     ty,
-                    result?,
+                    Some(result?),
                     Destination::Memory(retptr),
                 )?;
                 None
@@ -568,17 +568,19 @@ where
     fn lower_result_and_exit_call(
         lower: &mut LowerContext<'_, T>,
         ty: TypeFuncIndex,
-        ret: R,
+        ret: Option<R>,
         dst: Destination<'_>,
     ) -> Result<()> {
-        let caller_instance = lower.options().instance;
-        let mut flags = lower.instance_mut().instance_flags(caller_instance);
-        unsafe {
-            flags.set_may_leave(false);
-        }
-        Self::lower_result(lower, ty, ret, dst)?;
-        unsafe {
-            flags.set_may_leave(true);
+        if let Some(ret) = ret {
+            let caller_instance = lower.options().instance;
+            let mut flags = lower.instance_mut().instance_flags(caller_instance);
+            unsafe {
+                flags.set_may_leave(false);
+            }
+            Self::lower_result(lower, ty, ret, dst)?;
+            unsafe {
+                flags.set_may_leave(true);
+            }
         }
         lower.validate_scope_exit()?;
         Ok(())

crates/wast/src/spectest.rs

@@ -197,5 +197,26 @@ pub fn link_component_spectest<T>(linker: &mut component::Linker<T>) -> Result<(
             Ok(())
         },
     )?;
+    i.func_wrap_concurrent("never-return", |_, _: ()| {
+        Box::pin(async move { std::future::pending::<Result<()>>().await })
+    })?;
+    i.func_wrap_concurrent("return-two-slowly", |_, _: ()| {
+        Box::pin(async move {
+            tokio::task::yield_now().await;
+            Ok((2,))
+        })
+    })?;
+    i.func_wrap_concurrent("echo-slowly", |_, (a,): (u32,)| {
+        Box::pin(async move {
+            tokio::task::yield_now().await;
+            Ok((a,))
+        })
+    })?;
+    i.func_wrap_concurrent(
+        "[method]resource1.never-return",
+        |_, (_,): (Resource<Resource1>,)| {
+            Box::pin(async move { std::future::pending::<Result<()>>().await })
+        },
+    )?;
     Ok(())
 }