From e014bb1a9229814ca9779bb1c028ad6afb7cbb79 Mon Sep 17 00:00:00 2001 From: Joel Dice Date: Wed, 23 Sep 2026 15:18:29 -0600 Subject: [PATCH 1/5] add (off-by-default) `task-group-hook` feature This addresses #14247 and roughly resembles the API sketch Alex posted as a comment on that issue. Although the Component Model has no notion of a "task group", we define one here in order to enable embedders to associate guest->host calls with corresponding host->guest calls in a predictable way. See the doc comments on `TaskGroupHook` for details on what it means for a task to be part of a task group and what notifications a hook will receive for a given group. Note that, unlike with `CallHookHandler`, the functions in `TaskGroupHook` do not have a `StoreContextMut` parameter. Supporting such a parameter is technically possible, but would require converting a _lot_ of functions so that they accept `&mut dyn VMStore` parameters instead of `&mut StoreOpaque` parameters, including parts of the public API. This is something we can revisit later, if desired. Meanwhile, types implementing `TaskGroupHook` may need to use e.g. `Arc` to share state with the `T` in `Store`. This feature supersedes `StoreContextMut::async_call_stack`, which I've removed given that it is prone to use-after-free of a table slot as described in the aforementioned issue. Fixes #14247 --- Cargo.toml | 1 + crates/wasmtime/Cargo.toml | 2 + .../src/runtime/component/concurrent.rs | 328 +++++++++++++----- crates/wasmtime/src/runtime/component/mod.rs | 2 + tests/all/component_model/async.rs | 323 ++++++++--------- 5 files changed, 419 insertions(+), 237 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2a82d3745e9c..7f234b13c2ed 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -164,6 +164,7 @@ features = [ 'call-hook', 'memory-protection-keys', 'component-model-async', + 'task-group-hook', 'component-model-bytes', 'gc-copying', ] diff --git a/crates/wasmtime/Cargo.toml b/crates/wasmtime/Cargo.toml index 5bbbb645f20a..cbd69c44860e 100644 --- a/crates/wasmtime/Cargo.toml +++ b/crates/wasmtime/Cargo.toml @@ -431,3 +431,5 @@ rr = [ "component-model", "wasmtime-environ/rr" ] + +task-group-hook = [ "component-model", "component-model-async" ] diff --git a/crates/wasmtime/src/runtime/component/concurrent.rs b/crates/wasmtime/src/runtime/component/concurrent.rs index 9223b620ccfc..a5d8ae975bad 100644 --- a/crates/wasmtime/src/runtime/component/concurrent.rs +++ b/crates/wasmtime/src/runtime/component/concurrent.rs @@ -657,6 +657,76 @@ where } } +#[cfg(feature = "task-group-hook")] +struct TaskGroup { + ref_count: usize, +} + +#[cfg(feature = "task-group-hook")] +impl TableDebug for TaskGroup { + fn type_name() -> &'static str { + "TaskGroup" + } +} + +/// Represents a "task group" containing the "root" task of a host->guest call, +/// plus any subtasks transitively created by that task. +/// +/// See [TaskGroupHook] for details. +#[cfg(feature = "task-group-hook")] +#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)] +pub struct TaskGroupId(TableId); + +/// Trait for being notified by the runtime of activity concerning a "task group". +/// +/// The Component Model specification has no notion of a "task group", but we +/// define one here in order to enable embedders to associate guest->host calls +/// with corresponding host->guest calls in a predictable way. +/// +/// A `TaskGroupId` is allocated whenever a guest task is created for a +/// host->guest call, and `handle_start` is called. That task is considered the +/// "root" task for the task group, and any subtasks transitively created by it +/// will also be considered part of that task group. Whenever the runtime +/// switches (Component-Model-level) threads, it will call `handle_exit` for the +/// group to which the old thread belonged, if any, and call `handle_enter` for +/// the group to which the new thread belongs. Only once all the threads of all +/// those tasks have exited (and the guest has dropped any subtask handles +/// referring to any of those tasks) will the `TaskGroupId` be deallocated, at +/// which point `handle_finish` will be called. +/// +/// Note that a given `TaskGroupId` may be reused after `handle_finish` is +/// called, so implementations of this trait must take care to reset any state +/// associated with it. +/// +/// As of this writing, https://github.com/WebAssembly/component-model/pull/730 +/// (which adds `thread.set-task` and related intrinsics) has not yet been +/// merged. Once it has, and Wasmtime adds support for that feature, it will be +/// possible for guest threads to change their task; in that case, the thread +/// will effectively join whatever group the new task belongs to, which might +/// not be the same as that of the old task. In addition, the new +/// `thread.get-task` intrinsic will give the guest another way (besides subtask +/// handles) to keep tasks alive beyond the point when all their threads have +/// exited or switched tasks, in which case the group it belongs to will not be +/// disposed until all such tasks have been dropped using `task.drop`. +#[cfg(feature = "task-group-hook")] +pub trait TaskGroupHook: Send + Sync + 'static { + /// Handle notification that a new task group has been created (i.e. a + /// host->guest call has been prepared). + fn handle_start(&mut self, id: TaskGroupId) -> Result<()>; + /// Handle notification that the runtime has switched to a thread belonging to + /// the specified task group. + fn handle_enter(&mut self, id: TaskGroupId) -> Result<()>; + /// Handle notification that the runtime has switched away from a thread + /// belonging to the specified task group. + fn handle_exit(&mut self, id: TaskGroupId) -> Result<()>; + /// Handle notification that the specified task group has been disposed of + /// (i.e. the task created for the host->guest call for which the task group + /// was created has exited, along with any and all subtasks transitively + /// created by that task, and the guest has dropped any and all handles to + /// those tasks). + fn handle_finish(&mut self, id: TaskGroupId) -> Result<()>; +} + /// Represents parameter and result metadata for the caller side of a /// guest->guest call orchestrated by a fused adapter. enum CallerInfo { @@ -1059,6 +1129,17 @@ impl Store { { self.as_context_mut().spawn(task) } + + /// Set a [TaskGroupHook] for this store. + /// + /// This will overwrite any hook that was previously set. + #[cfg(feature = "task-group-hook")] + pub fn task_group_hook(&mut self, hook: impl TaskGroupHook) { + self.as_context_mut() + .0 + .concurrent_state_mut_without_forcing_current_thread() + .task_group_hook = Some(Box::new(hook)); + } } impl StoreContextMut<'_, T> { @@ -1677,42 +1758,6 @@ impl StoreContextMut<'_, T> { } } - /// Returns an iterator over the current async call stack defined by the - /// component model. - /// - /// This can be used, for example to correlate a host import call with which - /// root export task originally called it. - /// - /// Tasks are yielded "youngest first" where the first item in the iterator - /// is the current task, and the last item in the iterator is the original - /// call. - /// - /// # Omitted Tasks - /// - /// This stack might not contain an entry for every guest-to-guest component - /// call that is currently executing: if those guest-to-guest calls are not - /// asynchronous and provably cannot access async task state, then Wasmtime - /// can avoid materializing the guest task state for that call, and such - /// tasks will not have a `GuestTaskId` and will not appear in this stack - /// trace. - /// - /// The `GuestTaskId`s for host-to-guest and guest-to-host calls, however, - /// will always appear in this stack trace. - pub fn async_call_stack(&mut self) -> Result> { - let mut cur = Some(self.0.current_thread()?); - let state = self.0.concurrent_state_mut()?; - Ok(core::iter::from_fn(move || { - while let Some(t) = cur { - cur = state.parent(t); - if let Some(task) = t.guest_task() { - return Some(GuestTaskId(task)); - } - } - - None - })) - } - pub(crate) async fn start_instance( &mut self, instance: ModuleInstance, @@ -1924,6 +1969,21 @@ impl StoreOpaque { caller: self.materialize_host_task_id()?, } }; + + #[cfg(feature = "task-group-hook")] + let group = { + let state = self.concurrent_state_mut()?; + + match caller { + Caller::Guest { thread } => { + let group = state.get_mut(thread.task)?.group; + state.increment_group_ref_count(group)?; + group + } + Caller::Host { .. } => state.make_task_group()?, + } + }; + let state = self.concurrent_state_mut()?; let guest_thread = GuestTask::new( state, @@ -1939,6 +1999,8 @@ impl StoreOpaque { callee, callee_async_typed, true, + #[cfg(feature = "task-group-hook")] + group, )?; Instance::from_wasmtime(self, callee.instance).add_guest_thread_to_instance_table( @@ -2018,6 +2080,14 @@ impl StoreOpaque { } let caller = self.current_guest_thread()?; log::trace!("new deferred host task with caller {caller:?}"); + + #[cfg(feature = "task-group-hook")] + { + let state = self.concurrent_state_mut()?; + let group = state.get_mut(caller.task)?.group; + state.increment_group_ref_count(group)?; + } + self.set_thread(CurrentThread::DeferredHost(caller))?; let state = self.concurrent_state_mut()?; debug_assert!(state.deferred_host_call_context.is_none()); @@ -2039,6 +2109,13 @@ impl StoreOpaque { ) -> Result<()> { match original_task { Some(caller) => { + #[cfg(feature = "task-group-hook")] + { + let state = self.concurrent_state_mut()?; + let group = state.get_mut(caller.task)?.group; + state.decrement_group_ref_count(group)?; + } + self.set_thread(caller)?; if materialized_task.is_none() { let state = self.concurrent_state_mut()?; @@ -2082,6 +2159,21 @@ impl StoreOpaque { state.debug_assert_deferred_host_invariant(); let old_thread = mem::replace(&mut state.unforced_current_thread, thread); + #[cfg(feature = "task-group-hook")] + { + let old_group = old_thread.group(state)?; + let new_group = thread.group(state)?; + if let (true, Some(hook)) = ((old_group != new_group), &mut state.task_group_hook) { + if let Some(group) = old_group { + hook.handle_exit(TaskGroupId(group))?; + } + + if let Some(group) = new_group { + hook.handle_enter(TaskGroupId(group))?; + } + } + } + // First thing to do after swapping threads is updating the context // slots for this thread within the store. This restores the behavior of // `context.{get,set}`. This involves taking the old state out of the @@ -2440,12 +2532,13 @@ impl StoreOpaque { for waitable in mem::take(&mut state.get_mut(sync_call_set)?.ready) { if let Some(Event::Subtask { status: Status::Returned | Status::ReturnCancelled, - }) = waitable.common(state)?.event + }) = waitable.common(self.concurrent_state_mut()?)?.event { - waitable.delete_from(state)?; + waitable.delete_from(self)?; } } + let state = self.concurrent_state_mut()?; state.delete(guest_thread.thread)?; state.delete(sync_call_set)?; let task = state.get_mut(guest_thread.task)?; @@ -2471,7 +2564,7 @@ impl StoreOpaque { match cleanup_task { CleanupTask::Yes => { if ready_to_delete { - Waitable::Guest(guest_thread.task).delete_from(state)?; + Waitable::Guest(guest_thread.task).delete_from(self)?; } } CleanupTask::No => {} @@ -2543,7 +2636,6 @@ impl StoreOpaque { Ok(match self.current_thread()? { CurrentThread::Guest(id) => Some(CurrentScope::Id(Scope::Id(id.task.rep()))), - CurrentThread::GuestTask(id) => Some(CurrentScope::Id(Scope::Id(id.rep()))), CurrentThread::Host(id) => Some(CurrentScope::Id(Scope::HostId(id.rep()))), CurrentThread::DeferredHost(_) => Some(CurrentScope::DeferredHost), CurrentThread::None => return Ok(None), @@ -3125,6 +3217,15 @@ impl Instance { let return_ = SendSyncPtr::new(return_); let token = StoreToken::new(store.as_context_mut()); let old_thread = store.0.current_guest_thread()?; + + #[cfg(feature = "task-group-hook")] + let group = { + let state = store.0.concurrent_state_mut()?; + let group = state.get_mut(old_thread.task)?.group; + state.increment_group_ref_count(group)?; + group + }; + let state = store.0.concurrent_state_mut()?; debug_assert_eq!( @@ -3238,6 +3339,8 @@ impl Instance { // We don't know whether the callee export was lifted sync or async // yet, but we'll update this in `start_call`: false, + #[cfg(feature = "task-group-hook")] + group, )?; // Make the new thread the current one so that `Self::start_call` knows @@ -3478,7 +3581,7 @@ impl Instance { } if task.exited && task.ready_to_delete() { - Waitable::Guest(guest_thread.task).delete_from(state)?; + Waitable::Guest(guest_thread.task).delete_from(store.0)?; } } } @@ -3853,6 +3956,13 @@ impl Instance { bail_bug!("invalid state for callee in `subtask.drop`") } } + + #[cfg(feature = "task-group-hook")] + { + let group = task.group; + concurrent_state.decrement_group_ref_count(group)?; + } + (Waitable::Host(id), true) } else { let id = TableId::::new(rep); @@ -3875,7 +3985,7 @@ impl Instance { } if delete { - waitable.delete_from(concurrent_state)?; + waitable.delete_from(store)?; } log::trace!("subtask_drop {waitable:?} (handle {task_id})"); @@ -4968,19 +5078,14 @@ where pub(crate) struct HostTask { common: WaitableCommon, - /// The calling guest task of this host task. Note that this is only used - /// for backtrace purposes and is additionally not updated when the guest - /// task finishes. - /// - /// TODO: this field should probably get deleted entirely and the - /// backtrace-purposes here should move to an auxiliary table. - caller: TableId, - /// State of borrows/etc the host needs to track. Used when the guest passes /// borrows to the host, for example. call_context: CallContext, state: HostTaskState, + + #[cfg(feature = "task-group-hook")] + group: TableId, } enum HostTaskState { @@ -5007,12 +5112,16 @@ enum HostTaskState { } impl HostTask { - fn new(caller: TableId, state: HostTaskState) -> Self { + fn new( + state: HostTaskState, + #[cfg(feature = "task-group-hook")] group: TableId, + ) -> Self { Self { common: WaitableCommon::default(), call_context: CallContext::default(), - caller, state, + #[cfg(feature = "task-group-hook")] + group, } } } @@ -5296,6 +5405,9 @@ pub(crate) struct GuestTask { async_lifted: bool, decremented_interesting_task_count: bool, + + #[cfg(feature = "task-group-hook")] + group: TableId, } impl GuestTask { @@ -5341,6 +5453,7 @@ impl GuestTask { instance: RuntimeInstance, async_typed: bool, async_lifted: bool, + #[cfg(feature = "task-group-hook")] group: TableId, ) -> Result { let host_future_state = match &caller { Caller::Guest { .. } => HostFutureState::NotApplicable, @@ -5374,6 +5487,8 @@ impl GuestTask { async_typed, async_lifted, decremented_interesting_task_count: false, + #[cfg(feature = "task-group-hook")] + group, })?; let new_thread = GuestThread::new_implicit(state, task)?; let thread = state.push(new_thread)?; @@ -5555,17 +5670,21 @@ impl Waitable { Ok(()) } - /// Remove this waitable from the instance's rep table. - fn delete_from(&self, state: &mut ConcurrentState) -> Result<()> { + /// Remove this waitable from the store's rep table. + fn delete_from(&self, store: &mut StoreOpaque) -> Result<()> { match self { Self::Host(task) => { log::trace!("delete host task {task:?}"); - state.delete(*task)?; + store.concurrent_state_mut()?.delete(*task)?; } Self::Guest(task) => { log::trace!("delete guest task {task:?}"); + let state = store.concurrent_state_mut()?; let task = state.delete(*task)?; + #[cfg(feature = "task-group-hook")] + state.decrement_group_ref_count(task.group)?; + // When a guest task is created it increments the // `ConcurrentState::interesting_tasks` counter, and that needs // to be paired with a decrement. There are a few situations in @@ -5575,7 +5694,7 @@ impl Waitable { debug_assert!(task.decremented_interesting_task_count); } Self::Transmit(task) => { - state.delete(*task)?; + store.concurrent_state_mut()?.delete(*task)?; } } @@ -5660,10 +5779,6 @@ pub(crate) enum CurrentThread { /// The currently running thread is a host call whose task has not yet been /// materialized. The contained ID identifies its guest caller. DeferredHost(QualifiedThreadId), - /// A bit of a kludge to get `StoreOpaque::parent` working with backtraces - /// and this serves as the parent node of a `Host` task. This ideally would - /// get removed in favor of separate backtrace storage. - GuestTask(TableId), /// There is no currently running thread because we are in the main event /// loop or concurrency is disabled. None, @@ -5680,7 +5795,6 @@ impl CurrentThread { fn guest_task(&self) -> Option> { match self { Self::Guest(id) => Some(id.task), - Self::GuestTask(id) => Some(*id), _ => None, } } @@ -5688,6 +5802,17 @@ impl CurrentThread { fn is_none(&self) -> bool { matches!(self, Self::None) } + + #[cfg(feature = "task-group-hook")] + fn group(&self, state: &mut ConcurrentState) -> Result>> { + Ok(match self { + Self::Guest(thread) | Self::DeferredHost(thread) => { + Some(state.get_mut(thread.task)?.group) + } + Self::Host(task) => Some(state.get_mut(*task)?.group), + Self::None => None, + }) + } } impl From for CurrentThread { @@ -5802,6 +5927,10 @@ pub struct ConcurrentState { /// Whether the `StoreContextMut::poll_until` event loop is running. event_loop_running: bool, + + /// See [TaskGroupHook]. + #[cfg(feature = "task-group-hook")] + task_group_hook: Option>, } impl Default for ConcurrentState { @@ -5823,6 +5952,8 @@ impl Default for ConcurrentState { interesting_tasks_empty_waker: None, ready_for_concurrent_call_waker: None, event_loop_running: false, + #[cfg(feature = "task-group-hook")] + task_group_hook: None, } } } @@ -5944,6 +6075,9 @@ impl ConcurrentState { interesting_tasks_empty_waker: _, ready_for_concurrent_call_waker: _, event_loop_running: _, + + #[cfg(feature = "task-group-hook")] + task_group_hook: _, } = self; for entry in table.get_mut().iter_mut() { @@ -6180,24 +6314,6 @@ impl ConcurrentState { self.table.get_mut() } - /// Returns the parent thread, if any, of `cur`. - fn parent(&mut self, cur: CurrentThread) -> Option { - let task = match cur { - CurrentThread::GuestTask(task) => task, - CurrentThread::Guest(thread) => thread.task, - CurrentThread::Host(id) => { - return Some(CurrentThread::GuestTask(self.get_mut(id).ok()?.caller)); - } - CurrentThread::DeferredHost(caller) => return Some(caller.into()), - CurrentThread::None => return None, - }; - let task = self.get_mut(task).ok()?; - Some(match task.caller { - Caller::Host { caller, .. } => caller.map_or(CurrentThread::None, CurrentThread::Host), - Caller::Guest { thread } => thread.into(), - }) - } - fn debug_assert_deferred_host_invariant(&self) { debug_assert_eq!( self.deferred_host_call_context.is_some(), @@ -6213,8 +6329,18 @@ impl ConcurrentState { thread => return Ok(thread), }; + #[cfg(feature = "task-group-hook")] + let group = self.get_mut(caller.task)?.group; + + #[cfg(not(feature = "task-group-hook"))] + let _ = caller; + // Push first so allocation failure leaves the deferred state intact. - let task = self.push(HostTask::new(caller.task, HostTaskState::CalleeStarted))?; + let task = self.push(HostTask::new( + HostTaskState::CalleeStarted, + #[cfg(feature = "task-group-hook")] + group, + ))?; let call_context = self .deferred_host_call_context .take() @@ -6232,7 +6358,7 @@ impl ConcurrentState { match self.materialize_host_task()? { CurrentThread::Host(id) => Ok(Some(id)), CurrentThread::None => Ok(None), - CurrentThread::Guest(_) | CurrentThread::GuestTask(_) => { + CurrentThread::Guest(_) => { bail_bug!("tried to materialize a host task id from a guest thread") } CurrentThread::DeferredHost(_) => { @@ -6249,6 +6375,39 @@ impl ConcurrentState { _ => bail_bug!("current scope is not a deferred host scope"), } } + + #[cfg(feature = "task-group-hook")] + fn make_task_group(&mut self) -> Result> { + let group = self.push(TaskGroup { ref_count: 1 })?; + if let Some(hook) = &mut self.task_group_hook { + hook.handle_start(TaskGroupId(group))?; + } + log::trace!("new {group:?}"); + Ok(group) + } + + #[cfg(feature = "task-group-hook")] + fn increment_group_ref_count(&mut self, group: TableId) -> Result<()> { + let count = &mut self.get_mut(group)?.ref_count; + *count += 1; + log::trace!("increment {group:?} to {count}"); + Ok(()) + } + + #[cfg(feature = "task-group-hook")] + fn decrement_group_ref_count(&mut self, group: TableId) -> Result<()> { + let count = &mut self.get_mut(group)?.ref_count; + assert!(*count >= 1); + *count -= 1; + log::trace!("decrement {group:?} to {count}"); + if *count == 0 { + self.delete(group)?; + if let Some(hook) = &mut self.task_group_hook { + hook.handle_finish(TaskGroupId(group))?; + } + } + Ok(()) + } } /// Provide a type hint to compiler about the shape of a parameter lower @@ -6371,7 +6530,7 @@ impl TaskId { task.ready_to_delete() }; if delete { - Waitable::Guest(self.task).delete_from(store.concurrent_state_mut()?)? + Waitable::Guest(self.task).delete_from(store)? } Ok(()) } @@ -6419,6 +6578,9 @@ pub(crate) fn prepare_call( let caller = store.0.materialize_host_task_id()?; let state = store.0.concurrent_state_mut()?; + #[cfg(feature = "task-group-hook")] + let group = state.make_task_group()?; + let (tx, rx) = oneshot::channel(); let instance = handle.instance().runtime_instance(component_instance); @@ -6453,6 +6615,8 @@ pub(crate) fn prepare_call( instance, async_typed, async_lifted, + #[cfg(feature = "task-group-hook")] + group, )?; Ok(PreparedCall { diff --git a/crates/wasmtime/src/runtime/component/mod.rs b/crates/wasmtime/src/runtime/component/mod.rs index 7bf166b757ad..0190f8bab364 100644 --- a/crates/wasmtime/src/runtime/component/mod.rs +++ b/crates/wasmtime/src/runtime/component/mod.rs @@ -125,6 +125,8 @@ pub use self::concurrent::{ StreamAny, StreamConsumer, StreamProducer, StreamReader, StreamResult, TypedFuncCallConcurrent, VMComponentAsyncStore, VecBuffer, WriteBuffer, }; +#[cfg(feature = "task-group-hook")] +pub use self::concurrent::{TaskGroupHook, TaskGroupId}; pub use self::func::{ ComponentNamedList, ComponentType, Func, Lift, Lower, TypedFunc, WasmList, WasmStr, }; diff --git a/tests/all/component_model/async.rs b/tests/all/component_model/async.rs index 43c3ab23a0cc..0d80273aca88 100644 --- a/tests/all/component_model/async.rs +++ b/tests/all/component_model/async.rs @@ -1,7 +1,13 @@ use crate::async_functions::{PollOnce, execute_across_threads}; +use futures::stream::{FuturesUnordered, TryStreamExt as _}; +use std::collections::HashMap; +use std::iter; +use std::mem; use std::pin::Pin; +use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; use wasmtime::Result; +use wasmtime::component::{TaskGroupHook, TaskGroupId}; use wasmtime::{AsContextMut, Config, Engine, Store, StoreContextMut, Trap, component::*}; use wasmtime_component_util::REALLOC_AND_FREE; @@ -973,7 +979,8 @@ async fn bytes_stream_producer() -> Result<()> { #[tokio::test] #[cfg_attr(miri, ignore)] -async fn async_call_stack() -> Result<()> { +async fn task_group_hook() -> Result<()> { + _ = env_logger::try_init(); let mut config = Config::new(); config.wasm_component_model_async(true); let engine = Engine::new(&config)?; @@ -982,185 +989,191 @@ async fn async_call_stack() -> Result<()> { &engine, r#" (component - (import "a" (func $a)) - (core func $a (canon lower (func $a))) + (import "a" (func $a async)) + (core func $a (canon lower (func $a) async)) (core module $a - (import "" "a" (func $a)) - (func (export "a") call $a) + (import "" "a" (func $a (result i32))) + (import "" "waitable-set.new" (func $waitable-set.new (result i32))) + (import "" "waitable.join" (func $waitable.join (param i32 i32))) + (import "" "subtask.drop" (func $subtask.drop (param i32))) + (import "" "task.return" (func $task.return)) + (func (export "a") (result i32) + (local $ret i32) + (local $subtask i32) + (local $set i32) + + (local.set $ret (call $a)) + (if (i32.ne (i32.const 1 (; STARTED ;)) (i32.and (local.get $ret) (i32.const 0xf))) + (then unreachable)) + (local.set $subtask (i32.shr_u (local.get $ret) (i32.const 4))) + (local.set $set (call $waitable-set.new)) + (call $waitable.join (local.get $subtask) (local.get $set)) + (i32.or (i32.const 2 (; WAIT ;)) (i32.shl (local.get $set) (i32.const 4))) + ) + + (func (export "a-callback") (param $event_code i32) (param $waitable i32) (param $payload i32) (result i32) + (if (i32.ne (local.get $event_code) (i32.const 1 (; SUBTASK ;))) + (then unreachable)) + (if (i32.ne (local.get $payload) (i32.const 2 (; RETURNED ;))) + (then unreachable)) + (call $waitable.join (local.get $waitable) (i32.const 0)) + (call $subtask.drop (local.get $waitable)) + (call $task.return) + (i32.const 0 (; EXIT ;)) + ) ) + (canon waitable-set.new (core func $waitable-set.new)) + (canon waitable.join (core func $waitable.join)) + (canon subtask.drop (core func $subtask.drop)) + (canon task.return (core func $task.return)) (core instance $a (instantiate $a - (with "" (instance (export "a" (func $a)))) + (with "" (instance + (export "a" (func $a)) + (export "waitable-set.new" (func $waitable-set.new)) + (export "waitable.join" (func $waitable.join)) + (export "subtask.drop" (func $subtask.drop)) + (export "task.return" (func $task.return)) + )) )) - (func (export "a") async (canon lift (core func $a "a"))) + (func (export "a") async (canon lift (core func $a "a") async (callback (core func $a "a-callback")))) ) "#, )?; - let mut linker = Linker::new(&engine); - linker.root().func_wrap( - "a", - |mut store: StoreContextMut>, (): ()| { - let stack = store.async_call_stack()?.collect::>(); - assert_eq!(stack, [store.data().unwrap()]); - Ok(()) - }, - )?; - let mut store = Store::new(&engine, None); - let instance = linker.instantiate_async(&mut store, &component).await?; - let func = instance.get_typed_func::<(), ()>(&mut store, "a")?; + #[derive(Copy, Clone, PartialEq, Eq, Debug)] + enum Event { + Start, + PollHostCall, + Finish, + } - let call = func.start_call_concurrent(&mut store, ())?; - *store.data_mut() = Some(call.task()); - store - .run_concurrent(async |store| func.finish_call_concurrent(store, call).await) - .await??; + #[derive(Default)] + struct HookInner { + events: Vec<(Event, usize)>, + next_id: usize, + ids: HashMap, + entered: Option, + } - let component = Component::new( - &engine, - r#" - (component - (import "a" (func $a)) - (component $a - (import "a" (func $a)) - (core func $a (canon lower (func $a))) - - (core module $a - (import "" "a" (func $a)) - (func (export "a") call $a) - ) - (core instance $a (instantiate $a - (with "" (instance (export "a" (func $a)))) - )) - (func (export "a") (canon lift (core func $a "a"))) - ) + #[derive(Clone)] + struct Hook(Arc>); + + impl TaskGroupHook for Hook { + fn handle_start(&mut self, group: TaskGroupId) -> Result<()> { + let mut inner = self.0.try_lock().unwrap(); + let id = inner.next_id; + inner.next_id += 1; + // This `group` should not be equal to any other + // started-and-not-yet-finished group: + assert!(inner.ids.insert(group, id).is_none()); + inner.events.push((Event::Start, id)); + Ok(()) + } - (instance $a (instantiate $a (with "a" (func $a)))) - (instance $b (instantiate $a (with "a" (func $a "a")))) - (export "a" (func $b "a")) - ) - "#, - )?; + fn handle_enter(&mut self, group: TaskGroupId) -> Result<()> { + let mut inner = self.0.try_lock().unwrap(); + let id = *inner.ids.get(&group).unwrap(); + // We should not have entered-but-not-yet-exited this or any other + // group: + assert!(inner.entered.replace(id).is_none()); + Ok(()) + } - let mut linker = Linker::new(&engine); - linker.root().func_wrap( - "a", - |mut store: StoreContextMut>, (): ()| { - let stack = store.async_call_stack()?.collect::>(); - assert_eq!(stack.len(), 2); - assert_eq!(stack.last(), store.data().as_ref()); + fn handle_exit(&mut self, group: TaskGroupId) -> Result<()> { + let mut inner = self.0.try_lock().unwrap(); + let id = *inner.ids.get(&group).unwrap(); + // We should have entered-but-not-yet-exited this group: + assert_eq!(Some(id), inner.entered); + inner.entered = None; Ok(()) - }, - )?; - let instance = linker.instantiate_async(&mut store, &component).await?; - let func = instance.get_typed_func::<(), ()>(&mut store, "a")?; + } - let call = func.start_call_concurrent(&mut store, ())?; - *store.data_mut() = Some(call.task()); - store - .run_concurrent(async |store| func.finish_call_concurrent(store, call).await) - .await??; - Ok(()) -} + fn handle_finish(&mut self, group: TaskGroupId) -> Result<()> { + let mut inner = self.0.try_lock().unwrap(); + let id = inner.ids.remove(&group).unwrap(); + inner.events.push((Event::Finish, id)); + Ok(()) + } + } -#[tokio::test] -#[cfg_attr(miri, ignore)] -async fn async_call_stack_omits_transparent_adapters() -> Result<()> { - async fn call_stack(taint: bool) -> Result<(Vec, GuestTaskId)> { - let mut config = Config::new(); - config.wasm_component_model_async(true); - let engine = Engine::new(&config)?; - - let (taint_canon, taint_import, taint_arg) = if taint { - ( - r#"(core func $ctx (canon context.get i32 0))"#, - r#"(import "" "ctx" (func $ctx (result i32)))"#, - r#"(export "ctx" (func $ctx))"#, - ) - } else { - ("", "", "") - }; - - let component = Component::new( - &engine, - &format!( - r#" - (component - (import "a" (func $a)) + let hook = Hook(Arc::new(Mutex::new(HookInner::default()))); + let yield_count = 5; - ;; Lowers the host import, so this is never thread-transparent. - (component $Deep - (import "a" (func $a)) - (core func $a (canon lower (func $a))) - (core module $m - (import "" "a" (func $a)) - (func (export "a") call $a)) - (core instance $m (instantiate $m - (with "" (instance (export "a" (func $a)))))) - (func (export "a") (canon lift (core func $m "a"))) - ) + let mut linker = Linker::new(&engine); + linker.root().func_wrap_concurrent("a", { + let hook = hook.clone(); + move |_, ()| { + let hook = hook.clone(); + Box::pin(async move { + let mut previous_id = None; + let mut on_poll = move || { + let mut inner = hook.0.try_lock().unwrap(); + // We should only ever be polled after a `handle_enter` and + // before a `handle_exit` event: + let id = inner.entered.unwrap(); + inner.events.push((Event::PollHostCall, id)); + + // Should see the same task group on every poll: + if let Some(previous_id) = previous_id { + assert_eq!(previous_id, id); + } else { + previous_id = Some(id); + } - ;; Declares nothing but a `canon lower` of an imported *lifted*, - ;; non-`async` function, so an instance of this component is - ;; thread-transparent -- unless the taint below is present. - (component $Mid - (import "a" (func $a)) - {taint_canon} - (core func $a (canon lower (func $a))) - (core module $m - (import "" "a" (func $a)) - {taint_import} - (func (export "a") call $a)) - (core instance $m (instantiate $m - (with "" (instance (export "a" (func $a)) {taint_arg})))) - (func (export "a") (canon lift (core func $m "a"))) - ) + id + }; + + on_poll(); + for _ in 0..yield_count { + tokio::task::yield_now().await; + on_poll(); + } - (instance $deep (instantiate $Deep (with "a" (func $a)))) - (instance $mid (instantiate $Mid (with "a" (func $deep "a")))) - (instance $top (instantiate $Mid (with "a" (func $mid "a")))) - (export "a" (func $top "a")) - ) - "# - ), - )?; - - let mut linker = Linker::new(&engine); - linker.root().func_wrap( - "a", - |mut store: StoreContextMut>>, (): ()| { - let stack = store.async_call_stack()?.collect::>(); - *store.data_mut() = Some(stack); Ok(()) - }, - )?; + }) + } + })?; + + let mut store = Store::new(&engine, ()); + let instance = linker.instantiate_async(&mut store, &component).await?; + let func = instance.get_typed_func::<(), ()>(&mut store, "a")?; - let mut store = Store::new(&engine, None); - let instance = linker.instantiate_async(&mut store, &component).await?; - let func = instance.get_typed_func::<(), ()>(&mut store, "a")?; + store.task_group_hook(hook.clone()); - let call = func.start_call_concurrent(&mut store, ())?; - let root = call.task(); - store - .run_concurrent(async |store| func.finish_call_concurrent(store, call).await) - .await??; + let call_count = 3; + + // Call the guest export `call_count` times concurrently, recording the + // events in `hook`: + store + .run_concurrent(async |store| { + iter::repeat_with(|| func.call_concurrent(store, ())) + .take(call_count) + .collect::>() + .try_collect::<()>() + .await + }) + .await??; + + // Group the events by task group, preserving order: + let mut map = HashMap::<_, Vec<_>>::new(); + for (event, id) in mem::take(&mut hook.0.try_lock().unwrap().events) { + map.entry(id).or_default().push(event); + } + + // Assert that the events start with a `Start`, followed by the expected + // number of `PollHostCall`s, followed by a `Finish`: + let expected = iter::once(Event::Start) + .chain(iter::repeat(Event::PollHostCall).take(yield_count + 1)) + .chain(Some(Event::Finish)) + .collect::>(); - Ok((store.data_mut().take().unwrap(), root)) + for id in 0..call_count { + assert_eq!(&map.get(&id).unwrap()[..], &expected[..]) } - // With the intermediate instances tainted into opacity, all three - // guest-to-guest calls materialize a task and all three show up. - let (stack, root) = call_stack(true).await?; - assert_eq!(stack.len(), 3); - assert_eq!(stack.last(), Some(&root)); - - // Without the taint, `$top -> $mid` is a thread-transparent call, its task - // is never materialized, and so it is omitted from the stack. The root - // host-to-guest call is always materialized, so it is still reported, and - // `$mid -> $deep` still is too since `$deep` is opaque. - let (stack, root) = call_stack(false).await?; - assert_eq!(stack.len(), 2); - assert_eq!(stack.last(), Some(&root)); + assert_eq!(None, hook.0.try_lock().unwrap().entered); + assert!(hook.0.try_lock().unwrap().ids.is_empty()); Ok(()) } From 67c201f820d27d4e0a159257de1cb49c90fde447 Mon Sep 17 00:00:00 2001 From: Joel Dice Date: Fri, 25 Sep 2026 16:56:28 -0600 Subject: [PATCH 2/5] address review feedback - Ensure `TaskGroupHook` events are delivered even if the guest traps and/or leaks subtasks - Expand doc comments per review feedback - Move most `TaskGroupHook`-related code into its own module - Add `task_group_hook` function to `StoreContextMut` - Deduplicate code and make it harder to forget to increment or decrement `TaskGroup` ref counts --- Cargo.toml | 1 + crates/wasi-http/Cargo.toml | 1 + crates/wasi-http/src/handler.rs | 23 +- .../src/runtime/component/concurrent.rs | 276 ++++-------------- .../src/runtime/component/concurrent/func.rs | 41 ++- .../component/concurrent/task_group_hook.rs | 208 +++++++++++++ crates/wasmtime/src/runtime/component/mod.rs | 4 +- .../src/runtime/component/resource_table.rs | 13 +- .../wasmtime/src/runtime/component/store.rs | 5 + src/commands/serve.rs | 6 +- tests/all/component_model/async.rs | 113 +++++-- 11 files changed, 406 insertions(+), 285 deletions(-) create mode 100644 crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs diff --git a/Cargo.toml b/Cargo.toml index 7f234b13c2ed..be252f11aa58 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -556,6 +556,7 @@ component-model-async = [ "dep:futures", ] rr = ["wasmtime/rr", "component-model", "wasmtime-cli-flags/rr", "run"] +task-group-hook = ["wasmtime/task-group-hook", "wasmtime-wasi-http/task-group-hook"] # Enabled in release builds. We cannot enable the built-in gdbstub component in # the default build because the published-to-crates.io crates cannot build the diff --git a/crates/wasi-http/Cargo.toml b/crates/wasi-http/Cargo.toml index 664269d316f5..b031693db08e 100644 --- a/crates/wasi-http/Cargo.toml +++ b/crates/wasi-http/Cargo.toml @@ -20,6 +20,7 @@ default-send-request = ["dep:tokio-rustls", "dep:rustls", "dep:webpki-roots"] p2 = ["wasmtime-wasi/p2"] p3 = ["wasmtime-wasi/p3", "dep:tokio-util"] component-model-async = ["futures/alloc", "wasmtime/component-model-async"] +task-group-hook = ["wasmtime/task-group-hook"] [dependencies] async-trait = { workspace = true } diff --git a/crates/wasi-http/src/handler.rs b/crates/wasi-http/src/handler.rs index 937e2a41e96e..e2696f500cb7 100644 --- a/crates/wasi-http/src/handler.rs +++ b/crates/wasi-http/src/handler.rs @@ -30,7 +30,9 @@ use std::sync::{ use std::task::{Context, Poll}; use std::time::Instant; use tokio::sync::Notify; -use wasmtime::component::{Accessor, GuestTaskId, Resource, TypedFuncCallConcurrent}; +#[cfg(feature = "task-group-hook")] +use wasmtime::component::TaskGroupId; +use wasmtime::component::{Accessor, Resource, TypedFuncCallConcurrent}; #[cfg(feature = "p2")] use wasmtime::error::Context as _; use wasmtime::{AsContextMut, Result, Store, StoreContextMut, format_err}; @@ -184,8 +186,9 @@ pub trait WorkerState: 'static + Send + Sync { /// Notification that a request has been accepted by the worker. /// /// This method can be used to record anything within `store`, if necessary. - /// The `task` corresponding to the component-model-level async task about - /// to be created is additionally passed here. + /// If the `task-group-hook` feature is enabled, the task group + /// corresponding to the component-model-level async task about to be + /// created is additionally passed here. /// /// If the future returned by this function resolves before the guest has /// produced a response, the request will be considered "expired" and the @@ -217,7 +220,7 @@ pub trait WorkerState: 'static + Send + Sync { &self, store: StoreContextMut<'_, Self::StoreData>, data: Self::RequestData, - task: GuestTaskId, + #[cfg(feature = "task-group-hook")] task_group: TaskGroupId, ) -> Pin + 'static + Send + Sync>>; /// Dispose of the store belonging to the now-exited worker. @@ -508,7 +511,8 @@ where let expiration = dropper.state.on_request_start( store.as_context_mut(), request_data, - prepared.task(), + #[cfg(feature = "task-group-hook")] + prepared.group(), ); Ok((prepared, expiration)) } @@ -948,7 +952,7 @@ where } /// Representation of a "prepared" call for a guest, used to extract the -/// `GuestTaskId` before actually executing any handlers. +/// `TaskGroupId` before actually executing any handlers. /// /// Right now this is a bit gross since it has to type out a bunch of types by /// hand. @@ -1043,12 +1047,13 @@ impl<'a, T: Send> Prepared<'a, T> { } } - fn task(&self) -> GuestTaskId { + #[cfg(feature = "task-group-hook")] + fn group(&self) -> TaskGroupId { match self { #[cfg(feature = "p3")] - Prepared::P3 { call, .. } => call.task(), + Prepared::P3 { call, .. } => call.group(), #[cfg(feature = "p2")] - Prepared::P2 { call, .. } => call.task(), + Prepared::P2 { call, .. } => call.group(), } } diff --git a/crates/wasmtime/src/runtime/component/concurrent.rs b/crates/wasmtime/src/runtime/component/concurrent.rs index a5d8ae975bad..7f2285ccd709 100644 --- a/crates/wasmtime/src/runtime/component/concurrent.rs +++ b/crates/wasmtime/src/runtime/component/concurrent.rs @@ -86,6 +86,8 @@ use futures::channel::oneshot; use futures::stream::{FuturesUnordered, StreamExt}; use futures_and_streams::{FlatAbi, ReturnCode, TransmitHandle, TransmitIndex}; use table::{TableDebug, TableId}; +#[cfg(feature = "task-group-hook")] +use task_group_hook::TaskGroup; use wasmtime_environ::component::{ CanonicalAbiInfo, CanonicalOptions, CanonicalOptionsDataModel, MAX_FLAT_PARAMS, MAX_FLAT_RESULTS, OptionsIndex, PREPARE_ASYNC_NO_RESULT, PREPARE_ASYNC_WITH_RESULT, @@ -107,6 +109,8 @@ pub use futures_and_streams::{ StreamProducer, StreamReader, StreamResult, VecBuffer, WriteBuffer, }; pub(crate) use futures_and_streams::{ResourcePair, lower_error_context_to_index}; +#[cfg(feature = "task-group-hook")] +pub use task_group_hook::{TaskGroupHook, TaskGroupId}; mod abort; mod error_contexts; @@ -114,6 +118,8 @@ mod func; mod future_stream_any; mod futures_and_streams; pub(crate) mod table; +#[cfg(feature = "task-group-hook")] +mod task_group_hook; pub(crate) mod tls; /// Constant defined in the Component Model spec to indicate that the async @@ -657,76 +663,6 @@ where } } -#[cfg(feature = "task-group-hook")] -struct TaskGroup { - ref_count: usize, -} - -#[cfg(feature = "task-group-hook")] -impl TableDebug for TaskGroup { - fn type_name() -> &'static str { - "TaskGroup" - } -} - -/// Represents a "task group" containing the "root" task of a host->guest call, -/// plus any subtasks transitively created by that task. -/// -/// See [TaskGroupHook] for details. -#[cfg(feature = "task-group-hook")] -#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)] -pub struct TaskGroupId(TableId); - -/// Trait for being notified by the runtime of activity concerning a "task group". -/// -/// The Component Model specification has no notion of a "task group", but we -/// define one here in order to enable embedders to associate guest->host calls -/// with corresponding host->guest calls in a predictable way. -/// -/// A `TaskGroupId` is allocated whenever a guest task is created for a -/// host->guest call, and `handle_start` is called. That task is considered the -/// "root" task for the task group, and any subtasks transitively created by it -/// will also be considered part of that task group. Whenever the runtime -/// switches (Component-Model-level) threads, it will call `handle_exit` for the -/// group to which the old thread belonged, if any, and call `handle_enter` for -/// the group to which the new thread belongs. Only once all the threads of all -/// those tasks have exited (and the guest has dropped any subtask handles -/// referring to any of those tasks) will the `TaskGroupId` be deallocated, at -/// which point `handle_finish` will be called. -/// -/// Note that a given `TaskGroupId` may be reused after `handle_finish` is -/// called, so implementations of this trait must take care to reset any state -/// associated with it. -/// -/// As of this writing, https://github.com/WebAssembly/component-model/pull/730 -/// (which adds `thread.set-task` and related intrinsics) has not yet been -/// merged. Once it has, and Wasmtime adds support for that feature, it will be -/// possible for guest threads to change their task; in that case, the thread -/// will effectively join whatever group the new task belongs to, which might -/// not be the same as that of the old task. In addition, the new -/// `thread.get-task` intrinsic will give the guest another way (besides subtask -/// handles) to keep tasks alive beyond the point when all their threads have -/// exited or switched tasks, in which case the group it belongs to will not be -/// disposed until all such tasks have been dropped using `task.drop`. -#[cfg(feature = "task-group-hook")] -pub trait TaskGroupHook: Send + Sync + 'static { - /// Handle notification that a new task group has been created (i.e. a - /// host->guest call has been prepared). - fn handle_start(&mut self, id: TaskGroupId) -> Result<()>; - /// Handle notification that the runtime has switched to a thread belonging to - /// the specified task group. - fn handle_enter(&mut self, id: TaskGroupId) -> Result<()>; - /// Handle notification that the runtime has switched away from a thread - /// belonging to the specified task group. - fn handle_exit(&mut self, id: TaskGroupId) -> Result<()>; - /// Handle notification that the specified task group has been disposed of - /// (i.e. the task created for the host->guest call for which the task group - /// was created has exited, along with any and all subtasks transitively - /// created by that task, and the guest has dropped any and all handles to - /// those tasks). - fn handle_finish(&mut self, id: TaskGroupId) -> Result<()>; -} - /// Represents parameter and result metadata for the caller side of a /// guest->guest call orchestrated by a fused adapter. enum CallerInfo { @@ -1129,17 +1065,6 @@ impl Store { { self.as_context_mut().spawn(task) } - - /// Set a [TaskGroupHook] for this store. - /// - /// This will overwrite any hook that was previously set. - #[cfg(feature = "task-group-hook")] - pub fn task_group_hook(&mut self, hook: impl TaskGroupHook) { - self.as_context_mut() - .0 - .concurrent_state_mut_without_forcing_current_thread() - .task_group_hook = Some(Box::new(hook)); - } } impl StoreContextMut<'_, T> { @@ -1378,11 +1303,17 @@ impl StoreContextMut<'_, T> { // SAFETY: We never move `dropper` nor its `value` field. let future = unsafe { Pin::new_unchecked(dropper.value.deref_mut()) }; - dropper + let result = dropper .store .as_context_mut() .poll_until(future, trap_on_idle) - .await + .await; + + if result.is_err() { + dropper.store.0.set_trapped(); + } + + result } /// Run this store's event loop. @@ -1970,20 +1901,6 @@ impl StoreOpaque { } }; - #[cfg(feature = "task-group-hook")] - let group = { - let state = self.concurrent_state_mut()?; - - match caller { - Caller::Guest { thread } => { - let group = state.get_mut(thread.task)?.group; - state.increment_group_ref_count(group)?; - group - } - Caller::Host { .. } => state.make_task_group()?, - } - }; - let state = self.concurrent_state_mut()?; let guest_thread = GuestTask::new( state, @@ -1999,8 +1916,6 @@ impl StoreOpaque { callee, callee_async_typed, true, - #[cfg(feature = "task-group-hook")] - group, )?; Instance::from_wasmtime(self, callee.instance).add_guest_thread_to_instance_table( @@ -2081,13 +1996,6 @@ impl StoreOpaque { let caller = self.current_guest_thread()?; log::trace!("new deferred host task with caller {caller:?}"); - #[cfg(feature = "task-group-hook")] - { - let state = self.concurrent_state_mut()?; - let group = state.get_mut(caller.task)?.group; - state.increment_group_ref_count(group)?; - } - self.set_thread(CurrentThread::DeferredHost(caller))?; let state = self.concurrent_state_mut()?; debug_assert!(state.deferred_host_call_context.is_none()); @@ -2109,13 +2017,6 @@ impl StoreOpaque { ) -> Result<()> { match original_task { Some(caller) => { - #[cfg(feature = "task-group-hook")] - { - let state = self.concurrent_state_mut()?; - let group = state.get_mut(caller.task)?.group; - state.decrement_group_ref_count(group)?; - } - self.set_thread(caller)?; if materialized_task.is_none() { let state = self.concurrent_state_mut()?; @@ -2130,7 +2031,7 @@ impl StoreOpaque { "delete host task with caller {original_task:?} and materialized as {materialized_task:?}" ); if let Some(task) = materialized_task { - self.concurrent_state_mut()?.delete(task)?; + Waitable::Host(task).delete_from(self)?; } } None => { @@ -2153,7 +2054,7 @@ impl StoreOpaque { /// This will save off any state necessary for the previous thread, if /// applicable, and then it'll additionally update state for `thread` if /// needed too. - fn set_thread(&mut self, thread: impl Into) -> Result { + pub(crate) fn set_thread(&mut self, thread: impl Into) -> Result { let thread = thread.into(); let state = self.concurrent_state_mut()?; state.debug_assert_deferred_host_invariant(); @@ -2672,7 +2573,7 @@ impl StoreOpaque { .table .get_mut() .iter_mut() - .filter_map(|entry| { + .filter_map(|(_, entry)| { if let Some(task) = entry.downcast_ref::() { Some(task.instance) } else { @@ -3218,14 +3119,6 @@ impl Instance { let token = StoreToken::new(store.as_context_mut()); let old_thread = store.0.current_guest_thread()?; - #[cfg(feature = "task-group-hook")] - let group = { - let state = store.0.concurrent_state_mut()?; - let group = state.get_mut(old_thread.task)?.group; - state.increment_group_ref_count(group)?; - group - }; - let state = store.0.concurrent_state_mut()?; debug_assert_eq!( @@ -3339,8 +3232,6 @@ impl Instance { // We don't know whether the callee export was lifted sync or async // yet, but we'll update this in `start_call`: false, - #[cfg(feature = "task-group-hook")] - group, )?; // Make the new thread the current one so that `Self::start_call` knows @@ -3957,12 +3848,6 @@ impl Instance { } } - #[cfg(feature = "task-group-hook")] - { - let group = task.group; - concurrent_state.decrement_group_ref_count(group)?; - } - (Waitable::Host(id), true) } else { let id = TableId::::new(rep); @@ -5113,16 +4998,27 @@ enum HostTaskState { impl HostTask { fn new( + concurrent_state: &mut ConcurrentState, state: HostTaskState, - #[cfg(feature = "task-group-hook")] group: TableId, - ) -> Self { - Self { + caller: QualifiedThreadId, + ) -> Result { + #[cfg(feature = "task-group-hook")] + let group = { + let group = concurrent_state.get_mut(caller.task)?.group; + concurrent_state.increment_group_ref_count(group)?; + group + }; + + #[cfg(not(feature = "task-group-hook"))] + let _ = (concurrent_state, caller); + + Ok(Self { common: WaitableCommon::default(), call_context: CallContext::default(), state, #[cfg(feature = "task-group-hook")] group, - } + }) } } @@ -5453,7 +5349,6 @@ impl GuestTask { instance: RuntimeInstance, async_typed: bool, async_lifted: bool, - #[cfg(feature = "task-group-hook")] group: TableId, ) -> Result { let host_future_state = match &caller { Caller::Guest { .. } => HostFutureState::NotApplicable, @@ -5468,6 +5363,17 @@ impl GuestTask { } } }; + + #[cfg(feature = "task-group-hook")] + let group = match caller { + Caller::Guest { thread } => { + let group = state.get_mut(thread.task)?.group; + state.increment_group_ref_count(group)?; + group + } + Caller::Host { .. } => state.make_task_group()?, + }; + let task = state.push(Self { common: WaitableCommon::default(), lower_params: Some(lower_params), @@ -5675,7 +5581,14 @@ impl Waitable { match self { Self::Host(task) => { log::trace!("delete host task {task:?}"); - store.concurrent_state_mut()?.delete(*task)?; + let state = store.concurrent_state_mut()?; + let task = state.delete(*task)?; + + #[cfg(feature = "task-group-hook")] + state.decrement_group_ref_count(task.group)?; + + #[cfg(not(feature = "task-group-hook"))] + let _ = task; } Self::Guest(task) => { log::trace!("delete guest task {task:?}"); @@ -5802,17 +5715,6 @@ impl CurrentThread { fn is_none(&self) -> bool { matches!(self, Self::None) } - - #[cfg(feature = "task-group-hook")] - fn group(&self, state: &mut ConcurrentState) -> Result>> { - Ok(match self { - Self::Guest(thread) | Self::DeferredHost(thread) => { - Some(state.get_mut(thread.task)?.group) - } - Self::Host(task) => Some(state.get_mut(*task)?.group), - Self::None => None, - }) - } } impl From for CurrentThread { @@ -5981,7 +5883,7 @@ impl ConcurrentState { futures: &mut Vec>, ) { let mut items = Vec::new(); - for entry in self.table.get_mut().iter_mut() { + for (_, entry) in self.table.get_mut().iter_mut() { if let Some(set) = entry.downcast_mut::() { for mode in mem::take(&mut set.waiting).into_values() { match mode { @@ -6080,7 +5982,7 @@ impl ConcurrentState { task_group_hook: _, } = self; - for entry in table.get_mut().iter_mut() { + for (_, entry) in table.get_mut().iter_mut() { if let Some(set) = entry.downcast_mut::() { for mode in set.waiting.values_mut() { match mode { @@ -6329,18 +6231,9 @@ impl ConcurrentState { thread => return Ok(thread), }; - #[cfg(feature = "task-group-hook")] - let group = self.get_mut(caller.task)?.group; - - #[cfg(not(feature = "task-group-hook"))] - let _ = caller; - // Push first so allocation failure leaves the deferred state intact. - let task = self.push(HostTask::new( - HostTaskState::CalleeStarted, - #[cfg(feature = "task-group-hook")] - group, - ))?; + let task = HostTask::new(self, HostTaskState::CalleeStarted, caller)?; + let task = self.push(task)?; let call_context = self .deferred_host_call_context .take() @@ -6375,39 +6268,6 @@ impl ConcurrentState { _ => bail_bug!("current scope is not a deferred host scope"), } } - - #[cfg(feature = "task-group-hook")] - fn make_task_group(&mut self) -> Result> { - let group = self.push(TaskGroup { ref_count: 1 })?; - if let Some(hook) = &mut self.task_group_hook { - hook.handle_start(TaskGroupId(group))?; - } - log::trace!("new {group:?}"); - Ok(group) - } - - #[cfg(feature = "task-group-hook")] - fn increment_group_ref_count(&mut self, group: TableId) -> Result<()> { - let count = &mut self.get_mut(group)?.ref_count; - *count += 1; - log::trace!("increment {group:?} to {count}"); - Ok(()) - } - - #[cfg(feature = "task-group-hook")] - fn decrement_group_ref_count(&mut self, group: TableId) -> Result<()> { - let count = &mut self.get_mut(group)?.ref_count; - assert!(*count >= 1); - *count -= 1; - log::trace!("decrement {group:?} to {count}"); - if *count == 0 { - self.delete(group)?; - if let Some(hook) = &mut self.task_group_hook { - hook.handle_finish(TaskGroupId(group))?; - } - } - Ok(()) - } } /// Provide a type hint to compiler about the shape of a parameter lower @@ -6468,17 +6328,6 @@ enum WaitableCheck { Poll, } -/// An identifier representing a guest task within a component. -/// -/// This can be acquired by calling [`Func::start_call_concurrent`] or -/// [`TypedFunc::start_call_concurrent`] and then using the -/// [`FuncCallConcurrent::task`] accessor, for example. This can then be -/// reflected on with [`StoreContextMut::async_call_stack`]. -/// -/// [`TypedFunc::start_call_concurrent`]: crate::component::TypedFunc::start_call_concurrent -#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)] -pub struct GuestTaskId(TableId); - /// Represents a guest task called from the host, prepared using `prepare_call`. pub(crate) struct PreparedCall { /// The guest export to be called @@ -6578,9 +6427,6 @@ pub(crate) fn prepare_call( let caller = store.0.materialize_host_task_id()?; let state = store.0.concurrent_state_mut()?; - #[cfg(feature = "task-group-hook")] - let group = state.make_task_group()?; - let (tx, rx) = oneshot::channel(); let instance = handle.instance().runtime_instance(component_instance); @@ -6615,8 +6461,6 @@ pub(crate) fn prepare_call( instance, async_typed, async_lifted, - #[cfg(feature = "task-group-hook")] - group, )?; Ok(PreparedCall { @@ -6631,9 +6475,10 @@ pub(crate) fn prepare_call( pub(crate) struct StagedCall { store: StoreId, - task: TableId, rx: oneshot::Receiver, _marker: PhantomData R>, + #[cfg(feature = "task-group-hook")] + group: TableId, } impl StagedCall { @@ -6659,15 +6504,12 @@ impl StagedCall { Ok(StagedCall { store: store.0.id(), - task: thread.task, rx, _marker: PhantomData, + #[cfg(feature = "task-group-hook")] + group: store.0.concurrent_state_mut()?.get_mut(thread.task)?.group, }) } - - fn task(&self) -> GuestTaskId { - GuestTaskId(self.task) - } } impl Future for StagedCall diff --git a/crates/wasmtime/src/runtime/component/concurrent/func.rs b/crates/wasmtime/src/runtime/component/concurrent/func.rs index c1f823cd3d35..22a0fad41ca5 100644 --- a/crates/wasmtime/src/runtime/component/concurrent/func.rs +++ b/crates/wasmtime/src/runtime/component/concurrent/func.rs @@ -1,5 +1,5 @@ use crate::component::concurrent::TaskId; -use crate::component::concurrent::{self, GuestTaskId, PreparedCall}; +use crate::component::concurrent::{self, PreparedCall}; use crate::component::func::LowerContext; use crate::component::{AsAccessor, ComponentNamedList, Func, Lift, Lower, TypedFunc, Val}; use crate::prelude::*; @@ -13,7 +13,7 @@ use wasmtime_environ::component::{InterfaceType, MAX_FLAT_PARAMS, MAX_FLAT_RESUL /// Returned from [`Func::start_call_concurrent`] to represent a /// pending-but-not-yet-resolved call into wasm. pub struct FuncCallConcurrent<'a, T> { - call: concurrent::StagedCall>, + pub(crate) call: concurrent::StagedCall>, results: &'a mut [Val], _marker: marker::PhantomData, } @@ -133,6 +133,21 @@ impl Func { mut store: impl AsContextMut, params: &'a [Val], results: &'a mut [Val], + ) -> Result> { + let result = self.do_start_call_concurrent(&mut store, params, results); + + if result.is_err() { + store.as_context_mut().0.set_trapped(); + } + + result + } + + fn do_start_call_concurrent<'a, T: Send + 'static>( + self, + mut store: impl AsContextMut, + params: &'a [Val], + results: &'a mut [Val], ) -> Result> { self.check_params_results(store.as_context_mut(), params, results)?; let prepared = self.prepare_call_dynamic(store.as_context_mut(), params.to_vec())?; @@ -200,20 +215,10 @@ impl Func { } } -impl FuncCallConcurrent<'_, T> { - /// Returns the task that this invocation corresponds to. - /// - /// This can be later correlated with [`StoreContextMut::async_call_stack`] - /// for example. - pub fn task(&self) -> GuestTaskId { - self.call.task() - } -} - /// Returned from [`TypedFunc::start_call_concurrent`] to represent a /// pending-but-not-yet-resolved call into wasm. pub struct TypedFuncCallConcurrent { - call: concurrent::StagedCall, + pub(crate) call: concurrent::StagedCall, _marker: marker::PhantomData, } @@ -485,13 +490,3 @@ where ) } } - -impl TypedFuncCallConcurrent { - /// Returns the task that this invocation corresponds to. - /// - /// This can be later correlated with [`StoreContextMut::async_call_stack`] - /// for example. - pub fn task(&self) -> GuestTaskId { - self.call.task() - } -} diff --git a/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs new file mode 100644 index 000000000000..b9278c4b8e8b --- /dev/null +++ b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs @@ -0,0 +1,208 @@ +use crate::component::concurrent::table::{TableDebug, TableId}; +use crate::component::concurrent::{ConcurrentState, CurrentThread}; +use crate::component::{FuncCallConcurrent, TypedFuncCallConcurrent}; +use crate::error::Result; +use crate::store::StoreOpaque; +use crate::{AsContextMut as _, Store, StoreContextMut}; +use alloc::boxed::Box; +use alloc::vec::Vec; + +pub(super) struct TaskGroup { + ref_count: usize, +} + +impl TableDebug for TaskGroup { + fn type_name() -> &'static str { + "TaskGroup" + } +} + +/// Represents a "task group" containing the "root" task of a host->guest call, +/// plus any subtasks transitively created by that task. +/// +/// See [TaskGroupHook] for details. +#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)] +pub struct TaskGroupId(pub(super) TableId); + +/// Trait for being notified by the runtime of activity concerning a "task group". +/// +/// The Component Model specification has no notion of a "task group"[^1], but +/// we define one here in order to enable embedders to associate guest->host +/// calls with corresponding host->guest calls in a predictable way. +/// +/// A `TaskGroupId` is allocated whenever a guest task is created for a +/// host->guest call, and `handle_start` is called. That task is considered the +/// "root" task for the task group, and any subtasks transitively created by it +/// will also be considered part of that task group. Whenever the runtime +/// switches (Component-Model-level) threads, it will call `handle_exit` for the +/// group to which the old thread belonged, if any, and call `handle_enter` for +/// the group to which the new thread belongs. Only once all the threads of all +/// those tasks have exited (and the guest has dropped any subtask handles +/// referring to any of those tasks) will the `TaskGroupId` be deallocated, at +/// which point `handle_finish` will be called. +/// +/// Note that a given `TaskGroupId` may be reused after `handle_finish` is +/// called, so implementations of this trait must take care to reset any state +/// associated with it. +/// +/// Each of these functions may return an error, in which case any running guest +/// code will trap, the store will be poisoned such that it cannot be used to +/// run any further guest code, and the error will propagate back to the +/// host->guest caller. Furthermore, if and when the store is poisoned due to +/// any unrecoverable error (whether it was produced by the hook, the guest, or +/// the host), the currently-entered group will be exited (i.e. `handle_exit` +/// called), if any, and any started group will be finished +/// (i.e. `handle_finish` called). +/// +/// As of this writing, https://github.com/WebAssembly/component-model/pull/730 +/// (which adds `thread.set-task` and related intrinsics) has not yet been +/// merged. Once it has, and Wasmtime adds support for that feature, it will be +/// possible for guest threads to change their task; in that case, the thread +/// will effectively join whatever group the new task belongs to, which might +/// not be the same as that of the old task. In addition, the new +/// `thread.get-task` intrinsic will give the guest another way (besides subtask +/// handles) to keep tasks alive beyond the point when all their threads have +/// exited or switched tasks, in which case the group it belongs to will not be +/// disposed until all such tasks have been dropped using `task.drop`. +/// +/// [^1]: Although it does _imply_ such a notion in the discussion of ["semantic +/// tail +/// calls"](https://github.com/WebAssembly/component-model/blob/d1daf829e2da2293091c105121383ffbc3b3515b/design/mvp/Concurrency.md?plain=1#L339-L3410. +pub trait TaskGroupHook: Send + Sync + 'static { + /// Handle notification that a new task group has been created (i.e. a + /// host->guest call has been prepared). + fn handle_start(&mut self, id: TaskGroupId) -> Result<()>; + /// Handle notification that the runtime has switched to a thread belonging to + /// the specified task group. + fn handle_enter(&mut self, id: TaskGroupId) -> Result<()>; + /// Handle notification that the runtime has switched away from a thread + /// belonging to the specified task group. + fn handle_exit(&mut self, id: TaskGroupId) -> Result<()>; + /// Handle notification that the specified task group has been disposed of + /// (i.e. the task created for the host->guest call for which the task group + /// was created has exited, along with any and all subtasks transitively + /// created by that task, and the guest has dropped any and all handles to + /// those tasks). + fn handle_finish(&mut self, id: TaskGroupId) -> Result<()>; +} + +impl FuncCallConcurrent<'_, T> { + /// Returns the task group that this invocation corresponds to. + /// + /// This can be later correlated with calls to [`TaskGroupHook`]. + pub fn group(&self) -> TaskGroupId { + TaskGroupId(self.call.group) + } +} + +impl TypedFuncCallConcurrent { + /// Returns the task group that this invocation corresponds to. + /// + /// This can be later correlated with calls to [`TaskGroupHook`]. + pub fn group(&self) -> TaskGroupId { + TaskGroupId(self.call.group) + } +} + +impl Store { + /// Convenience wrapper for [`StoreContextMut::task_group_hook`] + pub fn task_group_hook(&mut self, hook: impl TaskGroupHook) { + self.as_context_mut().task_group_hook(hook); + } +} + +impl StoreContextMut<'_, T> { + /// Set a [`TaskGroupHook`] for this store. + /// + /// This will overwrite any hook that was previously set. + pub fn task_group_hook(self, hook: impl TaskGroupHook) { + self.0 + .concurrent_state_mut_without_forcing_current_thread() + .task_group_hook = Some(Box::new(hook)); + } +} + +impl StoreOpaque { + pub(crate) fn clean_up_task_groups(&mut self) { + if !self.concurrency_support() { + return; + } + + // Note that we ignore all errors here since, if we've reached here, + // it's either because we've already poisoned the store and are in the + // process of propagating the error which caused it to be poisoned or + // because we're disposing of the store entirely. + + let state = self.concurrent_state_mut_without_forcing_current_thread(); + if let Some(mut hook) = state.task_group_hook.take() { + let thread = state.unforced_current_thread; + if let Ok(Some(group)) = thread.group(state) { + _ = hook.handle_exit(TaskGroupId(group)); + } + + let groups = state + .table + .get_mut() + .iter_mut() + .filter_map(|(rep, entry)| { + if entry.downcast_mut::().is_some() { + Some(TableId::new(rep)) + } else { + None + } + }) + .collect::>(); + + for group in groups { + state.delete(group).unwrap(); + _ = hook.handle_finish(TaskGroupId(group)); + } + + state.task_group_hook = Some(hook); + } + } +} + +impl CurrentThread { + pub(super) fn group(&self, state: &mut ConcurrentState) -> Result>> { + Ok(match self { + Self::Guest(thread) | Self::DeferredHost(thread) => { + Some(state.get_mut(thread.task)?.group) + } + Self::Host(task) => Some(state.get_mut(*task)?.group), + Self::None => None, + }) + } +} + +impl ConcurrentState { + pub(super) fn make_task_group(&mut self) -> Result> { + let group = self.push(TaskGroup { ref_count: 1 })?; + if let Some(hook) = &mut self.task_group_hook { + hook.handle_start(TaskGroupId(group))?; + } + log::trace!("new {group:?}"); + Ok(group) + } + + pub(super) fn increment_group_ref_count(&mut self, group: TableId) -> Result<()> { + let count = &mut self.get_mut(group)?.ref_count; + *count += 1; + log::trace!("increment {group:?} to {count}"); + Ok(()) + } + + pub(super) fn decrement_group_ref_count(&mut self, group: TableId) -> Result<()> { + let count = &mut self.get_mut(group)?.ref_count; + assert!(*count >= 1); + *count -= 1; + log::trace!("decrement {group:?} to {count}"); + if *count == 0 { + self.delete(group)?; + if let Some(hook) = &mut self.task_group_hook { + hook.handle_finish(TaskGroupId(group))?; + } + } + Ok(()) + } +} diff --git a/crates/wasmtime/src/runtime/component/mod.rs b/crates/wasmtime/src/runtime/component/mod.rs index 0190f8bab364..61d7895170b2 100644 --- a/crates/wasmtime/src/runtime/component/mod.rs +++ b/crates/wasmtime/src/runtime/component/mod.rs @@ -121,8 +121,8 @@ pub use self::component::{Component, ComponentExportIndex, ExportLookup}; pub use self::concurrent::{ Access, Accessor, AccessorTask, AsAccessor, Destination, DirectDestination, DirectSource, ErrorContext, FuncCallConcurrent, FutureAny, FutureConsumer, FutureProducer, FutureReader, - GuardedFutureReader, GuardedStreamReader, GuestTaskId, JoinHandle, ReadBuffer, Source, - StreamAny, StreamConsumer, StreamProducer, StreamReader, StreamResult, TypedFuncCallConcurrent, + GuardedFutureReader, GuardedStreamReader, JoinHandle, ReadBuffer, Source, StreamAny, + StreamConsumer, StreamProducer, StreamReader, StreamResult, TypedFuncCallConcurrent, VMComponentAsyncStore, VecBuffer, WriteBuffer, }; #[cfg(feature = "task-group-hook")] diff --git a/crates/wasmtime/src/runtime/component/resource_table.rs b/crates/wasmtime/src/runtime/component/resource_table.rs index 6aad48b3846c..4f2a351ff5ac 100644 --- a/crates/wasmtime/src/runtime/component/resource_table.rs +++ b/crates/wasmtime/src/runtime/component/resource_table.rs @@ -406,11 +406,14 @@ impl ResourceTable { } /// Iterate over all the entries in this table. - pub fn iter_mut(&mut self) -> impl Iterator { - self.entries.iter_mut().filter_map(|entry| match entry { - Entry::Occupied { entry } => Some(&mut *entry.entry), - Entry::Free { .. } => None, - }) + pub fn iter_mut(&mut self) -> impl Iterator { + self.entries + .iter_mut() + .enumerate() + .filter_map(|(rep, entry)| match entry { + Entry::Occupied { entry } => Some((u32::try_from(rep).unwrap(), &mut *entry.entry)), + Entry::Free { .. } => None, + }) } } diff --git a/crates/wasmtime/src/runtime/component/store.rs b/crates/wasmtime/src/runtime/component/store.rs index 8006c4581fca..f7c7eab7a3df 100644 --- a/crates/wasmtime/src/runtime/component/store.rs +++ b/crates/wasmtime/src/runtime/component/store.rs @@ -122,6 +122,9 @@ impl ComponentStoreData { // be called from with in the context of a `tls::set` closure. #[cfg(feature = "component-model-async")] if store.0.component_data().task_state.is_concurrent() { + #[cfg(feature = "task-group-hook")] + store.0.clean_up_task_groups(); + ComponentStoreData::drop_fibers_and_futures(store.0); } #[cfg(not(feature = "component-model-async"))] @@ -288,6 +291,8 @@ impl StoreOpaque { pub(crate) fn set_trapped(&mut self) { self.store_data_mut().components.trapped = true; + #[cfg(feature = "task-group-hook")] + self.clean_up_task_groups(); } /// Determine whether an instance may be entered from the host. diff --git a/src/commands/serve.rs b/src/commands/serve.rs index f4e88d33fbd6..4bde5246314d 100644 --- a/src/commands/serve.rs +++ b/src/commands/serve.rs @@ -21,7 +21,9 @@ use std::{ use tokio::io::{self, AsyncRead, AsyncWrite}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::Semaphore; -use wasmtime::component::{Component, GuestTaskId, Linker}; +#[cfg(feature = "task-group-hook")] +use wasmtime::component::TaskGroupId; +use wasmtime::component::{Component, Linker}; use wasmtime::error::Context as _; use wasmtime::{ AsContextMut as _, Engine, Result, Store, StoreContextMut, StoreLimits, UpdateDeadline, bail, @@ -907,7 +909,7 @@ impl WorkerState for HostWorkerState { &self, _store: StoreContextMut, request_id: u64, - _task_id: GuestTaskId, + #[cfg(feature = "task-group-hook")] _task_group: TaskGroupId, ) -> Pin + 'static + Send + Sync>> { log::info!( "Instance {} handling request {request_id}", diff --git a/tests/all/component_model/async.rs b/tests/all/component_model/async.rs index 0d80273aca88..cf2e38a198af 100644 --- a/tests/all/component_model/async.rs +++ b/tests/all/component_model/async.rs @@ -977,17 +977,49 @@ async fn bytes_stream_producer() -> Result<()> { Ok(()) } +#[derive(Copy, Clone)] +enum TaskGroupHookTestVariant { + WithDrop, + WithoutDrop, + Trap, +} + +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn task_group_hook_with_drop() -> Result<()> { + test_task_group_hook(TaskGroupHookTestVariant::WithDrop).await +} + +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn task_group_hook_without_drop() -> Result<()> { + test_task_group_hook(TaskGroupHookTestVariant::WithoutDrop).await +} + #[tokio::test] #[cfg_attr(miri, ignore)] -async fn task_group_hook() -> Result<()> { +async fn task_group_hook_trap() -> Result<()> { + test_task_group_hook(TaskGroupHookTestVariant::Trap).await +} + +/// Verify that all the expected `TaskGroupHook` events are delivered, +/// regardless of whether the guest tasks exit normally, leak subtasks, or trap. +async fn test_task_group_hook(variant: TaskGroupHookTestVariant) -> Result<()> { _ = env_logger::try_init(); let mut config = Config::new(); config.wasm_component_model_async(true); let engine = Engine::new(&config)?; + let call_drop = if let TaskGroupHookTestVariant::WithDrop = variant { + "(call $subtask.drop (local.get $waitable))" + } else { + "" + }; + let component = Component::new( &engine, - r#" + format!( + r#" (component (import "a" (func $a async)) (core func $a (canon lower (func $a) async)) @@ -1012,13 +1044,17 @@ async fn task_group_hook() -> Result<()> { (i32.or (i32.const 2 (; WAIT ;)) (i32.shl (local.get $set) (i32.const 4))) ) - (func (export "a-callback") (param $event_code i32) (param $waitable i32) (param $payload i32) (result i32) + (func (export "a-callback") + (param $event_code i32) + (param $waitable i32) + (param $payload i32) + (result i32) (if (i32.ne (local.get $event_code) (i32.const 1 (; SUBTASK ;))) (then unreachable)) (if (i32.ne (local.get $payload) (i32.const 2 (; RETURNED ;))) (then unreachable)) (call $waitable.join (local.get $waitable) (i32.const 0)) - (call $subtask.drop (local.get $waitable)) + {call_drop} (call $task.return) (i32.const 0 (; EXIT ;)) ) @@ -1038,7 +1074,8 @@ async fn task_group_hook() -> Result<()> { )) (func (export "a") async (canon lift (core func $a "a") async (callback (core func $a "a-callback")))) ) - "#, + "# + ), )?; #[derive(Copy, Clone, PartialEq, Eq, Debug)] @@ -1054,6 +1091,7 @@ async fn task_group_hook() -> Result<()> { next_id: usize, ids: HashMap, entered: Option, + call_count: usize, } #[derive(Clone)] @@ -1097,8 +1135,10 @@ async fn task_group_hook() -> Result<()> { } } - let hook = Hook(Arc::new(Mutex::new(HookInner::default()))); let yield_count = 5; + let call_count = 3; + + let hook = Hook(Arc::new(Mutex::new(HookInner::default()))); let mut linker = Linker::new(&engine); linker.root().func_wrap_concurrent("a", { @@ -1107,7 +1147,7 @@ async fn task_group_hook() -> Result<()> { let hook = hook.clone(); Box::pin(async move { let mut previous_id = None; - let mut on_poll = move || { + let mut on_poll = || { let mut inner = hook.0.try_lock().unwrap(); // We should only ever be polled after a `handle_enter` and // before a `handle_exit` event: @@ -1130,30 +1170,49 @@ async fn task_group_hook() -> Result<()> { on_poll(); } - Ok(()) + let mut inner = hook.0.try_lock().unwrap(); + inner.call_count += 1; + if let TaskGroupHookTestVariant::Trap = variant + && inner.call_count == call_count + { + Err(wasmtime::format_err!("yikes")) + } else { + Ok(()) + } }) } })?; - let mut store = Store::new(&engine, ()); - let instance = linker.instantiate_async(&mut store, &component).await?; - let func = instance.get_typed_func::<(), ()>(&mut store, "a")?; - - store.task_group_hook(hook.clone()); - - let call_count = 3; - - // Call the guest export `call_count` times concurrently, recording the - // events in `hook`: - store - .run_concurrent(async |store| { - iter::repeat_with(|| func.call_concurrent(store, ())) - .take(call_count) - .collect::>() - .try_collect::<()>() - .await - }) - .await??; + { + let mut store = Store::new(&engine, ()); + let instance = linker.instantiate_async(&mut store, &component).await?; + let func = instance.get_typed_func::<(), ()>(&mut store, "a")?; + + store.task_group_hook(hook.clone()); + + // Call the guest export `call_count` times concurrently, recording the + // events in `hook`: + let result = store + .run_concurrent(async |store| { + iter::repeat_with(|| func.call_concurrent(store, ())) + .take(call_count) + .collect::>() + .try_collect::<()>() + .await + }) + .await; + + match (result, variant) { + ( + Ok(result), + TaskGroupHookTestVariant::WithDrop | TaskGroupHookTestVariant::WithoutDrop, + ) => result?, + (Err(error), TaskGroupHookTestVariant::Trap) => { + assert!(error.to_string().contains("yikes"), "{error}"); + } + _ => unreachable!(), + } + } // Group the events by task group, preserving order: let mut map = HashMap::<_, Vec<_>>::new(); From 967165799df35168e86a6f76033e12185e523ba3 Mon Sep 17 00:00:00 2001 From: Joel Dice Date: Tue, 29 Sep 2026 16:03:20 -0600 Subject: [PATCH 3/5] address review feedback --- Cargo.toml | 1 - crates/wasi-http/Cargo.toml | 1 - crates/wasi-http/src/handler.rs | 8 +-- crates/wasmtime/Cargo.toml | 3 + .../src/runtime/component/concurrent.rs | 54 ++++------------- .../src/runtime/component/concurrent/func.rs | 21 ++++++- .../component/concurrent/task_group_hook.rs | 60 +++++++++---------- .../concurrent/task_group_hook_disabled.rs | 34 +++++++++++ .../runtime/component/concurrent_disabled.rs | 7 +++ crates/wasmtime/src/runtime/component/mod.rs | 5 +- .../wasmtime/src/runtime/component/store.rs | 1 - src/commands/serve.rs | 6 +- 12 files changed, 110 insertions(+), 91 deletions(-) create mode 100644 crates/wasmtime/src/runtime/component/concurrent/task_group_hook_disabled.rs diff --git a/Cargo.toml b/Cargo.toml index be252f11aa58..7f234b13c2ed 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -556,7 +556,6 @@ component-model-async = [ "dep:futures", ] rr = ["wasmtime/rr", "component-model", "wasmtime-cli-flags/rr", "run"] -task-group-hook = ["wasmtime/task-group-hook", "wasmtime-wasi-http/task-group-hook"] # Enabled in release builds. We cannot enable the built-in gdbstub component in # the default build because the published-to-crates.io crates cannot build the diff --git a/crates/wasi-http/Cargo.toml b/crates/wasi-http/Cargo.toml index b031693db08e..664269d316f5 100644 --- a/crates/wasi-http/Cargo.toml +++ b/crates/wasi-http/Cargo.toml @@ -20,7 +20,6 @@ default-send-request = ["dep:tokio-rustls", "dep:rustls", "dep:webpki-roots"] p2 = ["wasmtime-wasi/p2"] p3 = ["wasmtime-wasi/p3", "dep:tokio-util"] component-model-async = ["futures/alloc", "wasmtime/component-model-async"] -task-group-hook = ["wasmtime/task-group-hook"] [dependencies] async-trait = { workspace = true } diff --git a/crates/wasi-http/src/handler.rs b/crates/wasi-http/src/handler.rs index e2696f500cb7..412c3bdbc727 100644 --- a/crates/wasi-http/src/handler.rs +++ b/crates/wasi-http/src/handler.rs @@ -30,9 +30,7 @@ use std::sync::{ use std::task::{Context, Poll}; use std::time::Instant; use tokio::sync::Notify; -#[cfg(feature = "task-group-hook")] -use wasmtime::component::TaskGroupId; -use wasmtime::component::{Accessor, Resource, TypedFuncCallConcurrent}; +use wasmtime::component::{Accessor, Resource, TaskGroupId, TypedFuncCallConcurrent}; #[cfg(feature = "p2")] use wasmtime::error::Context as _; use wasmtime::{AsContextMut, Result, Store, StoreContextMut, format_err}; @@ -220,7 +218,7 @@ pub trait WorkerState: 'static + Send + Sync { &self, store: StoreContextMut<'_, Self::StoreData>, data: Self::RequestData, - #[cfg(feature = "task-group-hook")] task_group: TaskGroupId, + task_group: TaskGroupId, ) -> Pin + 'static + Send + Sync>>; /// Dispose of the store belonging to the now-exited worker. @@ -511,7 +509,6 @@ where let expiration = dropper.state.on_request_start( store.as_context_mut(), request_data, - #[cfg(feature = "task-group-hook")] prepared.group(), ); Ok((prepared, expiration)) @@ -1047,7 +1044,6 @@ impl<'a, T: Send> Prepared<'a, T> { } } - #[cfg(feature = "task-group-hook")] fn group(&self) -> TaskGroupId { match self { #[cfg(feature = "p3")] diff --git a/crates/wasmtime/Cargo.toml b/crates/wasmtime/Cargo.toml index cbd69c44860e..d1636ec04d11 100644 --- a/crates/wasmtime/Cargo.toml +++ b/crates/wasmtime/Cargo.toml @@ -432,4 +432,7 @@ rr = [ "wasmtime-environ/rr" ] +# Enable support for the `Store::task_group_hook` function which registers a +# hook to be notified when host->guest task "groups" are created, entered, +# exited, and disposed of. task-group-hook = [ "component-model", "component-model-async" ] diff --git a/crates/wasmtime/src/runtime/component/concurrent.rs b/crates/wasmtime/src/runtime/component/concurrent.rs index 7f2285ccd709..eaaee6131616 100644 --- a/crates/wasmtime/src/runtime/component/concurrent.rs +++ b/crates/wasmtime/src/runtime/component/concurrent.rs @@ -86,8 +86,8 @@ use futures::channel::oneshot; use futures::stream::{FuturesUnordered, StreamExt}; use futures_and_streams::{FlatAbi, ReturnCode, TransmitHandle, TransmitIndex}; use table::{TableDebug, TableId}; -#[cfg(feature = "task-group-hook")] -use task_group_hook::TaskGroup; +#[cfg(not(feature = "task-group-hook"))] +use task_group_hook_disabled as task_group_hook; use wasmtime_environ::component::{ CanonicalAbiInfo, CanonicalOptions, CanonicalOptionsDataModel, MAX_FLAT_PARAMS, MAX_FLAT_RESULTS, OptionsIndex, PREPARE_ASYNC_NO_RESULT, PREPARE_ASYNC_WITH_RESULT, @@ -110,7 +110,8 @@ pub use futures_and_streams::{ }; pub(crate) use futures_and_streams::{ResourcePair, lower_error_context_to_index}; #[cfg(feature = "task-group-hook")] -pub use task_group_hook::{TaskGroupHook, TaskGroupId}; +pub use task_group_hook::TaskGroupHook; +pub use task_group_hook::TaskGroupId; mod abort; mod error_contexts; @@ -120,6 +121,8 @@ mod futures_and_streams; pub(crate) mod table; #[cfg(feature = "task-group-hook")] mod task_group_hook; +#[cfg(not(feature = "task-group-hook"))] +mod task_group_hook_disabled; pub(crate) mod tls; /// Constant defined in the Component Model spec to indicate that the async @@ -2060,20 +2063,7 @@ impl StoreOpaque { state.debug_assert_deferred_host_invariant(); let old_thread = mem::replace(&mut state.unforced_current_thread, thread); - #[cfg(feature = "task-group-hook")] - { - let old_group = old_thread.group(state)?; - let new_group = thread.group(state)?; - if let (true, Some(hook)) = ((old_group != new_group), &mut state.task_group_hook) { - if let Some(group) = old_group { - hook.handle_exit(TaskGroupId(group))?; - } - - if let Some(group) = new_group { - hook.handle_enter(TaskGroupId(group))?; - } - } - } + state.switch_threads(old_thread, thread)?; // First thing to do after swapping threads is updating the context // slots for this thread within the store. This restores the behavior of @@ -4969,8 +4959,7 @@ pub(crate) struct HostTask { state: HostTaskState, - #[cfg(feature = "task-group-hook")] - group: TableId, + group: TaskGroupId, } enum HostTaskState { @@ -5002,21 +4991,13 @@ impl HostTask { state: HostTaskState, caller: QualifiedThreadId, ) -> Result { - #[cfg(feature = "task-group-hook")] - let group = { - let group = concurrent_state.get_mut(caller.task)?.group; - concurrent_state.increment_group_ref_count(group)?; - group - }; - - #[cfg(not(feature = "task-group-hook"))] - let _ = (concurrent_state, caller); + let group = concurrent_state.get_mut(caller.task)?.group; + concurrent_state.increment_group_ref_count(group)?; Ok(Self { common: WaitableCommon::default(), call_context: CallContext::default(), state, - #[cfg(feature = "task-group-hook")] group, }) } @@ -5302,8 +5283,7 @@ pub(crate) struct GuestTask { decremented_interesting_task_count: bool, - #[cfg(feature = "task-group-hook")] - group: TableId, + group: TaskGroupId, } impl GuestTask { @@ -5364,7 +5344,6 @@ impl GuestTask { } }; - #[cfg(feature = "task-group-hook")] let group = match caller { Caller::Guest { thread } => { let group = state.get_mut(thread.task)?.group; @@ -5393,7 +5372,6 @@ impl GuestTask { async_typed, async_lifted, decremented_interesting_task_count: false, - #[cfg(feature = "task-group-hook")] group, })?; let new_thread = GuestThread::new_implicit(state, task)?; @@ -5584,18 +5562,13 @@ impl Waitable { let state = store.concurrent_state_mut()?; let task = state.delete(*task)?; - #[cfg(feature = "task-group-hook")] state.decrement_group_ref_count(task.group)?; - - #[cfg(not(feature = "task-group-hook"))] - let _ = task; } Self::Guest(task) => { log::trace!("delete guest task {task:?}"); let state = store.concurrent_state_mut()?; let task = state.delete(*task)?; - #[cfg(feature = "task-group-hook")] state.decrement_group_ref_count(task.group)?; // When a guest task is created it increments the @@ -5977,7 +5950,6 @@ impl ConcurrentState { interesting_tasks_empty_waker: _, ready_for_concurrent_call_waker: _, event_loop_running: _, - #[cfg(feature = "task-group-hook")] task_group_hook: _, } = self; @@ -6477,8 +6449,7 @@ pub(crate) struct StagedCall { store: StoreId, rx: oneshot::Receiver, _marker: PhantomData R>, - #[cfg(feature = "task-group-hook")] - group: TableId, + group: TaskGroupId, } impl StagedCall { @@ -6506,7 +6477,6 @@ impl StagedCall { store: store.0.id(), rx, _marker: PhantomData, - #[cfg(feature = "task-group-hook")] group: store.0.concurrent_state_mut()?.get_mut(thread.task)?.group, }) } diff --git a/crates/wasmtime/src/runtime/component/concurrent/func.rs b/crates/wasmtime/src/runtime/component/concurrent/func.rs index 22a0fad41ca5..280e0a972999 100644 --- a/crates/wasmtime/src/runtime/component/concurrent/func.rs +++ b/crates/wasmtime/src/runtime/component/concurrent/func.rs @@ -1,5 +1,4 @@ -use crate::component::concurrent::TaskId; -use crate::component::concurrent::{self, PreparedCall}; +use crate::component::concurrent::{self, PreparedCall, TaskGroupId, TaskId}; use crate::component::func::LowerContext; use crate::component::{AsAccessor, ComponentNamedList, Func, Lift, Lower, TypedFunc, Val}; use crate::prelude::*; @@ -490,3 +489,21 @@ where ) } } + +impl FuncCallConcurrent<'_, T> { + /// Returns the task group that this invocation corresponds to. + /// + /// This can be later correlated with calls to [`TaskGroupHook`]. + pub fn group(&self) -> TaskGroupId { + self.call.group + } +} + +impl TypedFuncCallConcurrent { + /// Returns the task group that this invocation corresponds to. + /// + /// This can be later correlated with calls to [`TaskGroupHook`]. + pub fn group(&self) -> TaskGroupId { + self.call.group + } +} diff --git a/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs index b9278c4b8e8b..8e2d08501faa 100644 --- a/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs +++ b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs @@ -1,13 +1,12 @@ use crate::component::concurrent::table::{TableDebug, TableId}; use crate::component::concurrent::{ConcurrentState, CurrentThread}; -use crate::component::{FuncCallConcurrent, TypedFuncCallConcurrent}; use crate::error::Result; use crate::store::StoreOpaque; use crate::{AsContextMut as _, Store, StoreContextMut}; use alloc::boxed::Box; use alloc::vec::Vec; -pub(super) struct TaskGroup { +struct TaskGroup { ref_count: usize, } @@ -22,7 +21,7 @@ impl TableDebug for TaskGroup { /// /// See [TaskGroupHook] for details. #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)] -pub struct TaskGroupId(pub(super) TableId); +pub struct TaskGroupId(TableId); /// Trait for being notified by the runtime of activity concerning a "task group". /// @@ -86,24 +85,6 @@ pub trait TaskGroupHook: Send + Sync + 'static { fn handle_finish(&mut self, id: TaskGroupId) -> Result<()>; } -impl FuncCallConcurrent<'_, T> { - /// Returns the task group that this invocation corresponds to. - /// - /// This can be later correlated with calls to [`TaskGroupHook`]. - pub fn group(&self) -> TaskGroupId { - TaskGroupId(self.call.group) - } -} - -impl TypedFuncCallConcurrent { - /// Returns the task group that this invocation corresponds to. - /// - /// This can be later correlated with calls to [`TaskGroupHook`]. - pub fn group(&self) -> TaskGroupId { - TaskGroupId(self.call.group) - } -} - impl Store { /// Convenience wrapper for [`StoreContextMut::task_group_hook`] pub fn task_group_hook(&mut self, hook: impl TaskGroupHook) { @@ -137,7 +118,7 @@ impl StoreOpaque { if let Some(mut hook) = state.task_group_hook.take() { let thread = state.unforced_current_thread; if let Ok(Some(group)) = thread.group(state) { - _ = hook.handle_exit(TaskGroupId(group)); + _ = hook.handle_exit(group); } let groups = state @@ -164,7 +145,7 @@ impl StoreOpaque { } impl CurrentThread { - pub(super) fn group(&self, state: &mut ConcurrentState) -> Result>> { + fn group(&self, state: &mut ConcurrentState) -> Result> { Ok(match self { Self::Guest(thread) | Self::DeferredHost(thread) => { Some(state.get_mut(thread.task)?.group) @@ -176,31 +157,46 @@ impl CurrentThread { } impl ConcurrentState { - pub(super) fn make_task_group(&mut self) -> Result> { - let group = self.push(TaskGroup { ref_count: 1 })?; + pub(super) fn switch_threads(&mut self, old: CurrentThread, new: CurrentThread) -> Result<()> { + let old_group = old.group(self)?; + let new_group = new.group(self)?; + if let (true, Some(hook)) = ((old_group != new_group), &mut self.task_group_hook) { + if let Some(group) = old_group { + hook.handle_exit(group)?; + } + + if let Some(group) = new_group { + hook.handle_enter(group)?; + } + } + Ok(()) + } + + pub(super) fn make_task_group(&mut self) -> Result { + let group = TaskGroupId(self.push(TaskGroup { ref_count: 1 })?); if let Some(hook) = &mut self.task_group_hook { - hook.handle_start(TaskGroupId(group))?; + hook.handle_start(group)?; } log::trace!("new {group:?}"); Ok(group) } - pub(super) fn increment_group_ref_count(&mut self, group: TableId) -> Result<()> { - let count = &mut self.get_mut(group)?.ref_count; + pub(super) fn increment_group_ref_count(&mut self, group: TaskGroupId) -> Result<()> { + let count = &mut self.get_mut(group.0)?.ref_count; *count += 1; log::trace!("increment {group:?} to {count}"); Ok(()) } - pub(super) fn decrement_group_ref_count(&mut self, group: TableId) -> Result<()> { - let count = &mut self.get_mut(group)?.ref_count; + pub(super) fn decrement_group_ref_count(&mut self, group: TaskGroupId) -> Result<()> { + let count = &mut self.get_mut(group.0)?.ref_count; assert!(*count >= 1); *count -= 1; log::trace!("decrement {group:?} to {count}"); if *count == 0 { - self.delete(group)?; + self.delete(group.0)?; if let Some(hook) = &mut self.task_group_hook { - hook.handle_finish(TaskGroupId(group))?; + hook.handle_finish(group)?; } } Ok(()) diff --git a/crates/wasmtime/src/runtime/component/concurrent/task_group_hook_disabled.rs b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook_disabled.rs new file mode 100644 index 000000000000..3264b79f5051 --- /dev/null +++ b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook_disabled.rs @@ -0,0 +1,34 @@ +use crate::component::concurrent::{ConcurrentState, CurrentThread}; +use crate::error::Result; +use crate::store::StoreOpaque; + +/// Represents a "task group" containing the "root" task of a host->guest call, +/// plus any subtasks transitively created by that task. +#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)] +pub struct TaskGroupId; + +impl StoreOpaque { + pub(crate) fn clean_up_task_groups(&mut self) {} +} + +impl ConcurrentState { + pub(super) fn switch_threads( + &mut self, + _old: CurrentThread, + _new: CurrentThread, + ) -> Result<()> { + Ok(()) + } + + pub(super) fn make_task_group(&mut self) -> Result { + Ok(TaskGroupId) + } + + pub(super) fn increment_group_ref_count(&mut self, _group: TaskGroupId) -> Result<()> { + Ok(()) + } + + pub(super) fn decrement_group_ref_count(&mut self, _group: TaskGroupId) -> Result<()> { + Ok(()) + } +} diff --git a/crates/wasmtime/src/runtime/component/concurrent_disabled.rs b/crates/wasmtime/src/runtime/component/concurrent_disabled.rs index 2e7f88562781..7ba9630203fd 100644 --- a/crates/wasmtime/src/runtime/component/concurrent_disabled.rs +++ b/crates/wasmtime/src/runtime/component/concurrent_disabled.rs @@ -8,6 +8,11 @@ use core::convert::Infallible; use core::mem::MaybeUninit; use wasmtime_environ::component::{CanonicalAbiInfo, InterfaceType}; +/// Represents a "task group" containing the "root" task of a host->guest call, +/// plus any subtasks transitively created by that task. +#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)] +pub struct TaskGroupId; + fn should_have_failed_validation(what: &str) -> Result { // This should be unreachable; if we trap here, it indicates a // bug in Wasmtime rather than in the guest. @@ -179,4 +184,6 @@ impl StoreOpaque { .current_scope_id_not_concurrent()? .map(|id| CurrentScope::Id(Scope::Id(id)))) } + + pub(crate) fn clean_up_task_groups(&mut self) {} } diff --git a/crates/wasmtime/src/runtime/component/mod.rs b/crates/wasmtime/src/runtime/component/mod.rs index 61d7895170b2..2467367f8cca 100644 --- a/crates/wasmtime/src/runtime/component/mod.rs +++ b/crates/wasmtime/src/runtime/component/mod.rs @@ -117,6 +117,9 @@ pub(crate) mod store; pub mod types; mod values; pub use self::component::{Component, ComponentExportIndex, ExportLookup}; +#[cfg(feature = "task-group-hook")] +pub use self::concurrent::TaskGroupHook; +pub use self::concurrent::TaskGroupId; #[cfg(feature = "component-model-async")] pub use self::concurrent::{ Access, Accessor, AccessorTask, AsAccessor, Destination, DirectDestination, DirectSource, @@ -125,8 +128,6 @@ pub use self::concurrent::{ StreamConsumer, StreamProducer, StreamReader, StreamResult, TypedFuncCallConcurrent, VMComponentAsyncStore, VecBuffer, WriteBuffer, }; -#[cfg(feature = "task-group-hook")] -pub use self::concurrent::{TaskGroupHook, TaskGroupId}; pub use self::func::{ ComponentNamedList, ComponentType, Func, Lift, Lower, TypedFunc, WasmList, WasmStr, }; diff --git a/crates/wasmtime/src/runtime/component/store.rs b/crates/wasmtime/src/runtime/component/store.rs index f7c7eab7a3df..4613f658ffb3 100644 --- a/crates/wasmtime/src/runtime/component/store.rs +++ b/crates/wasmtime/src/runtime/component/store.rs @@ -291,7 +291,6 @@ impl StoreOpaque { pub(crate) fn set_trapped(&mut self) { self.store_data_mut().components.trapped = true; - #[cfg(feature = "task-group-hook")] self.clean_up_task_groups(); } diff --git a/src/commands/serve.rs b/src/commands/serve.rs index 4bde5246314d..b67197b81898 100644 --- a/src/commands/serve.rs +++ b/src/commands/serve.rs @@ -21,9 +21,7 @@ use std::{ use tokio::io::{self, AsyncRead, AsyncWrite}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::Semaphore; -#[cfg(feature = "task-group-hook")] -use wasmtime::component::TaskGroupId; -use wasmtime::component::{Component, Linker}; +use wasmtime::component::{Component, Linker, TaskGroupId}; use wasmtime::error::Context as _; use wasmtime::{ AsContextMut as _, Engine, Result, Store, StoreContextMut, StoreLimits, UpdateDeadline, bail, @@ -909,7 +907,7 @@ impl WorkerState for HostWorkerState { &self, _store: StoreContextMut, request_id: u64, - #[cfg(feature = "task-group-hook")] _task_group: TaskGroupId, + _task_group: TaskGroupId, ) -> Pin + 'static + Send + Sync>> { log::info!( "Instance {} handling request {request_id}", From 305fe158ecaa54ef56527844184ca5ffa5a42637 Mon Sep 17 00:00:00 2001 From: Joel Dice Date: Wed, 30 Sep 2026 16:45:27 -0600 Subject: [PATCH 4/5] address review feedback --- crates/wasmtime/src/runtime/component/concurrent.rs | 6 +++--- .../src/runtime/component/concurrent/task_group_hook.rs | 6 +++++- .../component/concurrent/task_group_hook_disabled.rs | 8 +++++--- crates/wasmtime/src/runtime/component/store.rs | 1 - 4 files changed, 13 insertions(+), 8 deletions(-) diff --git a/crates/wasmtime/src/runtime/component/concurrent.rs b/crates/wasmtime/src/runtime/component/concurrent.rs index eaaee6131616..ff770e966571 100644 --- a/crates/wasmtime/src/runtime/component/concurrent.rs +++ b/crates/wasmtime/src/runtime/component/concurrent.rs @@ -86,8 +86,6 @@ use futures::channel::oneshot; use futures::stream::{FuturesUnordered, StreamExt}; use futures_and_streams::{FlatAbi, ReturnCode, TransmitHandle, TransmitIndex}; use table::{TableDebug, TableId}; -#[cfg(not(feature = "task-group-hook"))] -use task_group_hook_disabled as task_group_hook; use wasmtime_environ::component::{ CanonicalAbiInfo, CanonicalOptions, CanonicalOptionsDataModel, MAX_FLAT_PARAMS, MAX_FLAT_RESULTS, OptionsIndex, PREPARE_ASYNC_NO_RESULT, PREPARE_ASYNC_WITH_RESULT, @@ -123,6 +121,8 @@ pub(crate) mod table; mod task_group_hook; #[cfg(not(feature = "task-group-hook"))] mod task_group_hook_disabled; +#[cfg(not(feature = "task-group-hook"))] +use task_group_hook_disabled as task_group_hook; pub(crate) mod tls; /// Constant defined in the Component Model spec to indicate that the async @@ -2063,7 +2063,7 @@ impl StoreOpaque { state.debug_assert_deferred_host_invariant(); let old_thread = mem::replace(&mut state.unforced_current_thread, thread); - state.switch_threads(old_thread, thread)?; + state.handle_thread_switch(old_thread, thread)?; // First thing to do after swapping threads is updating the context // slots for this thread within the store. This restores the behavior of diff --git a/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs index 8e2d08501faa..828a4ec00fe8 100644 --- a/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs +++ b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs @@ -157,7 +157,11 @@ impl CurrentThread { } impl ConcurrentState { - pub(super) fn switch_threads(&mut self, old: CurrentThread, new: CurrentThread) -> Result<()> { + pub(super) fn handle_thread_switch( + &mut self, + old: CurrentThread, + new: CurrentThread, + ) -> Result<()> { let old_group = old.group(self)?; let new_group = new.group(self)?; if let (true, Some(hook)) = ((old_group != new_group), &mut self.task_group_hook) { diff --git a/crates/wasmtime/src/runtime/component/concurrent/task_group_hook_disabled.rs b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook_disabled.rs index 3264b79f5051..999ef28a4079 100644 --- a/crates/wasmtime/src/runtime/component/concurrent/task_group_hook_disabled.rs +++ b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook_disabled.rs @@ -5,14 +5,16 @@ use crate::store::StoreOpaque; /// Represents a "task group" containing the "root" task of a host->guest call, /// plus any subtasks transitively created by that task. #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)] -pub struct TaskGroupId; +pub struct TaskGroupId { + _priv: (), +} impl StoreOpaque { pub(crate) fn clean_up_task_groups(&mut self) {} } impl ConcurrentState { - pub(super) fn switch_threads( + pub(super) fn handle_thread_switch( &mut self, _old: CurrentThread, _new: CurrentThread, @@ -21,7 +23,7 @@ impl ConcurrentState { } pub(super) fn make_task_group(&mut self) -> Result { - Ok(TaskGroupId) + Ok(TaskGroupId { _priv: () }) } pub(super) fn increment_group_ref_count(&mut self, _group: TaskGroupId) -> Result<()> { diff --git a/crates/wasmtime/src/runtime/component/store.rs b/crates/wasmtime/src/runtime/component/store.rs index 4613f658ffb3..decb41133690 100644 --- a/crates/wasmtime/src/runtime/component/store.rs +++ b/crates/wasmtime/src/runtime/component/store.rs @@ -122,7 +122,6 @@ impl ComponentStoreData { // be called from with in the context of a `tls::set` closure. #[cfg(feature = "component-model-async")] if store.0.component_data().task_state.is_concurrent() { - #[cfg(feature = "task-group-hook")] store.0.clean_up_task_groups(); ComponentStoreData::drop_fibers_and_futures(store.0); From 98ebd10ec0acb36b22ad143e911e5f3b1f8e0e87 Mon Sep 17 00:00:00 2001 From: Joel Dice Date: Wed, 30 Sep 2026 18:33:30 -0600 Subject: [PATCH 5/5] fix docs issues --- .../src/runtime/component/concurrent/func.rs | 6 +++-- .../component/concurrent/task_group_hook.rs | 23 ++++++++++--------- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/crates/wasmtime/src/runtime/component/concurrent/func.rs b/crates/wasmtime/src/runtime/component/concurrent/func.rs index 280e0a972999..b0976868e2bd 100644 --- a/crates/wasmtime/src/runtime/component/concurrent/func.rs +++ b/crates/wasmtime/src/runtime/component/concurrent/func.rs @@ -493,7 +493,8 @@ where impl FuncCallConcurrent<'_, T> { /// Returns the task group that this invocation corresponds to. /// - /// This can be later correlated with calls to [`TaskGroupHook`]. + /// This can be later correlated with calls to `TaskGroupHook` if the + /// `task-group-hook` feature is enabled. pub fn group(&self) -> TaskGroupId { self.call.group } @@ -502,7 +503,8 @@ impl FuncCallConcurrent<'_, T> { impl TypedFuncCallConcurrent { /// Returns the task group that this invocation corresponds to. /// - /// This can be later correlated with calls to [`TaskGroupHook`]. + /// This can be later correlated with calls to `TaskGroupHook` if the + /// `task-group-hook` feature is enabled. pub fn group(&self) -> TaskGroupId { self.call.group } diff --git a/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs index 828a4ec00fe8..c941f64bc4d0 100644 --- a/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs +++ b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs @@ -53,20 +53,21 @@ pub struct TaskGroupId(TableId); /// called), if any, and any started group will be finished /// (i.e. `handle_finish` called). /// -/// As of this writing, https://github.com/WebAssembly/component-model/pull/730 -/// (which adds `thread.set-task` and related intrinsics) has not yet been -/// merged. Once it has, and Wasmtime adds support for that feature, it will be -/// possible for guest threads to change their task; in that case, the thread -/// will effectively join whatever group the new task belongs to, which might -/// not be the same as that of the old task. In addition, the new -/// `thread.get-task` intrinsic will give the guest another way (besides subtask -/// handles) to keep tasks alive beyond the point when all their threads have -/// exited or switched tasks, in which case the group it belongs to will not be -/// disposed until all such tasks have been dropped using `task.drop`. +/// As of this writing, +/// (which adds +/// `thread.set-task` and related intrinsics) has not yet been merged. Once it +/// has, and Wasmtime adds support for that feature, it will be possible for +/// guest threads to change their task; in that case, the thread will +/// effectively join whatever group the new task belongs to, which might not be +/// the same as that of the old task. In addition, the new `thread.get-task` +/// intrinsic will give the guest another way (besides subtask handles) to keep +/// tasks alive beyond the point when all their threads have exited or switched +/// tasks, in which case the group it belongs to will not be disposed until all +/// such tasks have been dropped using `task.drop`. /// /// [^1]: Although it does _imply_ such a notion in the discussion of ["semantic /// tail -/// calls"](https://github.com/WebAssembly/component-model/blob/d1daf829e2da2293091c105121383ffbc3b3515b/design/mvp/Concurrency.md?plain=1#L339-L3410. +/// calls"](https://github.com/WebAssembly/component-model/blob/d1daf829e2da2293091c105121383ffbc3b3515b/design/mvp/Concurrency.md?plain=1#L339-L3410). pub trait TaskGroupHook: Send + Sync + 'static { /// Handle notification that a new task group has been created (i.e. a /// host->guest call has been prepared).