save and restore next_switch_item for sync-to-sync calls

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

@@ -1125,6 +1125,7 @@ impl<T> StoreContextMut<'_, T> {
         assert!(state.next_switch_item.is_none());
         assert!(state.high_priority.is_empty());
         assert!(state.low_priority.is_empty());
+        assert!(state.saved_next_switch_items.is_empty());
         assert!(state.unforced_current_thread.is_none());
         assert!(state.deferred_host_call_context.is_none());
         assert!(state.futures_mut().unwrap().is_empty());
@@ -1945,6 +1946,13 @@ impl StoreOpaque {
             return self.enter_call_not_concurrent();
         }
 
+        // Stash the next switch item (if any) for the duration of this call
+        // since it's for the caller's caller, not for us.  We'll restore this
+        // in `exit_guest_sync_call`.
+        let state = self.concurrent_state_mut()?;
+        let item = state.next_switch_item.take();
+        state.saved_next_switch_items.push(item);
+
         let thread = self.current_thread()?;
         let caller = if let Some(thread) = thread.guest() {
             Caller::Guest { thread: *thread }
@@ -2034,6 +2042,17 @@ impl StoreOpaque {
 
         self.cleanup_thread(thread, instance, CleanupTask::Yes)?;
 
+        let state = self.concurrent_state_mut()?;
+        let Some(item) = state.saved_next_switch_items.pop() else {
+            bail_bug!("unable to pop from `saved_next_switch_items`");
+        };
+        if let Some(item) = mem::replace(&mut state.next_switch_item, item) {
+            // Stash it back in the store to ensure it's cleaned up on store
+            // drop:
+            state.push_high_priority(item);
+            bail_bug!("`next_switch_item` unexpectedly already set");
+        }
+
         Ok(())
     }
 
@@ -2318,6 +2337,7 @@ impl StoreOpaque {
                         fiber,
                     };
 
+                    log::trace!("set next switch item to {item:?}");
                     if state.next_switch_item.replace(item).is_some() {
                         // This should be unreachable per the save/restore code
                         // in `Self::suspend`.
@@ -5942,6 +5962,10 @@ pub struct ConcurrentState {
     /// Whether the `StoreContextMut::poll_until` event loop is running.
     event_loop_running: bool,
 
+    /// Stack of switch items to push to and pop from when entering and exiting
+    /// sync-to-sync calls.
+    saved_next_switch_items: Vec<Option<WorkItem>>,
+
     /// See [TaskGroupHook].
     #[cfg(feature = "task-group-hook")]
     task_group_hook: Option<Box<dyn TaskGroupHook>>,
@@ -5966,6 +5990,7 @@ impl Default for ConcurrentState {
             interesting_tasks_empty_waker: None,
             ready_for_concurrent_call_waker: None,
             event_loop_running: false,
+            saved_next_switch_items: Vec::new(),
             #[cfg(feature = "task-group-hook")]
             task_group_hook: None,
         }
@@ -6053,6 +6078,12 @@ impl ConcurrentState {
         for item in mem::take(&mut self.low_priority) {
             handle_item(item);
         }
+        for item in mem::take(&mut self.saved_next_switch_items)
+            .into_iter()
+            .filter_map(|v| v)
+        {
+            handle_item(item);
+        }
 
         if let Some(them) = self.futures.get_mut().take() {
             futures.push(them);
@@ -6073,6 +6104,7 @@ impl ConcurrentState {
             next_switch_item,
             high_priority,
             low_priority,
+            saved_next_switch_items,
 
             // TODO(cm-gc): This field contains `ValRaw`s, but they are never GC
             // references because the component model doesn't support GC yet. We
@@ -6145,6 +6177,12 @@ impl ConcurrentState {
         for item in low_priority {
             handle_item(item);
         }
+        for item in saved_next_switch_items
+            .iter_mut()
+            .filter_map(|v| v.as_mut())
+        {
+            handle_item(item);
+        }
     }
 
     fn push<V: Send + Sync + 'static>(

tests/misc_testsuite/component-model/async/default-flags-sync-cancel.wast

@@ -0,0 +1,105 @@
+;;! component_model_async = true
+;;! component_model_more_async_builtins = true
+
+;; $C async-lowers a call to $B.h.  $B, before it has suspended, makes a
+;; sync-lowered call to the sync-lifted $A.g (which runs to completion on
+;; $B's fiber) and then yields to one of its own subtasks by cancelling the
+;; started, yielding $A.f.
+;;
+;; definitions.py: nothing special happens; $B's cancel resolves $A.f as
+;; CANCELLED_BEFORE_RETURNED, $B returns, $C sees RETURNED (2).
+;;
+;; Wasmtime: when $A.g's thread exits, cleanup_thread calls
+;; take_next_switch_item(), which moves the event loop's next_switch_item
+;; (the fiber of $C, parked in start_call until $B first suspends) into
+;; switch_item.  When $B then suspends with YieldingToSubtask (the cancel's
+;; yield_), resume_fiber stores $B's fiber as the new next_switch_item, the
+;; stale switch item resumes $C inside start_call, and the save/restore in
+;; StoreOpaque::suspend overwrites next_switch_item with the value it saved,
+;; dropping $B's in-progress fiber: the `StoreFiber` Drop impl asserts, the
+;; raw fiber's Drop panics again during unwinding, and the process aborts.
+(component
+  (component $A
+    (core func $task.cancel (canon task.cancel))
+    (core func $ws.new (canon waitable-set.new))
+    (core func $ws.drop (canon waitable-set.drop))
+    (core module $m
+      (import "" "task.cancel" (func $task.cancel))
+      (import "" "waitable-set.new" (func $ws.new (result i32)))
+      (import "" "waitable-set.drop" (func $ws.drop (param i32)))
+      ;; f: yield on start (callback code 1)
+      (func (export "f") (result i32) (i32.const 1))
+      ;; cb: on TASK_CANCELLED (6) cancel the task and exit; anything else is unexpected
+      (func (export "cb") (param i32 i32 i32) (result i32)
+        (if (i32.ne (local.get 0) (i32.const 6)) (then unreachable))
+        (call $task.cancel)
+        (i32.const 0))
+      ;; g: a sync export that touches the concurrent state (any built-in or
+      ;; host call does; the fuzzer used resource.new), which makes Wasmtime
+      ;; materialize the deferred `enter_guest_sync_call` for this call and
+      ;; therefore run `cleanup_thread` when it exits
+      (func (export "g") (call $ws.drop (call $ws.new))))
+    (core instance $i (instantiate $m (with "" (instance
+      (export "task.cancel" (func $task.cancel))
+      (export "waitable-set.new" (func $ws.new))
+      (export "waitable-set.drop" (func $ws.drop))))))
+    (func (export "f") async (canon lift (core func $i "f") async (callback (core func $i "cb"))))
+    (func (export "g") (canon lift (core func $i "g"))))
+
+  (component $B
+    (import "f" (func $f async))
+    (import "g" (func $g))
+    (core func $f (canon lower (func $f) async))
+    (core func $g (canon lower (func $g)))
+    (core func $subtask.cancel (canon subtask.cancel))
+    (core func $subtask.drop (canon subtask.drop))
+    (core func $task.return (canon task.return))
+    (core module $m
+      (import "" "f" (func $f (result i32)))
+      (import "" "g" (func $g))
+      (import "" "subtask.cancel" (func $subtask.cancel (param i32) (result i32)))
+      (import "" "subtask.drop" (func $subtask.drop (param i32)))
+      (import "" "task.return" (func $task.return))
+      (func (export "h") (result i32) (local $r i32)
+        (local.set $r (call $f))
+        ;; the call must be STARTED (1) with a subtask handle
+        (if (i32.ne (i32.and (local.get $r) (i32.const 0xf)) (i32.const 1)) (then unreachable))
+        ;; a sync call to a sync-lifted export, made before this task ever suspended
+        (call $g)
+        ;; cancel the yielding subtask: definitions.py returns CANCELLED_BEFORE_RETURNED (4)
+        (if (i32.ne (call $subtask.cancel (i32.shr_u (local.get $r) (i32.const 4))) (i32.const 4)) (then unreachable))
+        (call $subtask.drop (i32.shr_u (local.get $r) (i32.const 4)))
+        (call $task.return)
+        (i32.const 0 (; EXIT ;)))
+      (func (export "cb") (param i32 i32 i32) (result i32) unreachable))
+    (core instance $i (instantiate $m (with "" (instance
+      (export "f" (func $f))
+      (export "g" (func $g))
+      (export "subtask.cancel" (func $subtask.cancel))
+      (export "subtask.drop" (func $subtask.drop))
+      (export "task.return" (func $task.return))))))
+    (func (export "h") async (canon lift (core func $i "h") async (callback (core func $i "cb")))))
+
+  (component $C
+    (import "h" (func $h async))
+    (core func $h (canon lower (func $h) async))
+    (core func $task.return (canon task.return (result u32)))
+    (core module $m
+      (import "" "h" (func $h (result i32)))
+      (import "" "task.return" (func $task.return (param i32)))
+      (func (export "main") (result i32)
+        ;; $B runs to completion before $C regains control: RETURNED (2), no handle
+        (call $task.return (call $h))
+        (i32.const 0 (; EXIT ;)))
+      (func (export "cb") (param i32 i32 i32) (result i32) unreachable))
+    (core instance $i (instantiate $m (with "" (instance
+      (export "h" (func $h))
+      (export "task.return" (func $task.return))))))
+    (func (export "main") async (result u32) (canon lift (core func $i "main") async (callback (core func $i "cb")))))
+
+  (instance $a (instantiate $A))
+  (instance $b (instantiate $B (with "f" (func $a "f")) (with "g" (func $a "g"))))
+  (instance $c (instantiate $C (with "h" (func $b "h"))))
+  (export "main" (func $c "main")))
+
+(assert_return (invoke "main") (u32.const 2))

tests/misc_testsuite/component-model/async/default-flags-variant-second-call.wast

@@ -0,0 +1,109 @@
+;;! component_model_async = true
+;;! component_model_more_async_builtins = true
+
+;; $C async-lowers a call to $B.h.  $B, before it has suspended, makes a
+;; sync-lowered call to the sync-lifted $A.g (which runs to completion on
+;; $B's fiber) and then yields to one of its own subtasks by cancelling the
+;; started, yielding $A.f.
+;;
+;; definitions.py: nothing special happens; $B's cancel resolves $A.f as
+;; CANCELLED_BEFORE_RETURNED, $B returns, $C sees RETURNED (2).
+;;
+;; Wasmtime: when $A.g's thread exits, cleanup_thread calls
+;; take_next_switch_item(), which moves the event loop's next_switch_item
+;; (the fiber of $C, parked in start_call until $B first suspends) into
+;; switch_item.  When $B then suspends with YieldingToSubtask (the cancel's
+;; yield_), resume_fiber stores $B's fiber as the new next_switch_item, the
+;; stale switch item resumes $C inside start_call, and the save/restore in
+;; StoreOpaque::suspend overwrites next_switch_item with the value it saved,
+;; dropping $B's in-progress fiber: the `StoreFiber` Drop impl asserts, the
+;; raw fiber's Drop panics again during unwinding, and the process aborts.
+(component
+  (component $A
+    (core func $task.cancel (canon task.cancel))
+    (core func $ws.new (canon waitable-set.new))
+    (core func $ws.drop (canon waitable-set.drop))
+    (core module $m
+      (import "" "task.cancel" (func $task.cancel))
+      (import "" "waitable-set.new" (func $ws.new (result i32)))
+      (import "" "waitable-set.drop" (func $ws.drop (param i32)))
+      ;; f: yield on start (callback code 1)
+      (func (export "f") (result i32) (i32.const 1))
+      ;; cb: on TASK_CANCELLED (6) cancel the task and exit; anything else is unexpected
+      (func (export "cb") (param i32 i32 i32) (result i32)
+        (if (i32.ne (local.get 0) (i32.const 6)) (then unreachable))
+        (call $task.cancel)
+        (i32.const 0))
+      ;; g: a sync export that touches the concurrent state (any built-in or
+      ;; host call does; the fuzzer used resource.new), which makes Wasmtime
+      ;; materialize the deferred `enter_guest_sync_call` for this call and
+      ;; therefore run `cleanup_thread` when it exits
+      (func (export "g") (call $ws.drop (call $ws.new))))
+    (core instance $i (instantiate $m (with "" (instance
+      (export "task.cancel" (func $task.cancel))
+      (export "waitable-set.new" (func $ws.new))
+      (export "waitable-set.drop" (func $ws.drop))))))
+    (func (export "f") async (canon lift (core func $i "f") async (callback (core func $i "cb"))))
+    (func (export "g") (canon lift (core func $i "g"))))
+
+  (component $B
+    (import "f" (func $f async))
+    (import "g" (func $g))
+    (core func $f (canon lower (func $f) async))
+    (core func $g (canon lower (func $g)))
+    (core func $subtask.cancel (canon subtask.cancel))
+    (core func $subtask.drop (canon subtask.drop))
+    (core func $task.return (canon task.return))
+    (core module $m
+      (import "" "f" (func $f (result i32)))
+      (import "" "g" (func $g))
+      (import "" "subtask.cancel" (func $subtask.cancel (param i32) (result i32)))
+      (import "" "subtask.drop" (func $subtask.drop (param i32)))
+      (import "" "task.return" (func $task.return))
+      (func (export "h") (result i32) (local $r i32) (local $r2 i32)
+        (local.set $r (call $f))
+        ;; the call must be STARTED (1) with a subtask handle
+        (if (i32.ne (i32.and (local.get $r) (i32.const 0xf)) (i32.const 1)) (then unreachable))
+        ;; a sync call to a sync-lifted export, made before this task ever suspended
+        (call $g)
+        ;; a second async call: STARTED (1) again in definitions.py
+        (local.set $r2 (call $f))
+        (if (i32.ne (i32.and (local.get $r2) (i32.const 0xf)) (i32.const 1)) (then unreachable))
+        (if (i32.ne (call $subtask.cancel (i32.shr_u (local.get $r) (i32.const 4))) (i32.const 4)) (then unreachable))
+        (if (i32.ne (call $subtask.cancel (i32.shr_u (local.get $r2) (i32.const 4))) (i32.const 4)) (then unreachable))
+        (call $subtask.drop (i32.shr_u (local.get $r) (i32.const 4)))
+        (call $subtask.drop (i32.shr_u (local.get $r2) (i32.const 4)))
+        (call $task.return)
+        (i32.const 0 (; EXIT ;)))
+      (func (export "cb") (param i32 i32 i32) (result i32) unreachable))
+    (core instance $i (instantiate $m (with "" (instance
+      (export "f" (func $f))
+      (export "g" (func $g))
+      (export "subtask.cancel" (func $subtask.cancel))
+      (export "subtask.drop" (func $subtask.drop))
+      (export "task.return" (func $task.return))))))
+    (func (export "h") async (canon lift (core func $i "h") async (callback (core func $i "cb")))))
+
+  (component $C
+    (import "h" (func $h async))
+    (core func $h (canon lower (func $h) async))
+    (core func $task.return (canon task.return (result u32)))
+    (core module $m
+      (import "" "h" (func $h (result i32)))
+      (import "" "task.return" (func $task.return (param i32)))
+      (func (export "main") (result i32)
+        ;; $B runs to completion before $C regains control: RETURNED (2), no handle
+        (call $task.return (call $h))
+        (i32.const 0 (; EXIT ;)))
+      (func (export "cb") (param i32 i32 i32) (result i32) unreachable))
+    (core instance $i (instantiate $m (with "" (instance
+      (export "h" (func $h))
+      (export "task.return" (func $task.return))))))
+    (func (export "main") async (result u32) (canon lift (core func $i "main") async (callback (core func $i "cb")))))
+
+  (instance $a (instantiate $A))
+  (instance $b (instantiate $B (with "f" (func $a "f")) (with "g" (func $a "g"))))
+  (instance $c (instantiate $C (with "h" (func $b "h"))))
+  (export "main" (func $c "main")))
+
+(assert_return (invoke "main") (u32.const 2))