wasmtime/runtime/component/concurrent/task_group_hook.rs
1use crate::bail_bug;
2use crate::component::concurrent::table::{TableDebug, TableId};
3use crate::component::concurrent::{ConcurrentState, CurrentThread};
4use crate::error::Result;
5use crate::store::StoreOpaque;
6use crate::{AsContextMut as _, Store, StoreContextMut};
7use alloc::boxed::Box;
8use alloc::vec::Vec;
9
10struct TaskGroup {
11 ref_count: usize,
12}
13
14impl TableDebug for TaskGroup {
15 fn type_name() -> &'static str {
16 "TaskGroup"
17 }
18}
19
20/// Represents a "task group" containing the "root" task of a host->guest call,
21/// plus any subtasks transitively created by that task.
22///
23/// See [TaskGroupHook] for details.
24#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)]
25pub struct TaskGroupId(TableId<TaskGroup>);
26
27/// Trait for being notified by the runtime of activity concerning a "task group".
28///
29/// The Component Model specification has no notion of a "task group"[^1], but
30/// we define one here in order to enable embedders to associate guest->host
31/// calls with corresponding host->guest calls in a predictable way.
32///
33/// A `TaskGroupId` is allocated whenever a guest task is created for a
34/// host->guest call, and `handle_start` is called. That task is considered the
35/// "root" task for the task group, and any subtasks transitively created by it
36/// will also be considered part of that task group. Whenever the runtime
37/// switches (Component-Model-level) threads, it will call `handle_exit` for the
38/// group to which the old thread belonged, if any, and call `handle_enter` for
39/// the group to which the new thread belongs. Only once all the threads of all
40/// those tasks have exited (and the guest has dropped any subtask handles
41/// referring to any of those tasks) will the `TaskGroupId` be deallocated, at
42/// which point `handle_finish` will be called.
43///
44/// Note that a given `TaskGroupId` may be reused after `handle_finish` is
45/// called, so implementations of this trait must take care to reset any state
46/// associated with it.
47///
48/// Each of these functions may return an error, in which case any running guest
49/// code will trap, the store will be poisoned such that it cannot be used to
50/// run any further guest code, and the error will propagate back to the
51/// host->guest caller. Furthermore, if and when the store is poisoned due to
52/// any unrecoverable error (whether it was produced by the hook, the guest, or
53/// the host), the currently-entered group will be exited (i.e. `handle_exit`
54/// called), if any, and any started group will be finished
55/// (i.e. `handle_finish` called).
56///
57/// As of this writing,
58/// <https://github.com/WebAssembly/component-model/pull/730> (which adds
59/// `thread.set-task` and related intrinsics) has not yet been merged. Once it
60/// has, and Wasmtime adds support for that feature, it will be possible for
61/// guest threads to change their task; in that case, the thread will
62/// effectively join whatever group the new task belongs to, which might not be
63/// the same as that of the old task. In addition, the new `thread.get-task`
64/// intrinsic will give the guest another way (besides subtask handles) to keep
65/// tasks alive beyond the point when all their threads have exited or switched
66/// tasks, in which case the group it belongs to will not be disposed until all
67/// such tasks have been dropped using `task.drop`.
68///
69/// [^1]: Although it does _imply_ such a notion in the discussion of ["semantic
70/// tail
71/// calls"](https://github.com/WebAssembly/component-model/blob/d1daf829e2da2293091c105121383ffbc3b3515b/design/mvp/Concurrency.md?plain=1#L339-L3410).
72pub trait TaskGroupHook: Send + Sync + 'static {
73 /// Handle notification that a new task group has been created (i.e. a
74 /// host->guest call has been prepared).
75 fn handle_start(&mut self, id: TaskGroupId) -> Result<()>;
76 /// Handle notification that the runtime has switched to a thread belonging to
77 /// the specified task group.
78 fn handle_enter(&mut self, id: TaskGroupId) -> Result<()>;
79 /// Handle notification that the runtime has switched away from a thread
80 /// belonging to the specified task group.
81 fn handle_exit(&mut self, id: TaskGroupId) -> Result<()>;
82 /// Handle notification that the specified task group has been disposed of
83 /// (i.e. the task created for the host->guest call for which the task group
84 /// was created has exited, along with any and all subtasks transitively
85 /// created by that task, and the guest has dropped any and all handles to
86 /// those tasks).
87 fn handle_finish(&mut self, id: TaskGroupId) -> Result<()>;
88}
89
90impl<T> Store<T> {
91 /// Convenience wrapper for [`StoreContextMut::task_group_hook`]
92 pub fn task_group_hook(&mut self, hook: impl TaskGroupHook) {
93 self.as_context_mut().task_group_hook(hook);
94 }
95}
96
97impl<T> StoreContextMut<'_, T> {
98 /// Set a [`TaskGroupHook`] for this store.
99 ///
100 /// This will overwrite any hook that was previously set.
101 pub fn task_group_hook(self, hook: impl TaskGroupHook) {
102 self.0
103 .concurrent_state_mut_without_forcing_current_thread()
104 .task_group_hook = Some(Box::new(hook));
105 }
106}
107
108impl StoreOpaque {
109 pub(crate) fn clean_up_task_groups(&mut self) {
110 if !self.concurrency_support() {
111 return;
112 }
113
114 // Note that we ignore all errors here since, if we've reached here,
115 // it's either because we've already poisoned the store and are in the
116 // process of propagating the error which caused it to be poisoned or
117 // because we're disposing of the store entirely.
118
119 let state = self.concurrent_state_mut_without_forcing_current_thread();
120 if let Some(mut hook) = state.task_group_hook.take() {
121 let thread = state.unforced_current_thread;
122 if let Ok(Some(group)) = thread.group(state) {
123 _ = hook.handle_exit(group);
124 }
125
126 let groups = state
127 .table
128 .get_mut()
129 .iter_mut()
130 .filter_map(|(rep, entry)| {
131 if entry.downcast_mut::<TaskGroup>().is_some() {
132 Some(TableId::new(rep))
133 } else {
134 None
135 }
136 })
137 .collect::<Vec<_>>();
138
139 for group in groups {
140 state.delete(group).unwrap();
141 _ = hook.handle_finish(TaskGroupId(group));
142 }
143
144 state.task_group_hook = Some(hook);
145 }
146 }
147}
148
149impl CurrentThread {
150 fn group(&self, state: &mut ConcurrentState) -> Result<Option<TaskGroupId>> {
151 Ok(match self {
152 Self::Guest(thread) | Self::DeferredHost(thread) => {
153 Some(state.get_mut(thread.task)?.group)
154 }
155 Self::Host(task) => Some(state.get_mut(*task)?.group),
156 Self::None => None,
157 })
158 }
159}
160
161impl ConcurrentState {
162 pub(super) fn handle_thread_switch(
163 &mut self,
164 old: CurrentThread,
165 new: CurrentThread,
166 ) -> Result<()> {
167 let old_group = old.group(self)?;
168 let new_group = new.group(self)?;
169 if let (true, Some(hook)) = ((old_group != new_group), &mut self.task_group_hook) {
170 if let Some(group) = old_group {
171 hook.handle_exit(group)?;
172 }
173
174 if let Some(group) = new_group {
175 hook.handle_enter(group)?;
176 }
177 }
178 Ok(())
179 }
180
181 pub(super) fn make_task_group(&mut self) -> Result<TaskGroupId> {
182 let group = TaskGroupId(self.push(TaskGroup { ref_count: 1 })?);
183 if let Some(hook) = &mut self.task_group_hook {
184 hook.handle_start(group)?;
185 }
186 log::trace!("new {group:?}");
187 Ok(group)
188 }
189
190 pub(super) fn increment_group_ref_count(&mut self, group: TaskGroupId) -> Result<()> {
191 let count = &mut self.get_mut(group.0)?.ref_count;
192 *count += 1;
193 log::trace!("increment {group:?} to {count}");
194 Ok(())
195 }
196
197 pub(super) fn decrement_group_ref_count(&mut self, group: TaskGroupId) -> Result<()> {
198 let count = &mut self.get_mut(group.0)?.ref_count;
199 if *count == 0 {
200 bail_bug!("can't decrement task group ref count that's already zero");
201 }
202 *count -= 1;
203 log::trace!("decrement {group:?} to {count}");
204 if *count == 0 {
205 self.delete(group.0)?;
206 if let Some(hook) = &mut self.task_group_hook {
207 hook.handle_finish(group)?;
208 }
209 }
210 Ok(())
211 }
212}