move waitable to original set after moving it to temporary one (#11357)

There are a few places in `concurrent.rs` where we use
`GuestTask::sync_call_set` to wait on waitables during synchronous calls.
However, they may have been members of another set before we joined them to
`sync_call_set`, in which case we need to move them back when we're done (or at
least remove them from `sync_call_set`).

Prior to this fix, we would panic when dropping a task which had subtasks which
had been synchronously cancelled.  I've updated `async_cancel_callee.rs` to
cover that case.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

Author: dicej

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

@@ -48,6 +48,11 @@ pub async fn async_cancel_caller() -> Result<()> {
     test_cancel(Mode::Normal).await
 }
 
+#[tokio::test]
+pub async fn async_cancel_caller_leak_task_after_cancel() -> Result<()> {
+    test_cancel(Mode::LeakTaskAfterCancel).await
+}
+
 #[tokio::test]
 pub async fn async_trap_cancel_guest_after_start_cancelled() -> Result<()> {
     test_cancel_trap(Mode::TrapCancelGuestAfterStartCancelled).await

crates/misc/component-async-tests/wit/test.wit

@@ -151,6 +151,7 @@ interface cancel {
     trap-cancel-guest-after-return,
     trap-cancel-host-after-return-cancelled,
     trap-cancel-host-after-return,
+    leak-task-after-cancel,
   }
 
   run: func(mode: mode, cancel-delay-millis: u64);

crates/test-programs/src/bin/async_cancel_callee.rs

@@ -57,6 +57,7 @@ const _MODE_TRAP_CANCEL_GUEST_AFTER_RETURN_CANCELLED: u8 = 2;
 const _MODE_TRAP_CANCEL_GUEST_AFTER_RETURN: u8 = 3;
 const MODE_TRAP_CANCEL_HOST_AFTER_RETURN_CANCELLED: u8 = 4;
 const MODE_TRAP_CANCEL_HOST_AFTER_RETURN: u8 = 5;
+const MODE_LEAK_TASK_AFTER_CANCEL: u8 = 6;
 
 #[derive(Clone, Copy)]
 struct SleepParams {
@@ -166,7 +167,10 @@ unsafe extern "C" fn callback_sleep_with_options_sleep_millis(
                 }
 
                 waitable_join(*waitable, 0);
-                subtask_drop(*waitable);
+
+                if params.mode != MODE_LEAK_TASK_AFTER_CANCEL {
+                    subtask_drop(*waitable);
+                }
 
                 if params.on_cancel_delay_millis == 0 {
                     match params.on_cancel {

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

@@ -2412,8 +2412,10 @@ impl Instance {
 
         // Use the caller's `GuestTask::sync_call_set` to register interest in
         // the subtask...
+        let guest_waitable = Waitable::Guest(guest_task);
+        let old_set = guest_waitable.common(state)?.set;
         let set = state.get_mut(caller)?.sync_call_set;
-        Waitable::Guest(guest_task).join(state, Some(set))?;
+        guest_waitable.join(state, Some(set))?;
 
         // ... and suspend this fiber temporarily while we wait for it to start.
         //
@@ -2468,7 +2470,7 @@ impl Instance {
 
         let state = self.concurrent_state_mut(store.0);
 
-        Waitable::Guest(guest_task).join(state, None)?;
+        guest_waitable.join(state, old_set)?;
 
         if let Some(storage) = storage {
             // The caller used a sync-lowered import to call an async-lifted
@@ -3167,10 +3169,14 @@ impl Instance {
                     if async_ {
                         return Ok(BLOCKED);
                     } else {
+                        let waitable = Waitable::Guest(guest_task);
+                        let old_set = waitable.common(concurrent_state)?.set;
                         let set = concurrent_state.get_mut(caller)?.sync_call_set;
-                        Waitable::Guest(guest_task).join(concurrent_state, Some(set))?;
+                        waitable.join(concurrent_state, Some(set))?;
 
                         self.suspend(store, SuspendReason::Waiting { set, task: caller })?;
+
+                        waitable.join(self.concurrent_state_mut(store), old_set)?;
                     }
                 }
             }
@@ -3945,6 +3951,8 @@ impl Waitable {
     /// remove it from any set it may currently belong to (when `set` is
     /// `None`).
     fn join(&self, state: &mut ConcurrentState, set: Option<TableId<WaitableSet>>) -> Result<()> {
+        log::trace!("waitable {self:?} join set {set:?}",);
+
         let old = mem::replace(&mut self.common(state)?.set, set);
 
         if let Some(old) = old {