cm-async: Fix cancelling a completed host task (#12797) (#12801)

Incorrect behavior on Wasmtime's part led to erroneously raising a trap
in the guest about the terminal status of a task being delivered
already. This commit refactors some internals to ensure that cancelling
a completed host task is not a failure, it instead just returns that the
host task is indeed completed.

Author: alexcrichton

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

@@ -810,7 +810,7 @@ pub(crate) fn poll_and_block<R: Send + Sync + 'static>(
 
     // Retrieve and return the result.
     let host_state = &mut store.concurrent_state_mut().get_mut(task)?.state;
-    match mem::replace(host_state, HostTaskState::CalleeDone) {
+    match mem::replace(host_state, HostTaskState::CalleeDone { cancelled: false }) {
         HostTaskState::CalleeFinished(result) => Ok(match result.downcast() {
             Ok(result) => *result,
             Err(_) => bail_bug!("host task finished with wrong type of result"),
@@ -2863,7 +2863,14 @@ impl Instance {
 
                 lower(store.as_context_mut(), result)?;
                 let state = store.0.concurrent_state_mut();
-                state.get_mut(task)?.state = HostTaskState::CalleeDone;
+                match &mut state.get_mut(task)?.state {
+                    // The task is already flagged as finished because it was
+                    // cancelled. No need to transition further.
+                    HostTaskState::CalleeDone { .. } => {}
+
+                    // Otherwise transition this task to the done state.
+                    other => *other = HostTaskState::CalleeDone { cancelled: false },
+                }
                 Waitable::Host(task).set_event(state, Some(Event::Subtask { status }))?;
 
                 store.0.set_thread(old)?;
@@ -3114,7 +3121,7 @@ impl Instance {
             let task = concurrent_state.get_mut(id)?;
             match &task.state {
                 HostTaskState::CalleeRunning(_) => bail!(Trap::SubtaskDropNotResolved),
-                HostTaskState::CalleeDone => {}
+                HostTaskState::CalleeDone { .. } => {}
                 HostTaskState::CalleeStarted | HostTaskState::CalleeFinished(_) => {
                     bail_bug!("invalid state for callee in `subtask.drop`")
                 }
@@ -3593,15 +3600,28 @@ impl Instance {
         let needs_block;
         if let Waitable::Host(host_task) = waitable {
             let state = &mut concurrent_state.get_mut(host_task)?.state;
-            match mem::replace(state, HostTaskState::CalleeDone) {
+            match mem::replace(state, HostTaskState::CalleeDone { cancelled: true }) {
                 // 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(),
+                //
+                // After cancelling this falls through to block waiting for the
+                // host task to actually finish assuming that `async_` is false.
+                // This blocking behavior resolves the race of `handle.abort()`
+                // with the task actually getting cancelled or finishing.
+                HostTaskState::CalleeRunning(handle) => {
+                    handle.abort();
+                    needs_block = true;
+                }
 
                 // Cancellation was already requested, so fail as the task can't
                 // be cancelled twice.
-                HostTaskState::CalleeDone => {
-                    bail!(Trap::SubtaskCancelAfterTerminal);
+                HostTaskState::CalleeDone { cancelled } => {
+                    if cancelled {
+                        bail!(Trap::SubtaskCancelAfterTerminal);
+                    } else {
+                        // The callee is already done so there's no need to
+                        // block further for an event.
+                        needs_block = false;
+                    }
                 }
 
                 // These states should not be possible for a subtask that's
@@ -3610,12 +3630,6 @@ impl Instance {
                     bail_bug!("invalid states for host callee")
                 }
             }
-
-            // 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.current_guest_thread()?;
             let guest_task = TableId::<GuestTask>::new(rep);
@@ -4254,7 +4268,7 @@ enum HostTaskState {
 
     /// Terminal state for host tasks meaning that the task was cancelled or the
     /// result was taken.
-    CalleeDone,
+    CalleeDone { cancelled: bool },
 }
 
 impl HostTask {

tests/misc_testsuite/component-model/async/cancel-host.wast

@@ -383,3 +383,94 @@
 )
 
 (assert_return (invoke "run"))
+
+;; If a host task completes, but the guest doesn't actually receive the
+;; notification that it's done, it should be possible to cancel it.
+(component
+  (import "host" (instance $host
+    (export "return-two-slowly" (func async (result s32)))
+  ))
+
+  (core module $Mem (memory (export "mem") 1))
+  (core instance $mem (instantiate $Mem))
+
+  (core module $m
+    (import "" "slow-sync" (func $slow-sync (result i32)))
+    (import "" "slow-async" (func $slow-async (param i32) (result i32)))
+    (import "" "subtask.cancel" (func $subtask.cancel (param i32) (result i32)))
+    (import "" "subtask.drop" (func $subtask.drop (param i32)))
+    (func (export "run")
+      (local $subtask i32)
+
+      ;; start a subtask for `return-two-slowly`
+      (local.set $subtask (call $start-slow))
+
+      ;; Call `return-two-slowly` synchronously a few times. This gives a
+      ;; chance for the host to finish the previous task, which shouldn't
+      ;; interfere with the below...
+      (call $assert-eq (call $slow-sync) (i32.const 2))
+      (call $assert-eq (call $slow-sync) (i32.const 2))
+      (call $assert-eq (call $slow-sync) (i32.const 2))
+
+      ;; We've never seen the result of `$subtask`, so it should be safe to
+      ;; cancel. This should either yield that the subtask returned or that
+      ;; it's return was cancelled, it's up to hosts.
+      (call $assert-either (call $subtask.cancel (local.get $subtask))
+                           (i32.const 4)  ;; RETURN_CANCELLED
+                           (i32.const 2)) ;; RETURNED
+      (call $subtask.drop (local.get $subtask))
+    )
+
+    ;; asserts param 0 == param 1
+    (func $assert-eq (param i32 i32)
+      local.get 0
+      local.get 1
+      i32.ne
+      if unreachable end)
+
+    ;; Asserts that param 0 is either equal to param 1 or param 2
+    (func $assert-either (param i32 i32 i32)
+      local.get 0
+      local.get 1
+      i32.eq
+      if return end
+      local.get 0
+      local.get 2
+      i32.eq
+      if return end
+      unreachable)
+
+    (func $start-slow (result i32)
+      (local $tmp i32)
+
+      ;; Start slow, expect STARTED
+      (call $slow-async (i32.const 100))
+      local.tee $tmp
+      i32.const 0xf
+      i32.and
+      i32.const 1 ;; STARTED
+      i32.ne
+      if unreachable end
+      local.get $tmp
+      i32.const 4
+      i32.shr_u
+    )
+  )
+  (core func $slow-sync (canon lower (func $host "return-two-slowly") (memory $mem "mem")))
+  (core func $slow-async (canon lower (func $host "return-two-slowly") async (memory $mem "mem")))
+  (core func $subtask.cancel (canon subtask.cancel))
+  (core func $subtask.drop (canon subtask.drop))
+  (core instance $i (instantiate $m
+    (with "" (instance
+      (export "slow-sync" (func $slow-sync))
+      (export "slow-async" (func $slow-async))
+      (export "subtask.cancel" (func $subtask.cancel))
+      (export "subtask.drop" (func $subtask.drop))
+    ))
+  ))
+
+  (func (export "run") async
+    (canon lift (core func $i "run")))
+)
+
+(assert_return (invoke "run"))