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/wasi-http/src/handler.rs b/crates/wasi-http/src/handler.rs index 937e2a41e96e..412c3bdbc727 100644 --- a/crates/wasi-http/src/handler.rs +++ b/crates/wasi-http/src/handler.rs @@ -30,7 +30,7 @@ use std::sync::{ use std::task::{Context, Poll}; use std::time::Instant; use tokio::sync::Notify; -use wasmtime::component::{Accessor, GuestTaskId, 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}; @@ -184,8 +184,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 +218,7 @@ pub trait WorkerState: 'static + Send + Sync { &self, store: StoreContextMut<'_, Self::StoreData>, data: Self::RequestData, - task: GuestTaskId, + task_group: TaskGroupId, ) -> Pin + 'static + Send + Sync>>; /// Dispose of the store belonging to the now-exited worker. @@ -508,7 +509,7 @@ where let expiration = dropper.state.on_request_start( store.as_context_mut(), request_data, - prepared.task(), + prepared.group(), ); Ok((prepared, expiration)) } @@ -948,7 +949,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 +1044,12 @@ impl<'a, T: Send> Prepared<'a, T> { } } - fn task(&self) -> GuestTaskId { + 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/Cargo.toml b/crates/wasmtime/Cargo.toml index 5bbbb645f20a..d1636ec04d11 100644 --- a/crates/wasmtime/Cargo.toml +++ b/crates/wasmtime/Cargo.toml @@ -431,3 +431,8 @@ rr = [ "component-model", "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 9223b620ccfc..ff770e966571 100644 --- a/crates/wasmtime/src/runtime/component/concurrent.rs +++ b/crates/wasmtime/src/runtime/component/concurrent.rs @@ -107,6 +107,9 @@ 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; +pub use task_group_hook::TaskGroupId; mod abort; mod error_contexts; @@ -114,6 +117,12 @@ mod func; mod future_stream_any; 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; +#[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 @@ -1297,11 +1306,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. @@ -1677,42 +1692,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 +1903,7 @@ impl StoreOpaque { caller: self.materialize_host_task_id()?, } }; + let state = self.concurrent_state_mut()?; let guest_thread = GuestTask::new( state, @@ -2018,6 +1998,7 @@ impl StoreOpaque { } let caller = self.current_guest_thread()?; log::trace!("new deferred host task with caller {caller:?}"); + self.set_thread(CurrentThread::DeferredHost(caller))?; let state = self.concurrent_state_mut()?; debug_assert!(state.deferred_host_call_context.is_none()); @@ -2053,7 +2034,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 => { @@ -2076,12 +2057,14 @@ 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(); let old_thread = mem::replace(&mut state.unforced_current_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 // `context.{get,set}`. This involves taking the old state out of the @@ -2440,12 +2423,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 +2455,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 +2527,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), @@ -2580,7 +2563,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 { @@ -3125,6 +3108,7 @@ impl Instance { let return_ = SendSyncPtr::new(return_); let token = StoreToken::new(store.as_context_mut()); let old_thread = store.0.current_guest_thread()?; + let state = store.0.concurrent_state_mut()?; debug_assert_eq!( @@ -3478,7 +3462,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 +3837,7 @@ impl Instance { bail_bug!("invalid state for callee in `subtask.drop`") } } + (Waitable::Host(id), true) } else { let id = TableId::::new(rep); @@ -3875,7 +3860,7 @@ impl Instance { } if delete { - waitable.delete_from(concurrent_state)?; + waitable.delete_from(store)?; } log::trace!("subtask_drop {waitable:?} (handle {task_id})"); @@ -4968,19 +4953,13 @@ 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, + + group: TaskGroupId, } enum HostTaskState { @@ -5007,13 +4986,20 @@ enum HostTaskState { } impl HostTask { - fn new(caller: TableId, state: HostTaskState) -> Self { - Self { + fn new( + concurrent_state: &mut ConcurrentState, + state: HostTaskState, + caller: QualifiedThreadId, + ) -> Result { + 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(), - caller, state, - } + group, + }) } } @@ -5296,6 +5282,8 @@ pub(crate) struct GuestTask { async_lifted: bool, decremented_interesting_task_count: bool, + + group: TaskGroupId, } impl GuestTask { @@ -5355,6 +5343,16 @@ impl GuestTask { } } }; + + 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), @@ -5374,6 +5372,7 @@ impl GuestTask { async_typed, async_lifted, decremented_interesting_task_count: false, + group, })?; let new_thread = GuestThread::new_implicit(state, task)?; let thread = state.push(new_thread)?; @@ -5555,17 +5554,23 @@ 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)?; + let state = store.concurrent_state_mut()?; + let task = state.delete(*task)?; + + state.decrement_group_ref_count(task.group)?; } Self::Guest(task) => { log::trace!("delete guest task {task:?}"); + let state = store.concurrent_state_mut()?; let task = state.delete(*task)?; + 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 +5580,7 @@ impl Waitable { debug_assert!(task.decremented_interesting_task_count); } Self::Transmit(task) => { - state.delete(*task)?; + store.concurrent_state_mut()?.delete(*task)?; } } @@ -5660,10 +5665,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 +5681,6 @@ impl CurrentThread { fn guest_task(&self) -> Option> { match self { Self::Guest(id) => Some(id.task), - Self::GuestTask(id) => Some(*id), _ => None, } } @@ -5802,6 +5802,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 +5827,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, } } } @@ -5850,7 +5856,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 { @@ -5944,9 +5950,11 @@ 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() { + for (_, entry) in table.get_mut().iter_mut() { if let Some(set) = entry.downcast_mut::() { for mode in set.waiting.values_mut() { match mode { @@ -6180,24 +6188,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(), @@ -6214,7 +6204,8 @@ impl ConcurrentState { }; // Push first so allocation failure leaves the deferred state intact. - let task = self.push(HostTask::new(caller.task, HostTaskState::CalleeStarted))?; + let task = HostTask::new(self, HostTaskState::CalleeStarted, caller)?; + let task = self.push(task)?; let call_context = self .deferred_host_call_context .take() @@ -6232,7 +6223,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(_) => { @@ -6309,17 +6300,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 @@ -6371,7 +6351,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(()) } @@ -6467,9 +6447,9 @@ pub(crate) fn prepare_call( pub(crate) struct StagedCall { store: StoreId, - task: TableId, rx: oneshot::Receiver, _marker: PhantomData R>, + group: TaskGroupId, } impl StagedCall { @@ -6495,15 +6475,11 @@ impl StagedCall { Ok(StagedCall { store: store.0.id(), - task: thread.task, rx, _marker: PhantomData, + 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..b0976868e2bd 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, GuestTaskId, 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::*; @@ -13,7 +12,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 +132,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 +214,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, } @@ -486,12 +490,22 @@ where } } +impl FuncCallConcurrent<'_, T> { + /// Returns the task group that this invocation corresponds to. + /// + /// 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 + } +} + impl TypedFuncCallConcurrent { - /// Returns the task that this invocation corresponds to. + /// Returns the task group 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() + /// 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 new file mode 100644 index 000000000000..c941f64bc4d0 --- /dev/null +++ b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook.rs @@ -0,0 +1,209 @@ +use crate::component::concurrent::table::{TableDebug, TableId}; +use crate::component::concurrent::{ConcurrentState, CurrentThread}; +use crate::error::Result; +use crate::store::StoreOpaque; +use crate::{AsContextMut as _, Store, StoreContextMut}; +use alloc::boxed::Box; +use alloc::vec::Vec; + +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(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, +/// (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 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(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 { + 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 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) { + 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(group)?; + } + log::trace!("new {group:?}"); + Ok(group) + } + + 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: 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.0)?; + if let Some(hook) = &mut self.task_group_hook { + 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..999ef28a4079 --- /dev/null +++ b/crates/wasmtime/src/runtime/component/concurrent/task_group_hook_disabled.rs @@ -0,0 +1,36 @@ +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 { + _priv: (), +} + +impl StoreOpaque { + pub(crate) fn clean_up_task_groups(&mut self) {} +} + +impl ConcurrentState { + pub(super) fn handle_thread_switch( + &mut self, + _old: CurrentThread, + _new: CurrentThread, + ) -> Result<()> { + Ok(()) + } + + pub(super) fn make_task_group(&mut self) -> Result { + Ok(TaskGroupId { _priv: () }) + } + + 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 7bf166b757ad..2467367f8cca 100644 --- a/crates/wasmtime/src/runtime/component/mod.rs +++ b/crates/wasmtime/src/runtime/component/mod.rs @@ -117,12 +117,15 @@ 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, 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, }; pub use self::func::{ 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..decb41133690 100644 --- a/crates/wasmtime/src/runtime/component/store.rs +++ b/crates/wasmtime/src/runtime/component/store.rs @@ -122,6 +122,8 @@ 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() { + store.0.clean_up_task_groups(); + ComponentStoreData::drop_fibers_and_futures(store.0); } #[cfg(not(feature = "component-model-async"))] @@ -288,6 +290,7 @@ impl StoreOpaque { pub(crate) fn set_trapped(&mut self) { self.store_data_mut().components.trapped = true; + 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..b67197b81898 100644 --- a/src/commands/serve.rs +++ b/src/commands/serve.rs @@ -21,7 +21,7 @@ use std::{ use tokio::io::{self, AsyncRead, AsyncWrite}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::Semaphore; -use wasmtime::component::{Component, GuestTaskId, Linker}; +use wasmtime::component::{Component, Linker, TaskGroupId}; use wasmtime::error::Context as _; use wasmtime::{ AsContextMut as _, Engine, Result, Store, StoreContextMut, StoreLimits, UpdateDeadline, bail, @@ -907,7 +907,7 @@ impl WorkerState for HostWorkerState { &self, _store: StoreContextMut, request_id: u64, - _task_id: GuestTaskId, + _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 43c3ab23a0cc..cf2e38a198af 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; @@ -971,196 +977,262 @@ 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 async_call_stack() -> Result<()> { +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_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)) - (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_drop} + (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, + call_count: usize, + } - 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 yield_count = 5; + let call_count = 3; - 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()))); - ;; 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 = || { + 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 + }; - (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(()) - }, - )?; + on_poll(); + for _ in 0..yield_count { + tokio::task::yield_now().await; + on_poll(); + } + + 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, None); + { + 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 call = func.start_call_concurrent(&mut store, ())?; - let root = call.task(); - store - .run_concurrent(async |store| func.finish_call_concurrent(store, call).await) - .await??; + 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(); + 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(()) }