Skip to main content

wasmtime/runtime/component/
concurrent.rs

1//! Runtime support for the Component Model Async ABI.
2//!
3//! This module and its submodules provide host runtime support for Component
4//! Model Async features such as async-lifted exports, async-lowered imports,
5//! streams, futures, and related intrinsics.  See [the Async
6//! Explainer](https://github.com/WebAssembly/component-model/blob/main/design/mvp/Concurrency.md)
7//! for a high-level overview.
8//!
9//! At the core of this support is an event loop which schedules and switches
10//! between guest tasks and any host tasks they create.  Each
11//! `Store` will have at most one event loop running at any given
12//! time, and that loop may be suspended and resumed by the host embedder using
13//! e.g. `StoreContextMut::run_concurrent`.  The `StoreContextMut::poll_until`
14//! function contains the loop itself, while the
15//! `StoreOpaque::concurrent_state` field holds its state.
16//!
17//! # Public API Overview
18//!
19//! ## Top-level API (e.g. kicking off host->guest calls and driving the event loop)
20//!
21//! - `[Typed]Func::call_concurrent`: Start a host->guest call to an
22//! async-lifted or sync-lifted import, creating a guest task.
23//!
24//! - `StoreContextMut::run_concurrent`: Run the event loop for the specified
25//! instance, allowing any and all tasks belonging to that instance to make
26//! progress.
27//!
28//! - `StoreContextMut::spawn`: Run a background task as part of the event loop
29//! for the specified instance.
30//!
31//! - `{Future,Stream}Reader::new`: Create a new Component Model `future` or
32//! `stream` which may be passed to the guest.  This takes a
33//! `{Future,Stream}Producer` implementation which will be polled for items when
34//! the consumer requests them.
35//!
36//! - `{Future,Stream}Reader::pipe`: Consume a `future` or `stream` by
37//! connecting it to a `{Future,Stream}Consumer` which will consume any items
38//! produced by the write end.
39//!
40//! ## Host Task API (e.g. implementing concurrent host functions and background tasks)
41//!
42//! - `LinkerInstance::func_wrap_concurrent`: Register a concurrent host
43//! function with the linker.  That function will take an `Accessor` as its
44//! first parameter, which provides access to the store between (but not across)
45//! await points.
46//!
47//! - `Accessor::with`: Access the store and its associated data.
48//!
49//! - `Accessor::spawn`: Run a background task as part of the event loop for the
50//! store.  This is equivalent to `StoreContextMut::spawn` but more convenient to use
51//! in host functions.
52
53use self::error_contexts::GlobalErrorContextRefCount;
54use crate::component::func::{Func, call_post_return};
55use crate::component::{
56    HasData, HasSelf, Instance, Resource, ResourceTable, ResourceTableError, RuntimeInstance,
57};
58use crate::fiber::{self, StoreFiber, StoreFiberYield};
59use crate::hash_set::HashSet;
60#[cfg(feature = "gc")]
61use crate::module::ModuleRegistry;
62use crate::prelude::*;
63use crate::store::{Store, StoreId, StoreInner, StoreOpaque, StoreToken};
64#[cfg(feature = "gc")]
65use crate::vm::GcRootsList;
66use crate::vm::component::{CallContext, ComponentInstance, CurrentScope, InstanceState, Scope};
67use crate::vm::{
68    AlwaysMut, SendSyncPtr, UncaughtException, VMFuncRef, VMLazyThread, VMMemoryDefinition, VMStore,
69};
70use crate::{AsContext, AsContextMut, Result, StoreContext, StoreContextMut, ValRaw, bail};
71use crate::{Instance as ModuleInstance, bail_bug};
72use alloc::borrow::ToOwned;
73use alloc::collections::{BTreeMap, BTreeSet, VecDeque};
74use core::any::Any;
75use core::cell::UnsafeCell;
76use core::fmt;
77use core::future;
78use core::future::Future;
79use core::marker::PhantomData;
80use core::mem::{self, ManuallyDrop, MaybeUninit};
81use core::ops::DerefMut;
82use core::pin::{Pin, pin};
83use core::ptr::{self, NonNull};
84use core::task::{Context, Poll, Waker};
85use futures::channel::oneshot;
86use futures::stream::{FuturesUnordered, StreamExt};
87use futures_and_streams::{FlatAbi, ReturnCode, TransmitHandle, TransmitIndex};
88use table::{TableDebug, TableId};
89use wasmtime_environ::component::{
90    CanonicalAbiInfo, CanonicalOptions, CanonicalOptionsDataModel, MAX_FLAT_PARAMS,
91    MAX_FLAT_RESULTS, OptionsIndex, PREPARE_ASYNC_NO_RESULT, PREPARE_ASYNC_WITH_RESULT,
92    RuntimeComponentInstanceIndex, RuntimeTableIndex, StringEncoding,
93    TypeComponentGlobalErrorContextTableIndex, TypeComponentLocalErrorContextTableIndex,
94    TypeFutureTableIndex, TypeStreamTableIndex, TypeTupleIndex,
95};
96use wasmtime_environ::packed_option::ReservedValue;
97use wasmtime_environ::{ModuleInternedTypeIndex, NUM_COMPONENT_CONTEXT_SLOTS, Trap};
98#[cfg(feature = "gc")]
99use wasmtime_unwinder::Unwind;
100
101pub use abort::JoinHandle;
102pub use func::{FuncCallConcurrent, TypedFuncCallConcurrent};
103pub use future_stream_any::{FutureAny, StreamAny};
104pub use futures_and_streams::{
105    Destination, DirectDestination, DirectSource, ErrorContext, FutureConsumer, FutureProducer,
106    FutureReader, GuardedFutureReader, GuardedStreamReader, ReadBuffer, Source, StreamConsumer,
107    StreamProducer, StreamReader, StreamResult, VecBuffer, WriteBuffer,
108};
109pub(crate) use futures_and_streams::{ResourcePair, lower_error_context_to_index};
110#[cfg(feature = "task-group-hook")]
111pub use task_group_hook::TaskGroupHook;
112pub use task_group_hook::TaskGroupId;
113
114mod abort;
115mod error_contexts;
116mod func;
117mod future_stream_any;
118mod futures_and_streams;
119pub(crate) mod table;
120#[cfg(feature = "task-group-hook")]
121mod task_group_hook;
122#[cfg(not(feature = "task-group-hook"))]
123mod task_group_hook_disabled;
124#[cfg(not(feature = "task-group-hook"))]
125use task_group_hook_disabled as task_group_hook;
126pub(crate) mod tls;
127
128/// Constant defined in the Component Model spec to indicate that the async
129/// intrinsic (e.g. `future.write`) has not yet completed.
130const BLOCKED: u32 = 0xffff_ffff;
131
132/// Corresponds to `CallState` in the upstream spec.
133#[derive(Clone, Copy, Eq, PartialEq, Debug)]
134pub enum Status {
135    Starting = 0,
136    Started = 1,
137    Returned = 2,
138    StartCancelled = 3,
139    ReturnCancelled = 4,
140}
141
142impl Status {
143    /// Packs this status and the optional `waitable` provided into a 32-bit
144    /// result that the canonical ABI requires.
145    ///
146    /// The low 4 bits are reserved for the status while the upper 28 bits are
147    /// the waitable, if present.
148    pub fn pack(self, waitable: Option<u32>) -> u32 {
149        assert!(matches!(self, Status::Returned) == waitable.is_none());
150        let waitable = waitable.unwrap_or(0);
151        assert!(waitable < (1 << 28));
152        (waitable << 4) | (self as u32)
153    }
154}
155
156/// Corresponds to `EventCode` in the Component Model spec, plus related payload
157/// data.
158#[derive(Clone, Copy, Debug)]
159enum Event {
160    None,
161    Subtask {
162        status: Status,
163    },
164    StreamRead {
165        code: ReturnCode,
166        pending: Option<(TypeStreamTableIndex, u32)>,
167    },
168    StreamWrite {
169        code: ReturnCode,
170        pending: Option<(TypeStreamTableIndex, u32)>,
171    },
172    FutureRead {
173        code: ReturnCode,
174        pending: Option<(TypeFutureTableIndex, u32)>,
175    },
176    FutureWrite {
177        code: ReturnCode,
178        pending: Option<(TypeFutureTableIndex, u32)>,
179    },
180    Cancelled,
181}
182
183impl Event {
184    /// Lower this event to core Wasm integers for delivery to the guest.
185    ///
186    /// Note that the waitable handle, if any, is assumed to be lowered
187    /// separately.
188    fn parts(self) -> (u32, u32) {
189        const EVENT_NONE: u32 = 0;
190        const EVENT_SUBTASK: u32 = 1;
191        const EVENT_STREAM_READ: u32 = 2;
192        const EVENT_STREAM_WRITE: u32 = 3;
193        const EVENT_FUTURE_READ: u32 = 4;
194        const EVENT_FUTURE_WRITE: u32 = 5;
195        const EVENT_CANCELLED: u32 = 6;
196        match self {
197            Event::None => (EVENT_NONE, 0),
198            Event::Cancelled => (EVENT_CANCELLED, 0),
199            Event::Subtask { status } => (EVENT_SUBTASK, status as u32),
200            Event::StreamRead { code, .. } => (EVENT_STREAM_READ, code.encode()),
201            Event::StreamWrite { code, .. } => (EVENT_STREAM_WRITE, code.encode()),
202            Event::FutureRead { code, .. } => (EVENT_FUTURE_READ, code.encode()),
203            Event::FutureWrite { code, .. } => (EVENT_FUTURE_WRITE, code.encode()),
204        }
205    }
206}
207
208/// Corresponds to `CallbackCode` in the spec.
209mod callback_code {
210    pub const EXIT: u32 = 0;
211    pub const YIELD: u32 = 1;
212    pub const WAIT: u32 = 2;
213}
214
215// Flags passed to the `start-call` intrinsic.
216const START_FLAG_ASYNC_CALLEE: u32 = wasmtime_environ::component::START_FLAG_ASYNC_CALLEE as u32;
217const START_FLAG_ASYNC_CALLER: u32 = wasmtime_environ::component::START_FLAG_ASYNC_CALLER as u32;
218
219/// Provides access to either store data (via the `get` method) or the store
220/// itself (via [`AsContext`]/[`AsContextMut`]), as well as the component
221/// instance to which the current host task belongs.
222///
223/// See [`Accessor::with`] for details.
224pub struct Access<'a, T: 'static, D: HasData + ?Sized = HasSelf<T>> {
225    store: StoreContextMut<'a, T>,
226    get_data: fn(&mut T) -> D::Data<'_>,
227}
228
229impl<'a, T, D> Access<'a, T, D>
230where
231    D: HasData + ?Sized,
232    T: 'static,
233{
234    /// Creates a new [`Access`] from its component parts.
235    pub fn new(store: StoreContextMut<'a, T>, get_data: fn(&mut T) -> D::Data<'_>) -> Self {
236        Self { store, get_data }
237    }
238
239    /// Get mutable access to the store data.
240    pub fn data_mut(&mut self) -> &mut T {
241        self.store.data_mut()
242    }
243
244    /// Get mutable access to the store data.
245    pub fn get(&mut self) -> D::Data<'_> {
246        (self.get_data)(self.data_mut())
247    }
248
249    /// Spawn a background task.
250    ///
251    /// See [`Accessor::spawn`] for details.
252    pub fn spawn(&mut self, task: impl for<'fut> AccessorTask<'fut, T, D>) -> Result<JoinHandle>
253    where
254        T: 'static,
255    {
256        let accessor = Accessor {
257            get_data: self.get_data,
258            token: StoreToken::new(self.store.as_context_mut()),
259        };
260        self.store
261            .as_context_mut()
262            .spawn_with_accessor(accessor, task)
263    }
264
265    /// Returns the getter this accessor is using to project from `T` into
266    /// `D::Data`.
267    pub fn getter(&self) -> fn(&mut T) -> D::Data<'_> {
268        self.get_data
269    }
270}
271
272impl<'a, T, D> AsContext for Access<'a, T, D>
273where
274    D: HasData + ?Sized,
275    T: 'static,
276{
277    type Data = T;
278
279    fn as_context(&self) -> StoreContext<'_, T> {
280        self.store.as_context()
281    }
282}
283
284impl<'a, T, D> AsContextMut for Access<'a, T, D>
285where
286    D: HasData + ?Sized,
287    T: 'static,
288{
289    fn as_context_mut(&mut self) -> StoreContextMut<'_, T> {
290        self.store.as_context_mut()
291    }
292}
293
294/// Provides scoped mutable access to store data in the context of a concurrent
295/// host task future.
296///
297/// This allows multiple host task futures to execute concurrently and access
298/// the store between (but not across) `await` points.
299///
300/// # Rationale
301///
302/// This structure is sort of like `&mut T` plus a projection from `&mut T` to
303/// `D::Data<'_>`. The problem this is solving, however, is that it does not
304/// literally store these values. The basic problem is that when a concurrent
305/// host future is being polled it has access to `&mut T` (and the whole
306/// `Store`) but when it's not being polled it does not have access to these
307/// values. This reflects how the store is only ever polling one future at a
308/// time so the store is effectively being passed between futures.
309///
310/// Rust's `Future` trait, however, has no means of passing a `Store`
311/// temporarily between futures. The [`Context`](core::task::Context) type does
312/// not have the ability to attach arbitrary information to it at this time.
313/// This type, [`Accessor`], is used to bridge this expressivity gap.
314///
315/// The [`Accessor`] type here represents the ability to acquire, temporarily in
316/// a synchronous manner, the current store. The [`Accessor::with`] function
317/// yields an [`Access`] which can be used to access [`StoreContextMut`], `&mut
318/// T`, or `D::Data<'_>`. Note though that [`Accessor::with`] intentionally does
319/// not take an `async` closure as its argument, instead it's a synchronous
320/// closure which must complete during on run of `Future::poll`. This reflects
321/// how the store is temporarily made available while a host future is being
322/// polled.
323///
324/// # Implementation
325///
326/// This type does not actually store `&mut T` nor `StoreContextMut<T>`, and
327/// this type additionally doesn't even have a lifetime parameter. This is
328/// instead a representation of proof of the ability to acquire these while a
329/// future is being polled. Wasmtime will, when it polls a host future,
330/// configure ambient state such that the `Accessor` that a future closes over
331/// will work and be able to access the store.
332///
333/// This has a number of implications for users such as:
334///
335/// * It's intentional that `Accessor` cannot be cloned, it needs to stay within
336///   the lifetime of a single future.
337/// * A future is expected to, however, close over an `Accessor` and keep it
338///   alive probably for the duration of the entire future.
339/// * Different host futures will be given different `Accessor`s, and that's
340///   intentional.
341/// * The `Accessor` type is `Send` and `Sync` irrespective of `T` which
342///   alleviates some otherwise required bounds to be written down.
343///
344/// # Using `Accessor` in `Drop`
345///
346/// The methods on `Accessor` are only expected to work in the context of
347/// `Future::poll` and are not guaranteed to work in `Drop`. This is because a
348/// host future can be dropped at any time throughout the system and Wasmtime
349/// store context is not necessarily available at that time. It's recommended to
350/// not use `Accessor` methods in anything connected to a `Drop` implementation
351/// as they will panic and have unintended results. If you run into this though
352/// feel free to file an issue on the Wasmtime repository.
353pub struct Accessor<T: 'static, D = HasSelf<T>>
354where
355    D: HasData + ?Sized,
356{
357    token: StoreToken<T>,
358    get_data: fn(&mut T) -> D::Data<'_>,
359}
360
361/// A helper trait to take any type of accessor-with-data in functions.
362///
363/// This trait is similar to [`AsContextMut`] except that it's used when
364/// working with an [`Accessor`] instead of a [`StoreContextMut`]. The
365/// [`Accessor`] is the main type used in concurrent settings and is passed to
366/// functions such as [`Func::call_concurrent`].
367///
368/// This trait is implemented for [`Accessor`] and `&T` where `T` implements
369/// this trait. This effectively means that regardless of the `D` in
370/// `Accessor<T, D>` it can still be passed to a function which just needs a
371/// store accessor.
372///
373/// Acquiring an [`Accessor`] can be done through
374/// [`StoreContextMut::run_concurrent`] for example or in a host function
375/// through
376/// [`Linker::func_wrap_concurrent`](crate::component::LinkerInstance::func_wrap_concurrent).
377pub trait AsAccessor {
378    /// The `T` in `Store<T>` that this accessor refers to.
379    type Data: 'static;
380
381    /// The `D` in `Accessor<T, D>`, or the projection out of
382    /// `Self::Data`.
383    type AccessorData: HasData + ?Sized;
384
385    /// Returns the accessor that this is referring to.
386    fn as_accessor(&self) -> &Accessor<Self::Data, Self::AccessorData>;
387}
388
389impl<T: AsAccessor + ?Sized> AsAccessor for &T {
390    type Data = T::Data;
391    type AccessorData = T::AccessorData;
392
393    fn as_accessor(&self) -> &Accessor<Self::Data, Self::AccessorData> {
394        T::as_accessor(self)
395    }
396}
397
398impl<T, D: HasData + ?Sized> AsAccessor for Accessor<T, D> {
399    type Data = T;
400    type AccessorData = D;
401
402    fn as_accessor(&self) -> &Accessor<T, D> {
403        self
404    }
405}
406
407// Note that it is intentional at this time that `Accessor` does not actually
408// store `&mut T` or anything similar. This distinctly enables the `Accessor`
409// structure to be both `Send` and `Sync` regardless of what `T` is (or `D` for
410// that matter). This is used to ergonomically simplify bindings where the
411// majority of the time `Accessor` is closed over in a future which then needs
412// to be `Send` and `Sync`. To avoid needing to write `T: Send` everywhere (as
413// you already have to write `T: 'static`...) it helps to avoid this.
414//
415// Note as well that `Accessor` doesn't actually store its data at all. Instead
416// it's more of a "proof" of what can be accessed from TLS. API design around
417// `Accessor` and functions like `Linker::func_wrap_concurrent` are
418// intentionally made to ensure that `Accessor` is ideally only used in the
419// context that TLS variables are actually set. For example host functions are
420// given `&Accessor`, not `Accessor`, and this prevents them from persisting
421// the value outside of a future. Within the future the TLS variables are all
422// guaranteed to be set while the future is being polled.
423//
424// Finally though this is not an ironclad guarantee, but nor does it need to be.
425// The TLS APIs are designed to panic or otherwise model usage where they're
426// called recursively or similar. It's hoped that code cannot be constructed to
427// actually hit this at runtime but this is not a safety requirement at this
428// time.
429const _: () = {
430    const fn assert<T: Send + Sync>() {}
431    assert::<Accessor<UnsafeCell<u32>>>();
432};
433
434impl<T> Accessor<T> {
435    /// Creates a new `Accessor` backed by the specified functions.
436    ///
437    /// - `get`: used to retrieve the store
438    ///
439    /// - `get_data`: used to "project" from the store's associated data to
440    /// another type (e.g. a field of that data or a wrapper around it).
441    ///
442    /// - `spawn`: used to queue spawned background tasks to be run later
443    pub(crate) fn new(token: StoreToken<T>) -> Self {
444        Self {
445            token,
446            get_data: |x| x,
447        }
448    }
449}
450
451impl<T, D> Accessor<T, D>
452where
453    D: HasData + ?Sized,
454{
455    /// Run the specified closure, passing it mutable access to the store.
456    ///
457    /// This function is one of the main building blocks of the [`Accessor`]
458    /// type. This yields synchronous, blocking, access to the store via an
459    /// [`Access`]. The [`Access`] implements [`AsContextMut`] in addition to
460    /// providing the ability to access `D` via [`Access::get`]. Note that the
461    /// `fun` here is given only temporary access to the store and `T`/`D`
462    /// meaning that the return value `R` here is not allowed to capture borrows
463    /// into the two. If access is needed to data within `T` or `D` outside of
464    /// this closure then it must be `clone`d out, for example.
465    ///
466    /// # Panics
467    ///
468    /// This function will panic if it is call recursively with any other
469    /// accessor already in scope. For example if `with` is called within `fun`,
470    /// then this function will panic. It is up to the embedder to ensure that
471    /// this does not happen.
472    pub fn with<R>(&self, fun: impl FnOnce(Access<'_, T, D>) -> R) -> R {
473        tls::get(|vmstore| {
474            fun(Access {
475                store: self.token.as_context_mut(vmstore),
476                get_data: self.get_data,
477            })
478        })
479    }
480
481    /// Returns the getter this accessor is using to project from `T` into
482    /// `D::Data`.
483    pub fn getter(&self) -> fn(&mut T) -> D::Data<'_> {
484        self.get_data
485    }
486
487    /// Changes this accessor to access `D2` instead of the current type
488    /// parameter `D`.
489    ///
490    /// This changes the underlying data access from `T` to `D2::Data<'_>`.
491    ///
492    /// # Panics
493    ///
494    /// When using this API the returned value is disconnected from `&self` and
495    /// the lifetime binding the `self` argument. An `Accessor` only works
496    /// within the context of the closure or async closure that it was
497    /// originally given to, however. This means that due to the fact that the
498    /// returned value has no lifetime connection it's possible to use the
499    /// accessor outside of `&self`, the original accessor, and panic.
500    ///
501    /// The returned value should only be used within the scope of the original
502    /// `Accessor` that `self` refers to.
503    pub fn with_getter<D2: HasData>(
504        &self,
505        get_data: fn(&mut T) -> D2::Data<'_>,
506    ) -> Accessor<T, D2> {
507        Accessor {
508            token: self.token,
509            get_data,
510        }
511    }
512
513    /// Spawn a background task which will receive an `&Accessor<T, D>` and
514    /// run concurrently with any other tasks in progress for the current
515    /// store.
516    ///
517    /// This is particularly useful for host functions which return a `stream`
518    /// or `future` such that the code to write to the write end of that
519    /// `stream` or `future` must run after the function returns.
520    ///
521    /// The returned [`JoinHandle`] may be used to cancel the task.
522    ///
523    /// # Panics
524    ///
525    /// Panics if called within a closure provided to the [`Accessor::with`]
526    /// function. This can only be called outside an active invocation of
527    /// [`Accessor::with`].
528    pub fn spawn(&self, task: impl for<'fut> AccessorTask<'fut, T, D>) -> Result<JoinHandle>
529    where
530        T: 'static,
531    {
532        let accessor = self.clone_for_spawn();
533        self.with(|mut access| access.as_context_mut().spawn_with_accessor(accessor, task))
534    }
535
536    fn clone_for_spawn(&self) -> Self {
537        Self {
538            token: self.token,
539            get_data: self.get_data,
540        }
541    }
542
543    /// Polls to see if this store contains any "interesting" tasks still within
544    /// it.
545    ///
546    /// Returns `Poll::Ready(())` if there are no more interesting tasks, and
547    /// otherwise returns `Poll::Pending`. If pending is returned then whenever
548    /// the last remaining "interesting" task has exited the provided context's
549    /// waker will be notified. Note that only the waker passed to the last call
550    /// to `poll_no_interesting_tasks` for the store will be notified, so this
551    /// is only appropriate to use once-at-a-time per store.
552    ///
553    /// The component model specification, as of this current date, does not
554    /// have a distinction between "interesting" tasks and not. The current
555    /// intention is that in a future revision of the component model this will
556    /// be distinguished at the component ABI level where tasks will be able to
557    /// flag themselves as "interesting" optionally. Additionally extra work can
558    /// be opted-in to being "interesting".
559    ///
560    /// For now what this means is that all component model tasks within this
561    /// store are considered interesting. This specifically includes the entire
562    /// duration of a task, so even all of the time after a task has returned
563    /// but before it has exited. This means that this function is, today,
564    /// effectively a proxy for "are there any more tasks still running in this
565    /// store". This can be used by embedders to determine whether there's any
566    /// more work going on, even in the background, for a particular guest.
567    /// Hosts can use this as a signal that the guest wants to stay alive a
568    /// little longer, even after a task has returned.
569    ///
570    /// In the future this predicate won't include all tasks in this store. Some
571    /// tasks will be able to flag themselves as not interesting, meaning that
572    /// when this returns ready it'd be possible that there are still tasks
573    /// remaining in the store.
574    ///
575    /// Note that at this time spawned threads within a task are always
576    /// considered uninteresting. If this function returns ready, then spawned
577    /// threads may still be in the store.
578    pub fn poll_no_interesting_tasks(&self, cx: &mut Context<'_>) -> Poll<()> {
579        self.with(|mut access| {
580            let store = access.as_context_mut().0;
581            let state = store.concurrent_state_mut_without_forcing_current_thread();
582            if state.interesting_tasks == 0 {
583                Poll::Ready(())
584            } else {
585                state.interesting_tasks_empty_waker = Some(cx.waker().clone());
586                Poll::Pending
587            }
588        })
589    }
590
591    /// Poll to see if the component instance corresponding to the specified
592    /// function is ready to run a concurrent call without queuing it (i.e. does
593    /// not have backpressure enabled and does not have a sync call in
594    /// progress).
595    ///
596    /// Returns `Poll::Ready(())` if the component instance is ready to run a
597    /// concurrent call, and otherwise returns `Poll::Pending`.  If pending is
598    /// returned then whenever the instance becomes ready for a call the
599    /// provided context's waker will be notified.  Note that only the waker
600    /// passed to the last call to `poll_ready_for_concurrent_call` for the
601    /// store will be notified (regardless of whether the same or different
602    /// `Func` is specified relative to earlier calls), so this is only
603    /// appropriate to use once-at-a-time per store.  Also note that the waker
604    /// may be notified when _any_ instance becomes callable (i.e. not
605    /// necessarily the last one polled), so this function must be called again
606    /// to determine if the instance of interest is ready.
607    pub fn poll_ready_for_concurrent_call(&self, func: Func, cx: &mut Context<'_>) -> Poll<()> {
608        self.with(|mut access| {
609            let store = access.as_context_mut().0;
610            let (_, _, _, raw_options) = func.abi_info(store);
611            let instance = func.instance().runtime_instance(raw_options.instance);
612            let state = store.instance_state(instance).concurrent_state();
613            if state.backpressure == 0 {
614                Poll::Ready(())
615            } else {
616                store
617                    .concurrent_state_mut_without_forcing_current_thread()
618                    .ready_for_concurrent_call_waker = Some(cx.waker().clone());
619                Poll::Pending
620            }
621        })
622    }
623}
624
625/// Represents an async closure which may be provided to `Accessor::spawn`,
626/// `Accessor::forward`, or `StoreContextMut::spawn`.
627// TODO: Replace this with `core::ops::AsyncFnOnce` when we are able to put `Send`
628// bound on the unnamed `Future` directly.
629//
630// As of this writing, it's not possible to specify e.g. `Send` and `Sync`
631// bounds on the `Future` type returned by an `AsyncFnOnce`.  Also, using `F:
632// Future<Output = Result<()>> + Send + Sync, FN: FnOnce(&Accessor<T>) -> F +
633// Send + Sync + 'static` fails with a type mismatch error as we cannot describe
634// that `F` should have `&Accessor<T>`'s unnamed lifetime
635//
636// Instead, this trait is used as a workaround for this limitation, the bound on `Self`
637// implementing `AsyncFnOnce()` is required for Rust to automatically infer that an async
638// closure is to be provided wherever we are accepting `impl for<'fut> AccessorTask<'fut, T, D>`
639// as argument. Otherwise, users will have to fully qualify the closure types before it
640// is accepted as a valid value. This also means that it is not intended for a user
641// to manually implement this trait for any arbitrary type, since they would first
642// have to implement `AsyncFnOnce`, which is unstable.
643//
644// The blanket implementation for this trait will ensure that `Self` is an async closure
645// that returns a `Future` that is `Send`, and lives as long as `&Accessor<T, D>`
646pub trait AccessorTask<'fut, T, D = HasSelf<T>>:
647    AsyncFnOnce(&Accessor<T, D>) -> Result<()> + Send + 'static
648where
649    D: HasData + ?Sized,
650{
651    /// Run the task.
652    fn run(self, accessor: &'fut Accessor<T, D>) -> impl Future<Output = Result<()>> + Send + 'fut;
653}
654
655impl<'fut, F, Fut, T, D> AccessorTask<'fut, T, D> for F
656where
657    T: 'static,
658    F: AsyncFnOnce(&Accessor<T, D>) -> Result<()>,
659    F: FnOnce(&'fut Accessor<T, D>) -> Fut + Send + 'static,
660    Fut: Future<Output = Result<()>> + Send + 'fut,
661    D: HasData,
662{
663    fn run(self, accessor: &'fut Accessor<T, D>) -> impl Future<Output = Result<()>> + Send + 'fut {
664        (self)(accessor)
665    }
666}
667
668/// Represents parameter and result metadata for the caller side of a
669/// guest->guest call orchestrated by a fused adapter.
670enum CallerInfo {
671    /// Metadata for a call to an async-lowered import
672    Async {
673        params: Vec<ValRaw>,
674        has_result: bool,
675    },
676    /// Metadata for a call to an sync-lowered import
677    Sync {
678        params: Vec<ValRaw>,
679        result_count: u32,
680    },
681}
682
683/// Indicates how a guest task is waiting on a waitable set.
684enum WaitMode {
685    /// The guest task is waiting using `task.wait`
686    Fiber(StoreFiber<'static>),
687    /// The guest task is waiting via a callback declared as part of an
688    /// async-lifted export.
689    Callback(Instance),
690}
691
692impl fmt::Debug for WaitMode {
693    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
694        match self {
695            Self::Fiber(_) => f.debug_tuple("Fiber").finish(),
696            Self::Callback(instance) => f.debug_tuple("Callback").field(instance).finish(),
697        }
698    }
699}
700
701/// Represents the reason a fiber is suspending itself.
702#[derive(Debug)]
703enum SuspendReason {
704    /// The fiber is waiting for an event to be delivered to the specified
705    /// waitable set or task.
706    Waiting {
707        set: TableId<WaitableSet>,
708        thread: QualifiedThreadId,
709    },
710    /// The fiber is waiting for a subtask to suspend or exit, e.g. for a
711    /// guest-to-guest call or a `subtask.cancel`.
712    YieldingToSubtask { thread: QualifiedThreadId },
713    /// The fiber has finished handling its most recent work item and is waiting
714    /// for another (or to be dropped if it is no longer needed).
715    NeedWork,
716    /// The fiber is yielding and should be resumed once other tasks have had a
717    /// chance to run.
718    Yielding { thread: QualifiedThreadId },
719    /// The fiber was explicitly suspended with a call to `thread.suspend` or
720    /// `thread.switch-to`.
721    ExplicitlySuspending { thread: QualifiedThreadId },
722}
723
724/// Represents a pending call into guest code for a given guest task.
725enum GuestCallKind {
726    /// Indicates there's an event to deliver to the task, possibly related to a
727    /// waitable set the task has been waiting on or polling.
728    DeliverEvent {
729        /// The instance to which the task belongs.
730        instance: Instance,
731        /// The waitable set the event belongs to, if any.
732        ///
733        /// If this is `None` the event will be waiting in the
734        /// `GuestTask::event` field for the task.
735        set: Option<TableId<WaitableSet>>,
736    },
737    /// Indicates that a new guest task call is pending and may be executed
738    /// using the specified closure.
739    ///
740    /// If the closure returns `Ok(Some(call))`, the `call` should be run
741    /// immediately using `handle_guest_call`.
742    StartImplicit(Box<dyn FnOnce(&mut dyn VMStore) -> Result<()> + Send + Sync>),
743    StartExplicit(Box<dyn FnOnce(&mut dyn VMStore) -> Result<()> + Send + Sync>),
744}
745
746impl fmt::Debug for GuestCallKind {
747    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
748        match self {
749            Self::DeliverEvent { instance, set } => f
750                .debug_struct("DeliverEvent")
751                .field("instance", instance)
752                .field("set", set)
753                .finish(),
754            Self::StartImplicit(_) => f.debug_tuple("StartImplicit").finish(),
755            Self::StartExplicit(_) => f.debug_tuple("StartExplicit").finish(),
756        }
757    }
758}
759
760/// The target of a suspension intrinsic.
761#[derive(Copy, Clone, Debug)]
762pub enum SuspensionTarget {
763    Resume(u32),
764    Promote(u32),
765    None,
766}
767
768/// Behavior for `resume_thread`.
769#[derive(Copy, Clone, Debug)]
770pub enum ResumeThread {
771    Promote,
772    Resume,
773    ResumeLater,
774}
775
776/// Represents a pending call into guest code for a given guest thread.
777#[derive(Debug)]
778struct GuestCall {
779    thread: QualifiedThreadId,
780    kind: GuestCallKind,
781}
782
783impl GuestCall {
784    /// Returns whether or not the call is ready to run.
785    ///
786    /// A call will not be ready to run if either:
787    ///
788    /// - the (sub-)component instance to be called has already been entered and
789    /// cannot be reentered until an in-progress call completes
790    ///
791    /// - the call is for a not-yet started task and the (sub-)component
792    /// instance to be called has backpressure enabled
793    fn is_ready(&self, store: &mut StoreOpaque) -> Result<bool> {
794        let task = store.concurrent_state_mut()?.get_mut(self.thread.task)?;
795        let async_typed = task.async_typed;
796        let instance = task.instance;
797        let state = store.instance_state(instance).concurrent_state();
798
799        let ready = match &self.kind {
800            GuestCallKind::DeliverEvent { .. } => !state.do_not_enter,
801            GuestCallKind::StartImplicit(_) => {
802                !async_typed || !(state.do_not_enter || state.backpressure > 0)
803            }
804            GuestCallKind::StartExplicit(_) => true,
805        };
806        log::trace!(
807            "call {self:?} ready? {ready} (do_not_enter: {}; backpressure: {})",
808            state.do_not_enter,
809            state.backpressure
810        );
811        Ok(ready)
812    }
813}
814
815/// Job to be run on a worker fiber.
816enum WorkerItem {
817    GuestCall(GuestCall),
818    Function(AlwaysMut<Box<dyn FnOnce(&mut dyn VMStore) -> Result<()> + Send>>),
819}
820
821/// Represents a pending work item to be handled by the event loop for a given
822/// component instance.
823enum WorkItem {
824    /// A host task to be pushed to `ConcurrentState::futures`.
825    PushFuture(AlwaysMut<HostTaskFuture>),
826    /// A fiber to resume.
827    ResumeFiber {
828        instance: RuntimeInstance,
829        thread: QualifiedThreadId,
830        fiber: StoreFiber<'static>,
831    },
832    /// A thread to resume.
833    ResumeThread {
834        instance: RuntimeInstance,
835        thread: QualifiedThreadId,
836    },
837    /// A pending call into guest code for a given guest task.
838    GuestCall {
839        instance: RuntimeInstance,
840        call: GuestCall,
841    },
842    /// A job to run on a worker fiber.
843    WorkerFunction(AlwaysMut<Box<dyn FnOnce(&mut dyn VMStore) -> Result<()> + Send>>),
844}
845
846impl fmt::Debug for WorkItem {
847    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
848        match self {
849            Self::PushFuture(_) => f.debug_tuple("PushFuture").finish(),
850            Self::ResumeFiber {
851                instance, thread, ..
852            } => f
853                .debug_struct("ResumeFiber")
854                .field("instance", instance)
855                .field("thread", thread)
856                .finish(),
857            Self::ResumeThread { instance, thread } => f
858                .debug_struct("ResumeThread")
859                .field("instance", instance)
860                .field("thread", thread)
861                .finish(),
862            Self::GuestCall { instance, call } => f
863                .debug_struct("GuestCall")
864                .field("instance", instance)
865                .field("call", call)
866                .finish(),
867            Self::WorkerFunction(_) => f.debug_tuple("WorkerFunction").finish(),
868        }
869    }
870}
871
872/// Whether a suspension intrinsic was cancelled or completed
873#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq)]
874pub(crate) enum WaitResult {
875    Cancelled,
876    Completed,
877}
878
879/// Poll the specified future until it completes on behalf of a guest->host call
880/// using a sync-lowered import.
881///
882/// This is similar to `Instance::first_poll` except it's for sync-lowered
883/// imports, meaning we don't need to handle cancellation and we can block the
884/// caller until the task completes, at which point the caller can handle
885/// lowering the result to the guest's stack and linear memory.
886pub(crate) fn poll_and_block<R: Send + Sync + 'static>(
887    store: &mut dyn VMStore,
888    host_task: EnteredHostTask,
889    future: impl Future<Output = Result<R>> + Send + 'static,
890) -> Result<R> {
891    // Poll the future once before creating a host task. The host task will be
892    // created lazily if it's needed during the poll and otherwise will be
893    // created if the future suspends.  We can use a dummy `Waker` here because
894    // we'll add the future to `ConcurrentState::futures` and poll it
895    // automatically from the event loop if it doesn't complete immediately
896    // here.
897    let mut future = Box::pin(future);
898    let poll = tls::set(store, || {
899        future
900            .as_mut()
901            .poll(&mut Context::from_waker(&Waker::noop()))
902    });
903
904    let caller = match host_task {
905        Some(caller) => caller,
906        None => bail_bug!("host task wasn't created but should have been"),
907    };
908
909    let task = match poll {
910        // It completed immediately, so no persistent host task is needed.
911        Poll::Ready(result) => return result,
912
913        // It did not complete immediately; create the host task and add it to
914        // `ConcurrentState::futures` so it will be polled via the event loop;
915        // then use `GuestThread::sync_call_set` to wait for the task to
916        // complete, suspending the current fiber until it does so.
917        Poll::Pending => {
918            let Some(task) = store.materialize_host_task_id()? else {
919                bail_bug!("current thread is not a host thread")
920            };
921
922            // Wrap the future in a closure which will stash its result in the
923            // host task and resume this fiber when it completes.
924            let future = Box::pin(async move {
925                let result = run_with_host_task_set(task, future).await??;
926                tls::get(move |store| {
927                    let state = store.concurrent_state_mut()?;
928                    let host_state = &mut state.get_mut(task)?.state;
929                    assert!(matches!(host_state, HostTaskState::CalleeStarted));
930                    *host_state = HostTaskState::CalleeFinished(Box::new(result));
931
932                    Waitable::Host(task).set_event(
933                        state,
934                        Some(Event::Subtask {
935                            status: Status::Returned,
936                        }),
937                    )?;
938
939                    Ok(())
940                })
941            }) as HostTaskFuture;
942
943            let caller_instance = store.concurrent_state_mut()?.get_mut(caller.task)?.instance;
944            store.switch_or_trap_if_may_not_suspend(caller_instance)?;
945
946            let state = store.concurrent_state_mut()?;
947            state.push_future(future);
948
949            let set = state.get_mut(caller.thread)?.sync_call_set;
950            Waitable::Host(task).join(state, Some(set))?;
951
952            store.suspend(SuspendReason::Waiting {
953                set,
954                thread: caller,
955            })?;
956
957            // Remove the `task` from the `sync_call_set` to ensure that when
958            // this function returns and the task is deleted that there are no
959            // more lingering references to this host task.
960            Waitable::Host(task).join(store.concurrent_state_mut()?, None)?;
961            task
962        }
963    };
964
965    // Retrieve and return the result.
966    let host_state = &mut store.concurrent_state_mut()?.get_mut(task)?.state;
967    match mem::replace(host_state, HostTaskState::CalleeDone { cancelled: false }) {
968        HostTaskState::CalleeFinished(result) => Ok(match result.downcast() {
969            Ok(result) => *result,
970            Err(_) => bail_bug!("host task finished with wrong type of result"),
971        }),
972        _ => bail_bug!("unexpected host task state after completion"),
973    }
974}
975
976/// Execute the specified guest call.
977fn handle_guest_call(store: &mut dyn VMStore, call: GuestCall) -> Result<()> {
978    match call.kind {
979        GuestCallKind::DeliverEvent { instance, set } => {
980            // This thread's wait on `set`, if any, which started in
981            // `Instance::wait_with_callback`, ends here.
982            if let Some(set) = set {
983                store.concurrent_state_mut()?.get_mut(set)?.stop_waiting()?;
984            }
985            let (event, waitable) = match instance.get_event(store, call.thread.task, set, true)? {
986                Some(pair) => pair,
987                None => match set {
988                    // The event which woke this thread up was taken by another
989                    // thread (e.g. via `waitable-set.poll`) or its waitable was
990                    // removed from the set before this thread had a chance to
991                    // run. Regardless, keep waiting.
992                    Some(set) => {
993                        log::trace!(
994                            "event for {:?} on {set:?} no longer present; waiting again",
995                            call.thread
996                        );
997                        return instance.wait_with_callback(
998                            store.store_opaque_mut(),
999                            call.thread,
1000                            set,
1001                        );
1002                    }
1003                    // When a thread yields it doesn't wait on a set and the
1004                    // event to be delivered isn't stored anywhere, and this is
1005                    // just handling that case so indicate that the yield is
1006                    // complete.
1007                    None => (Event::None, None),
1008                },
1009            };
1010            let state = store.concurrent_state_mut()?;
1011            let task = state.get_mut(call.thread.task)?;
1012            let runtime_instance = task.instance;
1013            let handle = waitable.map(|(_, v)| v).unwrap_or(0);
1014
1015            log::trace!(
1016                "use callback to deliver event {event:?} to {:?} for {waitable:?}",
1017                call.thread,
1018            );
1019
1020            let old_thread = store.set_thread(call.thread)?;
1021            log::trace!(
1022                "GuestCallKind::DeliverEvent: replaced {old_thread:?} with {:?} as current thread",
1023                call.thread
1024            );
1025
1026            store.enter_instance(runtime_instance);
1027
1028            let Some(callback) = store
1029                .concurrent_state_mut()?
1030                .get_mut(call.thread.task)?
1031                .callback
1032                .take()
1033            else {
1034                bail_bug!("guest task callback field not present")
1035            };
1036
1037            let code = callback(store, event, handle)?;
1038
1039            store
1040                .concurrent_state_mut()?
1041                .get_mut(call.thread.task)?
1042                .callback = Some(callback);
1043
1044            store.exit_instance(runtime_instance)?;
1045
1046            store.set_thread(old_thread)?;
1047
1048            instance.handle_callback_code(store, call.thread, runtime_instance.index, code)?;
1049
1050            log::trace!("GuestCallKind::DeliverEvent: restored {old_thread:?} as current thread");
1051        }
1052        GuestCallKind::StartImplicit(fun) => {
1053            fun(store)?;
1054        }
1055        GuestCallKind::StartExplicit(fun) => {
1056            fun(store)?;
1057        }
1058    }
1059
1060    Ok(())
1061}
1062
1063impl<T> Store<T> {
1064    /// Convenience wrapper for [`StoreContextMut::run_concurrent`].
1065    pub async fn run_concurrent<R>(&mut self, fun: impl AsyncFnOnce(&Accessor<T>) -> R) -> Result<R>
1066    where
1067        T: Send + 'static,
1068    {
1069        ensure!(
1070            self.as_context().0.concurrency_support(),
1071            "cannot use `run_concurrent` when Config::concurrency_support disabled",
1072        );
1073        self.as_context_mut().run_concurrent(fun).await
1074    }
1075
1076    #[doc(hidden)]
1077    pub fn assert_concurrent_state_empty(&mut self) {
1078        self.as_context_mut().assert_concurrent_state_empty();
1079    }
1080
1081    #[doc(hidden)]
1082    pub fn concurrent_state_table_size(&mut self) -> usize {
1083        self.as_context_mut().concurrent_state_table_size()
1084    }
1085
1086    /// Convenience wrapper for [`StoreContextMut::spawn`].
1087    pub fn spawn(
1088        &mut self,
1089        task: impl for<'fut> AccessorTask<'fut, T, HasSelf<T>>,
1090    ) -> Result<JoinHandle>
1091    where
1092        T: 'static,
1093    {
1094        self.as_context_mut().spawn(task)
1095    }
1096}
1097
1098impl<T> StoreContextMut<'_, T> {
1099    /// Assert that all the relevant tables and queues in the concurrent state
1100    /// for this store are empty.
1101    ///
1102    /// This is for sanity checking in integration tests
1103    /// (e.g. `component-async-tests`) that the relevant state has been cleared
1104    /// after each test concludes.  This should help us catch leaks, e.g. guest
1105    /// tasks which haven't been deleted despite having completed and having
1106    /// been dropped by their supertasks.
1107    ///
1108    /// Only intended for use in Wasmtime's own testing.
1109    #[doc(hidden)]
1110    pub fn assert_concurrent_state_empty(self) {
1111        let store = self.0;
1112        store
1113            .store_data_mut()
1114            .components
1115            .assert_instance_states_empty();
1116        let state = store.concurrent_state_mut().unwrap();
1117        assert!(
1118            state.table.get_mut().is_empty(),
1119            "non-empty table: {:?}",
1120            state.table.get_mut()
1121        );
1122        assert!(state.switch_item.is_none());
1123        assert!(state.next_switch_item.is_none());
1124        assert!(state.high_priority.is_empty());
1125        assert!(state.low_priority.is_empty());
1126        assert!(state.saved_next_switch_items.is_empty());
1127        assert!(state.unforced_current_thread.is_none());
1128        assert!(state.deferred_host_call_context.is_none());
1129        assert!(state.futures_mut().unwrap().is_empty());
1130        assert!(state.global_error_context_ref_counts.is_empty());
1131    }
1132
1133    /// Helper function to perform tests over the size of the concurrent state
1134    /// table which can be useful for detecting leaks.
1135    ///
1136    /// Only intended for use in Wasmtime's own testing.
1137    #[doc(hidden)]
1138    pub fn concurrent_state_table_size(&mut self) -> usize {
1139        self.0
1140            .concurrent_state_mut()
1141            .unwrap()
1142            .table
1143            .get_mut()
1144            .iter_mut()
1145            .count()
1146    }
1147
1148    /// Spawn a background task to run as part of this instance's event loop.
1149    ///
1150    /// The task will receive an `&Accessor<U>` and run concurrently with
1151    /// any other tasks in progress for the instance.
1152    ///
1153    /// Note that the task will only make progress if and when the event loop
1154    /// for this instance is run.
1155    ///
1156    /// The returned [`JoinHandle`] may be used to cancel the task.
1157    pub fn spawn(mut self, task: impl for<'fut> AccessorTask<'fut, T>) -> Result<JoinHandle>
1158    where
1159        T: 'static,
1160    {
1161        let accessor = Accessor::new(StoreToken::new(self.as_context_mut()));
1162        self.spawn_with_accessor(accessor, task)
1163    }
1164
1165    /// Internal implementation of `spawn` functions where a `store` is
1166    /// available along with an `Accessor`.
1167    fn spawn_with_accessor<D>(
1168        self,
1169        accessor: Accessor<T, D>,
1170        task: impl for<'fut> AccessorTask<'fut, T, D>,
1171    ) -> Result<JoinHandle>
1172    where
1173        T: 'static,
1174        D: HasData + ?Sized,
1175    {
1176        // Create an "abortable future" here where internally the future will
1177        // hook calls to poll and possibly spawn more background tasks on each
1178        // iteration.
1179        let (handle, future) = JoinHandle::run(async move { task.run(&accessor).await });
1180        self.0
1181            .concurrent_state_mut()?
1182            .push_future(Box::pin(async move { future.await.unwrap_or(Ok(())) }));
1183        Ok(handle)
1184    }
1185
1186    /// Run the specified closure `fun` to completion as part of this store's
1187    /// event loop.
1188    ///
1189    /// This will run `fun` as part of this store's event loop until it
1190    /// yields a result.  `fun` is provided an [`Accessor`], which provides
1191    /// controlled access to the store and its data.
1192    ///
1193    /// This function can be used to invoke [`Func::call_concurrent`] for
1194    /// example within the async closure provided here.
1195    ///
1196    /// This function will unconditionally return an error if
1197    /// [`Config::concurrency_support`] is disabled.
1198    ///
1199    /// [`Config::concurrency_support`]: crate::Config::concurrency_support
1200    ///
1201    /// # Store-blocking behavior
1202    ///
1203    /// At this time there are certain situations in which the `Future` returned
1204    /// by the `AsyncFnOnce` passed to this function will not be polled for an
1205    /// extended period of time, despite one or more `Waker::wake` events having
1206    /// occurred for the task to which it belongs.  This can manifest as the
1207    /// `Future` seeming to be "blocked" or "locked up", but is actually due to
1208    /// the `Store` being held by e.g. a blocking host function, preventing the
1209    /// `Future` from being polled. A canonical example of this is when the
1210    /// `fun` provided to this function attempts to set a timeout for an
1211    /// invocation of a wasm function. In this situation the async closure is
1212    /// waiting both on (a) the wasm computation to finish, and (b) the timeout
1213    /// to elapse. At this time this setup will not always work and the timeout
1214    /// may not reliably fire.
1215    ///
1216    /// This function will not block the current thread and as such is always
1217    /// suitable to run in an `async` context, but the current implementation of
1218    /// Wasmtime can lead to situations where a certain wasm computation is
1219    /// required to make progress the closure to make progress. This is an
1220    /// artifact of Wasmtime's historical implementation of `async` functions
1221    /// and is the topic of [#11869] and [#11870]. In the timeout example from
1222    /// above it means that Wasmtime can get "wedged" for a bit where (a) must
1223    /// progress for a readiness notification of (b) to get delivered.
1224    ///
1225    /// This effectively means that it's not possible to reliably perform a
1226    /// "select" operation within the `fun` closure, which timeouts for example
1227    /// are based on. Fixing this requires some relatively major refactoring
1228    /// work within Wasmtime itself. This is a known pitfall otherwise and one
1229    /// that is intended to be fixed one day. In the meantime it's recommended
1230    /// to apply timeouts or such to the entire `run_concurrent` call itself
1231    /// rather than internally.
1232    ///
1233    /// [#11869]: https://github.com/bytecodealliance/wasmtime/issues/11869
1234    /// [#11870]: https://github.com/bytecodealliance/wasmtime/issues/11870
1235    ///
1236    /// # Example
1237    ///
1238    /// ```
1239    /// # use {
1240    /// #   wasmtime::{
1241    /// #     error::{Result},
1242    /// #     component::{ Component, Linker, Resource, ResourceTable},
1243    /// #     Config, Engine, Store
1244    /// #   },
1245    /// # };
1246    /// #
1247    /// # struct MyResource(u32);
1248    /// # struct Ctx { table: ResourceTable }
1249    /// #
1250    /// # async fn foo() -> Result<()> {
1251    /// # let mut config = Config::new();
1252    /// # let engine = Engine::new(&config)?;
1253    /// # let mut store = Store::new(&engine, Ctx { table: ResourceTable::new() });
1254    /// # let mut linker = Linker::new(&engine);
1255    /// # let component = Component::new(&engine, "")?;
1256    /// # let instance = linker.instantiate_async(&mut store, &component).await?;
1257    /// # let foo = instance.get_typed_func::<(Resource<MyResource>,), (Resource<MyResource>,)>(&mut store, "foo")?;
1258    /// # let bar = instance.get_typed_func::<(u32,), ()>(&mut store, "bar")?;
1259    /// store.run_concurrent(async |accessor| -> wasmtime::Result<_> {
1260    ///    let resource = accessor.with(|mut access| access.get().table.push(MyResource(42)))?;
1261    ///    let (another_resource,) = foo.call_concurrent(accessor, (resource,)).await?;
1262    ///    let value = accessor.with(|mut access| access.get().table.delete(another_resource))?;
1263    ///    bar.call_concurrent(accessor, (value.0,)).await?;
1264    ///    Ok(())
1265    /// }).await??;
1266    /// # Ok(())
1267    /// # }
1268    /// ```
1269    pub async fn run_concurrent<R>(self, fun: impl AsyncFnOnce(&Accessor<T>) -> R) -> Result<R>
1270    where
1271        T: Send + 'static,
1272    {
1273        ensure!(
1274            self.0.concurrency_support(),
1275            "cannot use `run_concurrent` when Config::concurrency_support disabled",
1276        );
1277        self.do_run_concurrent(fun, false).await
1278    }
1279
1280    pub(super) async fn run_concurrent_trap_on_idle<R>(
1281        self,
1282        fun: impl AsyncFnOnce(&Accessor<T>) -> R,
1283    ) -> Result<R> {
1284        self.do_run_concurrent(fun, true).await
1285    }
1286
1287    async fn do_run_concurrent<R>(
1288        mut self,
1289        fun: impl AsyncFnOnce(&Accessor<T>) -> R,
1290        trap_on_idle: bool,
1291    ) -> Result<R> {
1292        debug_assert!(self.0.concurrency_support());
1293        let already_running = self
1294            .0
1295            .concurrent_state_mut_already_forced_current_thread()
1296            .event_loop_running;
1297        if already_running {
1298            bail!("Recursive `StoreContextMut::run_concurrent` calls not supported")
1299        }
1300        let token = StoreToken::new(self.as_context_mut());
1301
1302        struct Dropper<'a, T: 'static, V> {
1303            store: StoreContextMut<'a, T>,
1304            value: ManuallyDrop<V>,
1305        }
1306
1307        impl<'a, T, V> Drop for Dropper<'a, T, V> {
1308            fn drop(&mut self) {
1309                self.store
1310                    .0
1311                    .concurrent_state_mut_already_forced_current_thread()
1312                    .event_loop_running = false;
1313
1314                tls::set(self.store.0, || {
1315                    // SAFETY: Here we drop the value without moving it for the
1316                    // first and only time -- per the contract for `Drop::drop`,
1317                    // this code won't run again, and the `value` field will no
1318                    // longer be accessible.
1319                    unsafe { ManuallyDrop::drop(&mut self.value) }
1320                });
1321            }
1322        }
1323
1324        let accessor = &Accessor::new(token);
1325        self.0
1326            .concurrent_state_mut_already_forced_current_thread()
1327            .event_loop_running = true;
1328        let dropper = &mut Dropper {
1329            store: self,
1330            value: ManuallyDrop::new(fun(accessor)),
1331        };
1332        // SAFETY: We never move `dropper` nor its `value` field.
1333        let future = unsafe { Pin::new_unchecked(dropper.value.deref_mut()) };
1334
1335        let result = dropper
1336            .store
1337            .as_context_mut()
1338            .poll_until(future, trap_on_idle)
1339            .await;
1340
1341        if result.is_err() {
1342            dropper.store.0.set_trapped();
1343        }
1344
1345        result
1346    }
1347
1348    /// Run this store's event loop.
1349    ///
1350    /// The returned future will resolve when the specified future completes or,
1351    /// if `trap_on_idle` is true, when the event loop can't make further
1352    /// progress.
1353    async fn poll_until<R>(
1354        mut self,
1355        mut future: Pin<&mut impl Future<Output = R>>,
1356        trap_on_idle: bool,
1357    ) -> Result<R> {
1358        struct Reset<'a, T: 'static> {
1359            store: StoreContextMut<'a, T>,
1360            futures: Option<FuturesUnordered<HostTaskFuture>>,
1361        }
1362
1363        impl<'a, T> Drop for Reset<'a, T> {
1364            fn drop(&mut self) {
1365                if let Some(futures) = self.futures.take() {
1366                    *self
1367                        .store
1368                        .0
1369                        .concurrent_state_mut_already_forced_current_thread()
1370                        .futures
1371                        .get_mut() = Some(futures);
1372                }
1373            }
1374        }
1375
1376        // Keep a continuously busy event loop from holding one executor poll
1377        // indefinitely.  In particular, Tokio's cooperative budget is only
1378        // replenished when the future returns `Pending` to the executor.
1379        const MAX_TURNS_WITHOUT_YIELD: usize = 128;
1380        let mut turns_without_yield = 0;
1381
1382        loop {
1383            // Take `ConcurrentState::futures` out of the store so we can poll
1384            // it while also safely giving any of the futures inside access to
1385            // `self`.
1386            let futures = self.0.concurrent_state_mut()?.futures.get_mut().take();
1387            let mut reset = Reset {
1388                store: self.as_context_mut(),
1389                futures,
1390            };
1391            let mut next = match reset.futures.as_mut() {
1392                Some(f) => pin!(f.next()),
1393                None => bail_bug!("concurrent state missing futures field"),
1394            };
1395
1396            enum PollResult<R> {
1397                Complete(R),
1398                ProcessWork {
1399                    ready: Option<WorkItem>,
1400                    low_priority: bool,
1401                },
1402            }
1403
1404            let result = future::poll_fn(|cx| {
1405                // First, poll the future we were passed as an argument and
1406                // return immediately if it's ready.
1407                if let Poll::Ready(value) = tls::set(reset.store.0, || future.as_mut().poll(cx)) {
1408                    return Poll::Ready(Ok(PollResult::Complete(value)));
1409                }
1410
1411                // If the store has been poisoned by a trap, either before this
1412                // event loop started or while it was running, then refuse to
1413                // make any further progress. Any work items, suspended fibers,
1414                // or pending futures left over from the trap might otherwise
1415                // resume guest code or observe scheduler state which was never
1416                // unwound. Note that this is checked after polling `future`
1417                // above so that it can observe the original error, e.g. a
1418                // trap delivered over a oneshot channel by a work item.
1419                if reset.store.0.trapped() {
1420                    return Poll::Ready(Err(Trap::CannotEnterComponent.into()));
1421                }
1422
1423                // Next, poll `ConcurrentState::futures` (which includes any
1424                // pending host tasks and/or background tasks), returning
1425                // immediately if one of them fails.
1426                let next = match tls::set(reset.store.0, || next.as_mut().poll(cx)) {
1427                    Poll::Ready(Some(output)) => {
1428                        match output {
1429                            Err(e) => return Poll::Ready(Err(e)),
1430                            Ok(()) => {}
1431                        }
1432                        Poll::Ready(true)
1433                    }
1434                    Poll::Ready(None) => Poll::Ready(false),
1435                    Poll::Pending => Poll::Pending,
1436                };
1437
1438                // Next, identify the next work item to process, if any, using
1439                // the following priority order:
1440                //
1441                // - `switch_item`: Represents the guest thread we _must_ switch
1442                // to before any other thread runs per the determinism
1443                // requirements in the Component Model spec.
1444                //
1445                // - `high_priority`: "Urgent" work items, e.g. async calls have
1446                // become freshly unblocked due to backpressure clearing or
1447                // similar, stream or future state updates, etc.
1448                //
1449                // - `low_priority`: Work items such as resuming a fiber after
1450                // it yields, in which case the point is to let other items run
1451                // first.
1452                let state = reset.store.0.concurrent_state_mut()?;
1453                let mut ready = state.switch_item.take();
1454                let mut low_priority = false;
1455                if ready.is_none() {
1456                    ready = state.high_priority.pop_back();
1457                    if ready.is_none() {
1458                        ready = state.low_priority.pop_back();
1459                        low_priority = true;
1460                    }
1461                }
1462                if ready.is_some() {
1463                    return Poll::Ready(Ok(PollResult::ProcessWork {
1464                        ready,
1465                        low_priority,
1466                    }));
1467                }
1468
1469                // Finally, if we have nothing else to do right now, determine what to do
1470                // based on whether there are any pending futures in
1471                // `ConcurrentState::futures`.
1472                return match next {
1473                    Poll::Ready(true) => {
1474                        // In this case, one of the futures in
1475                        // `ConcurrentState::futures` completed
1476                        // successfully, so we return now and continue
1477                        // the outer loop in case there is another one
1478                        // ready to complete.
1479                        Poll::Ready(Ok(PollResult::ProcessWork {
1480                            ready: None,
1481                            low_priority: false,
1482                        }))
1483                    }
1484                    Poll::Ready(false) => {
1485                        // Poll the future we were passed one last time
1486                        // in case one of `ConcurrentState::futures` had
1487                        // the side effect of unblocking it.
1488                        if let Poll::Ready(value) =
1489                            tls::set(reset.store.0, || future.as_mut().poll(cx))
1490                        {
1491                            Poll::Ready(Ok(PollResult::Complete(value)))
1492                        } else {
1493                            // In this case, there are no more pending
1494                            // futures in `ConcurrentState::futures`,
1495                            // there are no remaining work items, _and_
1496                            // the future we were passed as an argument
1497                            // still hasn't completed.
1498                            if trap_on_idle {
1499                                // `trap_on_idle` is true, so we exit
1500                                // immediately.
1501
1502                                // If there are any tasks belonging to
1503                                // an instance which may not suspend, trap with
1504                                // `CannotBlockSyncTask`:
1505                                Poll::Ready(Err(if reset.store.0.any_may_not_suspend()? {
1506                                    Trap::CannotBlockSyncTask.into()
1507                                } else {
1508                                    // Otherwise, trap with `AsyncDeadlock`:
1509                                    Trap::AsyncDeadlock.into()
1510                                }))
1511                            } else {
1512                                // `trap_on_idle` is false, so we assume
1513                                // that future will wake up and give us
1514                                // more work to do when it's ready to.
1515                                Poll::Pending
1516                            }
1517                        }
1518                    }
1519                    // There is at least one pending future in
1520                    // `ConcurrentState::futures` and we have nothing
1521                    // else to do but wait for now, so we return
1522                    // `Pending`.
1523                    Poll::Pending => Poll::Pending,
1524                };
1525            })
1526            .await;
1527
1528            // Put the `ConcurrentState::futures` back into the store before we
1529            // return or handle any work items since one or more of those items
1530            // might append more futures.
1531            drop(reset);
1532
1533            match result? {
1534                // The future we were passed as an argument completed, so we
1535                // return the result.
1536                PollResult::Complete(value) => break Ok(value),
1537                // The future we were passed has not yet completed, so handle
1538                // any work items and then loop again.
1539                PollResult::ProcessWork {
1540                    ready,
1541                    low_priority,
1542                } => {
1543                    struct Dispose<'a, T: 'static> {
1544                        store: StoreContextMut<'a, T>,
1545                        ready: Option<WorkItem>,
1546                    }
1547
1548                    impl<'a, T> Drop for Dispose<'a, T> {
1549                        fn drop(&mut self) {
1550                            if let Some(item) = self.ready.take() {
1551                                match item {
1552                                    WorkItem::ResumeFiber { mut fiber, .. } => {
1553                                        fiber.dispose(self.store.0);
1554                                    }
1555                                    WorkItem::PushFuture(future) => {
1556                                        tls::set(self.store.0, move || drop(future))
1557                                    }
1558                                    _ => {}
1559                                }
1560                            }
1561                        }
1562                    }
1563
1564                    let mut dispose = Dispose {
1565                        store: self.as_context_mut(),
1566                        ready,
1567                    };
1568
1569                    // If we're about to run a low-priority task, first yield to
1570                    // the executor.  This ensures that it won't be starved of
1571                    // the ability to e.g. update the readiness of sockets,
1572                    // etc. which the guest may be using `thread.yield` along
1573                    // with `waitable-set.poll` to monitor in a CPU-heavy loop.
1574                    //
1575                    // This works for e.g. `thread.yield` and callbacks which
1576                    // return `CALLBACK_CODE_YIELD` because we queue a low
1577                    // priority item to resume the task (i.e. resume the thread
1578                    // or call the callback, respectively) just prior to
1579                    // suspending it.  Indeed, as of this writing those are the
1580                    // _only_ situations we queue low-priority tasks.
1581                    // Therefore, we interpret the guest's request to yield as
1582                    // meaning "yield to other guest tasks _and_/_or_ host
1583                    // operations such as updating socket readiness", the latter
1584                    // being the async runtime's responsibility.
1585                    //
1586                    // In the future, if this ends up causing measurable
1587                    // performance issues, this could be optimized such that we
1588                    // only yield periodically (e.g. for batches of low priority
1589                    // items) and not for each and every individual item.
1590                    if low_priority {
1591                        dispose.store.0.yield_now().await;
1592                        turns_without_yield = 0;
1593                    }
1594
1595                    if let Some(item) = dispose.ready.take() {
1596                        dispose
1597                            .store
1598                            .as_context_mut()
1599                            .handle_work_item(item)
1600                            .await?;
1601                    }
1602
1603                    turns_without_yield += 1;
1604                    if turns_without_yield == MAX_TURNS_WITHOUT_YIELD {
1605                        turns_without_yield = 0;
1606                        dispose.store.0.yield_now().await;
1607                    }
1608                }
1609            }
1610        }
1611    }
1612
1613    /// Handle the specified work item, possibly resuming a fiber if applicable.
1614    async fn handle_work_item(self, item: WorkItem) -> Result<()> {
1615        log::trace!("handle work item {item:?}");
1616        match item {
1617            WorkItem::PushFuture(future) => {
1618                self.0
1619                    .concurrent_state_mut()?
1620                    .futures_mut()?
1621                    .push(future.into_inner());
1622            }
1623            WorkItem::ResumeFiber { fiber, .. } => {
1624                self.0.resume_fiber(fiber).await?;
1625            }
1626            WorkItem::ResumeThread { thread, .. } => {
1627                if let GuestThreadState::Ready { fiber, .. } = mem::replace(
1628                    &mut self.0.concurrent_state_mut()?.get_mut(thread.thread)?.state,
1629                    GuestThreadState::Running,
1630                ) {
1631                    self.0.resume_fiber(fiber).await?;
1632                } else {
1633                    bail_bug!("cannot resume non-pending thread {thread:?}");
1634                }
1635            }
1636            WorkItem::GuestCall { call, .. } => {
1637                if call.is_ready(self.0)? {
1638                    self.0
1639                        .concurrent_state_mut()?
1640                        .get_mut(call.thread.thread)?
1641                        .wake_on_cancel = WakeOnCancel::None;
1642                    self.run_on_worker(WorkerItem::GuestCall(call)).await?;
1643                } else {
1644                    let state = self.0.concurrent_state_mut()?;
1645                    let task = state.get_mut(call.thread.task)?;
1646                    if !task.starting_sent {
1647                        task.starting_sent = true;
1648                        if let GuestCallKind::StartImplicit(_) = &call.kind {
1649                            Waitable::Guest(call.thread.task).set_event(
1650                                state,
1651                                Some(Event::Subtask {
1652                                    status: Status::Starting,
1653                                }),
1654                            )?;
1655                        }
1656                    }
1657
1658                    let instance = state.get_mut(call.thread.task)?.instance;
1659                    // If this call is for a callback-lifted task waiting on a
1660                    // waitable set, it can't run until `instance` is enterable
1661                    // again. Meanwhile, the set may have an event which another
1662                    // thread waiting on it could take, and that thread may
1663                    // even be the reason `instance` isn't enterable, so wake
1664                    // the next waiter up as well if there's still an event
1665                    // for it. Whichever thread runs first takes the event, and
1666                    // the other goes back to waiting if none is left.
1667                    if let GuestCallKind::DeliverEvent { set: Some(set), .. } = &call.kind {
1668                        let state = self.0.concurrent_state_mut()?;
1669                        if !state.get_mut(*set)?.ready.is_empty() {
1670                            state.wake_waiter(*set)?;
1671                        }
1672                    }
1673                    self.0
1674                        .instance_state(instance)
1675                        .concurrent_state()
1676                        .pending
1677                        .insert(call.thread, call.kind);
1678
1679                    // Switch back to the caller (or canceller) immediately if
1680                    // applicable since we aren't yet able to run the subtask it
1681                    // yielded to.
1682                    self.0.concurrent_state_mut()?.take_next_switch_item()?;
1683                }
1684            }
1685            WorkItem::WorkerFunction(fun) => {
1686                self.run_on_worker(WorkerItem::Function(fun)).await?;
1687            }
1688        }
1689
1690        Ok(())
1691    }
1692
1693    /// Execute the specified guest call on a worker fiber.
1694    async fn run_on_worker(self, item: WorkerItem) -> Result<()> {
1695        let worker = if let Some(fiber) = self.0.concurrent_state_mut()?.worker.take() {
1696            fiber
1697        } else {
1698            // SAFETY: the `make_fiber_unchecked` function is unsafe because the
1699            // returned fiber is unconditionally `Send` as opposed to being
1700            // conditionally send depending on the argument (in this case
1701            // `self.0`). This `async` function, however, is conditionally
1702            // `Send` depending on `self`, in this case `StoreContextMut<T>`,
1703            // which is already going to be conditionally `Send` depending on
1704            // `T`.
1705            //
1706            // The returned fiber is possibly stored within the `Store<T>` as
1707            // well. If `T: Send` then that's fine and everything's dandy. If
1708            // `T: !Send`, however, then the store is already not-`Send` meaning
1709            // that putting more actually-not-`Send` things inside of it isn't
1710            // an issue.
1711            //
1712            // The main issue here is that the returned fiber effectively can't
1713            // get transferred outside the context of the store. That's an
1714            // implementation detail we'll have to rely on, but is currently
1715            // true.
1716            unsafe {
1717                fiber::make_fiber_unchecked(self.0, move |store| {
1718                    loop {
1719                        let Some(item) = store.concurrent_state_mut()?.worker_item.take() else {
1720                            bail_bug!("worker_item not present when resuming fiber")
1721                        };
1722                        match item {
1723                            WorkerItem::GuestCall(call) => handle_guest_call(store, call)?,
1724                            WorkerItem::Function(fun) => fun.into_inner()(store)?,
1725                        }
1726
1727                        store.suspend(SuspendReason::NeedWork)?;
1728                    }
1729                })?
1730            }
1731        };
1732
1733        let worker_item = &mut self.0.concurrent_state_mut()?.worker_item;
1734        assert!(worker_item.is_none());
1735        *worker_item = Some(item);
1736
1737        self.0.resume_fiber(worker).await
1738    }
1739
1740    /// Wrap the specified host function in a future which will call it, passing
1741    /// it an `&Accessor<T>`.
1742    ///
1743    /// See the `Accessor` documentation for details.
1744    pub(crate) fn wrap_call<F, R>(self, closure: F) -> impl Future<Output = Result<R>> + 'static
1745    where
1746        T: 'static,
1747        F: FnOnce(&Accessor<T>) -> Pin<Box<dyn Future<Output = Result<R>> + Send + '_>>
1748            + Send
1749            + Sync
1750            + 'static,
1751        R: Send + Sync + 'static,
1752    {
1753        let token = StoreToken::new(self);
1754        async move {
1755            let mut accessor = Accessor::new(token);
1756            closure(&mut accessor).await
1757        }
1758    }
1759
1760    pub(crate) async fn start_instance(
1761        &mut self,
1762        instance: ModuleInstance,
1763        callee: Option<RuntimeInstance>,
1764    ) -> Result<ModuleInstance> {
1765        let (tx, rx) = oneshot::channel();
1766        let token = StoreToken::new(self.as_context_mut());
1767        self.0.queue_task(move |store| {
1768            _ = tx.send(
1769                super::instance::start_raw(&mut token.as_context_mut(store), instance, callee)
1770                    .map(|()| instance),
1771            );
1772            Ok(())
1773        })?;
1774        self.as_context_mut()
1775            .run_concurrent_trap_on_idle(async |_| {
1776                rx.await
1777                    .map_err(|_| format_err!("oneshot channel canceled"))
1778            })
1779            .await??
1780    }
1781}
1782
1783/// Return value of [`StoreOpaque::host_task_create`].
1784///
1785/// This is an `Option` to handle the dynamic `store.concurrency_support()`
1786/// property. When present this records the guest thread to restore when the
1787/// host call exits. The corresponding [`HostTask`] will need to be lazily
1788/// created if needed via [`StoreOpaque::materialize_host_task_id`].
1789pub type EnteredHostTask = Option<QualifiedThreadId>;
1790
1791impl StoreOpaque {
1792    /// Returns the currently-running thread, promoting any deferred lazy guest
1793    /// thread into a fully-materialized `CurrentThread`. Deferred [`HostTask`]s
1794    /// are not materialized.
1795    #[inline]
1796    pub(crate) fn current_thread(&mut self) -> Result<CurrentThread> {
1797        // Without concurrency support there is nothing to force.
1798        if !self.concurrency_support() {
1799            return Ok(CurrentThread::None);
1800        }
1801
1802        // If the JIT-visible current thread isn't a deferred thread then
1803        // `ConcurrentState` is already up to date.
1804        if !self
1805            .vm_store_context_mut()
1806            .current_thread_mut()
1807            .is_deferred()
1808        {
1809            return Ok(self
1810                .concurrent_state_mut_already_forced_current_thread()
1811                .unforced_current_thread);
1812        }
1813
1814        self.force_deferred_current_thread()
1815    }
1816
1817    /// Slow path of [`Self::current_thread`]: promote the deferred lazy
1818    /// thread into a fully-materialized `CurrentThread`.
1819    #[cold]
1820    fn force_deferred_current_thread(&mut self) -> Result<CurrentThread> {
1821        // The component instance whose adapters pushed the deferred frames; all
1822        // frames in a guest-to-guest, sync-to-sync call chain of fused adapters
1823        // live within a single `wasmtime::component::Instance` (because
1824        // cross-`wasmtime::component::Instance` calls don't go through fused
1825        // adapters), and guest code only ever runs as a guest thread, so the
1826        // chain's base thread is already materialized in `ConcurrentState` and
1827        // we can get the `ComponentInstanceId` shared by the whole chain from
1828        // here.
1829        let state = self.concurrent_state_mut_without_forcing_current_thread();
1830        let id = match state.unforced_current_thread.guest_task() {
1831            Some(task) => state.get_mut(task)?.instance.instance,
1832            None => bail_bug!("deferred component-model thread with non-guest base"),
1833        };
1834
1835        // Collect the deferred frames pushed inline by fused adapters, walking
1836        // the `parent` chain from innermost to the base.
1837        let mut frames = Vec::new();
1838        let mut cur = *self.vm_store_context_mut().current_thread_mut();
1839        while let Some(ptr) = cur.as_deferred() {
1840            // SAFETY: `ptr` points at a `VMDeferredThread` living in a fused
1841            // adapter's stack frame that is suspended below us on the stack
1842            // (mid-call, waiting for this nested call to return), so the
1843            // referent is still valid and exclusively ours to read.
1844            let deferred = unsafe { ptr.as_non_null().as_ref() };
1845            frames.push((
1846                deferred.callee_async != 0,
1847                deferred.callee_instance,
1848                deferred.saved_context,
1849            ));
1850            cur = deferred.parent;
1851        }
1852
1853        // Mark the current thread forced *before* replaying so that any
1854        // reentrant `force_current_thread` call short-circuits via the
1855        // non-deferred path above.
1856        *self.vm_store_context_mut().current_thread_mut() = VMLazyThread::forced();
1857
1858        // Save the current context, as we need to overwrite it while replaying
1859        // below.
1860        let current_context = *self.vm_store_context_mut().component_context_mut();
1861
1862        // Replay the deferred `enter_guest_sync_call`s outermost-first so that
1863        // the resulting `ConcurrentState` matches what the non-deferred path
1864        // would have otherwise produced.
1865        for (callee_async, callee_instance, saved_context) in frames.into_iter().rev() {
1866            // Restore the caller's context slots so that we save the correct
1867            // values into the caller's thread, exactly as the non-deferred path
1868            // would have on entry.
1869            *self.vm_store_context_mut().component_context_mut() = saved_context;
1870            let callee = RuntimeInstance {
1871                instance: id,
1872                index: RuntimeComponentInstanceIndex::from_u32(callee_instance),
1873            };
1874            self.enter_guest_sync_call(callee_async, callee)?;
1875        }
1876
1877        // Replaying done; restore the current context.
1878        *self.vm_store_context_mut().component_context_mut() = current_context;
1879
1880        Ok(self
1881            .concurrent_state_mut_without_forcing_current_thread()
1882            .unforced_current_thread)
1883    }
1884
1885    fn current_guest_thread(&mut self) -> Result<QualifiedThreadId> {
1886        match self.current_thread()?.guest() {
1887            Some(id) => Ok(*id),
1888            None => bail_bug!("current thread is not a guest thread"),
1889        }
1890    }
1891
1892    // A result of `None` may indicate that this is either the top-level event
1893    // loop, a deferred host task, or concurrency support is disabled. In all
1894    // cases we don't have an ID for the task.
1895    pub(crate) fn current_materialized_host_task(&mut self) -> Result<Option<TableId<HostTask>>> {
1896        match self.current_thread()? {
1897            CurrentThread::Host(id) => Ok(Some(id)),
1898            CurrentThread::DeferredHost(_) | CurrentThread::None => Ok(None),
1899            _ => bail_bug!("current thread is not a host thread"),
1900        }
1901    }
1902
1903    /// Returns the current host task ID, materializing a deferred host task if
1904    /// one is active. `None` represents a call from the top-level host.
1905    fn materialize_host_task_id(&mut self) -> Result<Option<TableId<HostTask>>> {
1906        Ok(self
1907            .concurrent_state_mut()?
1908            .materialize_current_host_task_id()?)
1909    }
1910
1911    fn enter_sync_call(&mut self, callee: RuntimeInstance) -> Result<()> {
1912        log::trace!("enter sync-typed call {callee:?}");
1913        let state = self.instance_state(callee).concurrent_state();
1914        let old_do_not_suspend = state.do_not_suspend;
1915        state.do_not_suspend = true;
1916
1917        let thread = self.current_guest_thread()?;
1918        let thread = self.concurrent_state_mut()?.get_mut(thread.thread)?;
1919        if thread.old_do_not_suspend.is_some() {
1920            bail_bug!("current thread already has `old_do_not_suspend` value");
1921        }
1922
1923        thread.old_do_not_suspend = Some(old_do_not_suspend);
1924
1925        Ok(())
1926    }
1927
1928    fn exit_sync_call(&mut self, callee: RuntimeInstance) -> Result<()> {
1929        log::trace!("exit sync-typed call {callee:?}");
1930        let thread = self.current_guest_thread()?;
1931        let thread = self.concurrent_state_mut()?.get_mut(thread.thread)?;
1932        let Some(old_do_not_suspend) = thread.old_do_not_suspend.take() else {
1933            bail_bug!("current thread missing `old_do_not_suspend` value");
1934        };
1935        let state = self.instance_state(callee).concurrent_state();
1936        state.do_not_suspend = old_do_not_suspend;
1937        Ok(())
1938    }
1939
1940    /// Push a `GuestTask` onto the task stack for either a sync-to-sync,
1941    /// guest-to-guest call or a sync host-to-guest call.
1942    ///
1943    /// This task will only be used for the purpose of handling calls to
1944    /// intrinsic functions; both parameter lowering and result lifting are
1945    /// assumed to be taken care of elsewhere.
1946    ///
1947    /// NB: for sync-to-sync, guest-to-guest calls we delay task construction in
1948    /// fused adapters, see `StoreOpaque::current_thread`, `VMDeferredThread`,
1949    /// and `lower_fact_enter_sync_call`. Make sure all this stuff stays in
1950    /// sync!
1951    pub(crate) fn enter_guest_sync_call(
1952        &mut self,
1953        callee_async_typed: bool,
1954        callee: RuntimeInstance,
1955    ) -> Result<()> {
1956        log::trace!("enter sync-lifted call {callee:?}");
1957        if !self.concurrency_support() {
1958            return self.enter_call_not_concurrent();
1959        }
1960
1961        // Stash the next switch item (if any) for the duration of this call
1962        // since it's for the caller's caller, not for us.  We'll restore this
1963        // in `exit_guest_sync_call`.
1964        let state = self.concurrent_state_mut()?;
1965        let item = state.next_switch_item.take();
1966        state.saved_next_switch_items.push(item);
1967
1968        let thread = self.current_thread()?;
1969        let caller = if let Some(thread) = thread.guest() {
1970            Caller::Guest { thread: *thread }
1971        } else {
1972            Caller::Host {
1973                tx: None,
1974                host_future_present: false,
1975                caller: self.materialize_host_task_id()?,
1976            }
1977        };
1978
1979        let state = self.concurrent_state_mut()?;
1980        let guest_thread = GuestTask::new(
1981            state,
1982            Box::new(move |_, _| bail_bug!("cannot lower params in sync call")),
1983            LiftResult {
1984                lift: Box::new(move |_, _| bail_bug!("cannot lift result in sync call")),
1985                ty: TypeTupleIndex::reserved_value(),
1986                memory: None,
1987                string_encoding: StringEncoding::Utf8,
1988            },
1989            caller,
1990            None,
1991            callee,
1992            callee_async_typed,
1993            false,
1994        )?;
1995
1996        Instance::from_wasmtime(self, callee.instance).add_guest_thread_to_instance_table(
1997            guest_thread.thread,
1998            self,
1999            callee.index,
2000        )?;
2001        self.set_thread(guest_thread)?;
2002
2003        if !callee_async_typed {
2004            self.enter_sync_call(callee)?;
2005        }
2006
2007        Ok(())
2008    }
2009
2010    /// Pop a `GuestTask` previously pushed using `enter_guest_sync_call`.
2011    ///
2012    /// NB: for sync-to-sync, guest-to-guest calls we delay task construction in
2013    /// fused adapters and then when the call returns we check to see if the
2014    /// task's contruction was forced and if not avoid calling out of the JIT
2015    /// code to this function. See `lower_fact_exit_sync_call`. Make sure all
2016    /// this stuff stays in sync!
2017    pub(crate) fn exit_guest_sync_call(&mut self) -> Result<()> {
2018        if !self.concurrency_support() {
2019            return Ok(self.exit_call_not_concurrent());
2020        }
2021
2022        let thread = match self.current_thread()?.guest() {
2023            Some(t) => *t,
2024            None => bail_bug!("expected task when exiting"),
2025        };
2026        let task = self.concurrent_state_mut()?.get_mut(thread.task)?;
2027        let instance = task.instance;
2028
2029        let caller = match &task.caller {
2030            &Caller::Guest { thread } => thread.into(),
2031            &Caller::Host { caller, .. } => caller
2032                .map(CurrentThread::Host)
2033                .unwrap_or(CurrentThread::None),
2034        };
2035        task.lift_result = None;
2036        task.exited = true;
2037        let async_typed = task.async_typed;
2038
2039        if !async_typed {
2040            self.exit_sync_call(instance)?;
2041        }
2042
2043        self.set_thread(caller)?;
2044
2045        log::trace!("exit sync-lifted call {instance:?}");
2046
2047        if async_typed {
2048            // If we're async-typed, returning control to our caller won't help
2049            // resolve any outstanding sync-typed call which might be in
2050            // progress, so we may need to switch or trap before exiting this
2051            // thread:
2052            self.switch_or_trap_if_may_not_suspend(instance)?;
2053        }
2054
2055        self.cleanup_thread(thread, instance, CleanupTask::Yes)?;
2056
2057        let state = self.concurrent_state_mut()?;
2058        let Some(item) = state.saved_next_switch_items.pop() else {
2059            bail_bug!("unable to pop from `saved_next_switch_items`");
2060        };
2061        if let Some(item) = mem::replace(&mut state.next_switch_item, item) {
2062            // Stash it back in the store to ensure it's cleaned up on store
2063            // drop:
2064            state.push_high_priority(item);
2065            bail_bug!("`next_switch_item` unexpectedly already set");
2066        }
2067
2068        Ok(())
2069    }
2070
2071    /// Similar to `enter_guest_sync_call` except for when the guest makes a
2072    /// transition to the host.
2073    ///
2074    /// This initially records a deferred host call. A full [`HostTask`] should
2075    /// be allocated later if needed via
2076    /// [`StoreOpaque::materialize_host_task_id`].
2077    pub(crate) fn host_task_create(&mut self) -> Result<EnteredHostTask> {
2078        if !self.concurrency_support() {
2079            self.enter_call_not_concurrent()?;
2080            return Ok(None);
2081        }
2082        let caller = self.current_guest_thread()?;
2083        log::trace!("new deferred host task with caller {caller:?}");
2084
2085        self.set_thread(CurrentThread::DeferredHost(caller))?;
2086        let state = self.concurrent_state_mut()?;
2087        debug_assert!(state.deferred_host_call_context.is_none());
2088        state.deferred_host_call_context = Some(CallContext::default());
2089        state.debug_assert_deferred_host_invariant();
2090        Ok(Some(caller))
2091    }
2092
2093    /// Dual of `host_task_create` and signifies that the host has finished and
2094    /// will be cleaned up.
2095    ///
2096    /// Note that this isn't invoked when the host is invoked asynchronously and
2097    /// the host isn't complete yet. In that situation the host task persists
2098    /// and will be cleaned up separately in `subtask_drop`
2099    pub(crate) fn host_task_delete(
2100        &mut self,
2101        original_task: EnteredHostTask,
2102        materialized_task: Option<TableId<HostTask>>,
2103    ) -> Result<()> {
2104        match original_task {
2105            Some(caller) => {
2106                self.set_thread(caller)?;
2107                if materialized_task.is_none() {
2108                    let state = self.concurrent_state_mut()?;
2109                    let context = state
2110                        .deferred_host_call_context
2111                        .take()
2112                        .expect("deferred host call context should be present");
2113                    debug_assert!(context.is_empty());
2114                    state.debug_assert_deferred_host_invariant();
2115                }
2116                log::trace!(
2117                    "delete host task with caller {original_task:?} and materialized as {materialized_task:?}"
2118                );
2119                if let Some(task) = materialized_task {
2120                    Waitable::Host(task).delete_from(self)?;
2121                }
2122            }
2123            None => {
2124                debug_assert!(materialized_task.is_none());
2125                self.exit_call_not_concurrent();
2126            }
2127        }
2128        Ok(())
2129    }
2130
2131    /// Helper function to retrieve the `InstanceState` for the
2132    /// specified instance.
2133    fn instance_state(&mut self, instance: RuntimeInstance) -> &mut InstanceState {
2134        self.component_instance_mut(instance.instance)
2135            .instance_state(instance.index)
2136    }
2137
2138    /// Configure the currently running `thread`.
2139    ///
2140    /// This will save off any state necessary for the previous thread, if
2141    /// applicable, and then it'll additionally update state for `thread` if
2142    /// needed too.
2143    pub(crate) fn set_thread(&mut self, thread: impl Into<CurrentThread>) -> Result<CurrentThread> {
2144        let thread = thread.into();
2145        let state = self.concurrent_state_mut()?;
2146        state.debug_assert_deferred_host_invariant();
2147        let old_thread = mem::replace(&mut state.unforced_current_thread, thread);
2148
2149        state.handle_thread_switch(old_thread, thread)?;
2150
2151        // First thing to do after swapping threads is updating the context
2152        // slots for this thread within the store. This restores the behavior of
2153        // `context.{get,set}`. This involves taking the old state out of the
2154        // store, saving it in the thread that's being swapped from, and doing
2155        // the inverse for the new thread. When debug assertions are enabled
2156        // this also leaves behind sentinel values to try to uncover bugs where
2157        // this may be forgotten.
2158        if let Some(old_thread) = old_thread.guest() {
2159            let old_context = *self.vm_store_context_mut().component_context_mut();
2160            self.concurrent_state_mut()?
2161                .get_mut(old_thread.thread)?
2162                .context = old_context;
2163        }
2164        if cfg!(debug_assertions) {
2165            *self.vm_store_context_mut().component_context_mut() =
2166                [u32::MAX; NUM_COMPONENT_CONTEXT_SLOTS];
2167        }
2168        if let Some(thread) = thread.guest() {
2169            let thread = self.concurrent_state_mut()?.get_mut(thread.thread)?;
2170            let context = thread.context;
2171            if cfg!(debug_assertions) {
2172                thread.context = [u32::MAX; NUM_COMPONENT_CONTEXT_SLOTS];
2173            }
2174            *self.vm_store_context_mut().component_context_mut() = context;
2175        }
2176
2177        // Keep the JIT-visible current-thread pointer in sync.
2178        *self.vm_store_context_mut().current_thread_mut() = if thread.is_none() {
2179            VMLazyThread::none()
2180        } else {
2181            VMLazyThread::forced()
2182        };
2183
2184        Ok(old_thread)
2185    }
2186
2187    /// Call `switch_if_may_not_suspend` and trap if it returns `false`.
2188    fn switch_or_trap_if_may_not_suspend(&mut self, instance: RuntimeInstance) -> Result<()> {
2189        if self.switch_if_may_not_suspend(instance)? {
2190            Ok(())
2191        } else {
2192            Err(Trap::CannotBlockSyncTask.into())
2193        }
2194    }
2195
2196    /// Check if the specified instance has a sync-typed call in progress; if so
2197    /// attempt to switch to another ready thread for that instance, and if no
2198    /// such thread exists, return false.
2199    fn switch_if_may_not_suspend(&mut self, instance: RuntimeInstance) -> Result<bool> {
2200        // Call this for the side effect of forcing any deferred task creation,
2201        // which may influence the value of `ConcurrentState::do_not_suspend`
2202        // below:
2203        self.concurrent_state_mut()?;
2204
2205        Ok(!self.concurrency_support()
2206            || !self
2207                .instance_state(instance)
2208                .concurrent_state()
2209                .do_not_suspend
2210            || self
2211                .concurrent_state_mut()?
2212                .promote_instance_local_thread_work_item(instance)?)
2213    }
2214
2215    /// Record that we're about to enter a (sub-)component instance which does
2216    /// not support more than one concurrent, stackful activation, meaning it
2217    /// cannot be entered again until the next call returns.
2218    fn enter_instance(&mut self, instance: RuntimeInstance) {
2219        log::trace!("enter {instance:?}");
2220        self.instance_state(instance)
2221            .concurrent_state()
2222            .do_not_enter = true;
2223    }
2224
2225    /// Record that we've exited a (sub-)component instance previously entered
2226    /// with `Self::enter_instance` and then calls `Self::partition_pending`.
2227    /// See the documentation for the latter for details.
2228    fn exit_instance(&mut self, instance: RuntimeInstance) -> Result<()> {
2229        log::trace!("exit {instance:?}");
2230        self.instance_state(instance)
2231            .concurrent_state()
2232            .do_not_enter = false;
2233        self.partition_pending(instance)
2234    }
2235
2236    /// Iterate over `InstanceState::pending`, moving any ready items into the
2237    /// "high priority" work item queue.
2238    ///
2239    /// Also, notify `ConcurrentState::ready_for_concurrent_call_waker` if
2240    /// present.
2241    ///
2242    /// See `GuestCall::is_ready` for details.
2243    fn partition_pending(&mut self, instance: RuntimeInstance) -> Result<()> {
2244        for (thread, kind) in
2245            mem::take(&mut self.instance_state(instance).concurrent_state().pending).into_iter()
2246        {
2247            let call = GuestCall { thread, kind };
2248            if call.is_ready(self)? {
2249                self.concurrent_state_mut()?
2250                    .push_high_priority(WorkItem::GuestCall { instance, call });
2251            } else {
2252                self.instance_state(instance)
2253                    .concurrent_state()
2254                    .pending
2255                    .insert(call.thread, call.kind);
2256            }
2257        }
2258
2259        if let Some(waker) = self
2260            .concurrent_state_mut()?
2261            .ready_for_concurrent_call_waker
2262            .take()
2263        {
2264            waker.wake();
2265        }
2266
2267        Ok(())
2268    }
2269
2270    /// Implements the `backpressure.{inc,dec}` intrinsics.
2271    pub(crate) fn backpressure_modify(
2272        &mut self,
2273        caller_instance: RuntimeInstance,
2274        modify: impl FnOnce(u16) -> Option<u16>,
2275    ) -> Result<()> {
2276        let state = self.instance_state(caller_instance).concurrent_state();
2277        let old = state.backpressure;
2278        let new = modify(old).ok_or_else(|| Trap::BackpressureOverflow)?;
2279        state.backpressure = new;
2280
2281        if old > 0 && new == 0 {
2282            // Backpressure was previously enabled and is now disabled; move any
2283            // newly-eligible guest calls to the "high priority" queue.
2284            self.partition_pending(caller_instance)?;
2285        }
2286
2287        Ok(())
2288    }
2289
2290    /// Resume the specified fiber, giving it exclusive access to the specified
2291    /// store.
2292    async fn resume_fiber(&mut self, fiber: StoreFiber<'static>) -> Result<()> {
2293        let old_thread = self.current_thread()?;
2294        log::trace!("resume_fiber: save current thread {old_thread:?}");
2295
2296        let fiber = fiber::resolve_or_release(self, fiber).await?;
2297
2298        self.set_thread(old_thread)?;
2299
2300        let state = self.concurrent_state_mut()?;
2301
2302        if let Some(ot) = old_thread.guest() {
2303            state.get_mut(ot.thread)?.state = GuestThreadState::Running;
2304        }
2305        log::trace!("resume_fiber: restore current thread {old_thread:?}");
2306
2307        if let Some(mut fiber) = fiber {
2308            log::trace!("resume_fiber: suspend reason {:?}", &state.suspend_reason);
2309            // See the `SuspendReason` documentation for what each case means.
2310            let reason = match state.suspend_reason.take() {
2311                Some(r) => r,
2312                None => bail_bug!("suspend reason missing when resuming fiber"),
2313            };
2314            match reason {
2315                SuspendReason::NeedWork => {
2316                    if state.worker.is_none() {
2317                        state.worker = Some(fiber);
2318                    } else {
2319                        fiber.dispose(self);
2320                    }
2321                }
2322                SuspendReason::Yielding { thread } => {
2323                    state.get_mut(thread.thread)?.state = GuestThreadState::Ready { fiber };
2324                    let instance = state.get_mut(thread.task)?.instance;
2325                    state.push_low_priority(WorkItem::ResumeThread { instance, thread });
2326                }
2327                SuspendReason::ExplicitlySuspending { thread } => {
2328                    state.get_mut(thread.thread)?.state = GuestThreadState::Suspended(fiber);
2329                }
2330                SuspendReason::Waiting { set, thread } => {
2331                    let old = state
2332                        .get_mut(set)?
2333                        .waiting
2334                        .insert(thread, WaitMode::Fiber(fiber));
2335                    assert!(old.is_none());
2336                }
2337                SuspendReason::YieldingToSubtask { thread } => {
2338                    // In this case, the thread has either invoked or sent a
2339                    // cancel request to a subtask, and is now yielding to that
2340                    // subtask.  According to the CM spec, that subtask may only
2341                    // yield back to the original thread the first time it
2342                    // suspends or exits (or a thread that it has resumed
2343                    // suspends or exits, etc.), which we ensure by setting
2344                    // `ConcurrentState::next_switch_item` here.
2345
2346                    let item = WorkItem::ResumeFiber {
2347                        instance: state.get_mut(thread.task)?.instance,
2348                        thread,
2349                        fiber,
2350                    };
2351
2352                    log::trace!("set next switch item to {item:?}");
2353                    if state.next_switch_item.replace(item).is_some() {
2354                        // This should be unreachable per the save/restore code
2355                        // in `Self::suspend`.
2356                        bail_bug!(
2357                            "`ConcurrentState::next_switch_item` was already `Some(_)` when \
2358                             a thread wanted to wait on a subtask"
2359                        );
2360                    }
2361                }
2362            };
2363        } else {
2364            log::trace!("resume_fiber: fiber has exited");
2365        }
2366
2367        Ok(())
2368    }
2369
2370    /// Suspend the current fiber, storing the reason in
2371    /// `ConcurrentState::suspend_reason` to indicate the conditions under which
2372    /// it should be resumed.
2373    ///
2374    /// See the `SuspendReason` documentation for details.
2375    fn suspend(&mut self, reason: SuspendReason) -> Result<()> {
2376        log::trace!("suspend fiber: {reason:?}");
2377
2378        let state = self.concurrent_state_mut()?;
2379
2380        // If we're yielding or waiting on behalf of a guest thread, we'll be
2381        // overwriting the current thread, so save it now and restore it once
2382        // we've resumed.
2383        //
2384        // Also, if we're yielding to a subtask, we're about to overwrite
2385        // `ConcurrentState::next_switch_item`, so also save and restore that.
2386        let (save_and_restore_thread, save_and_restore_next_switch_item) = match &reason {
2387            SuspendReason::Yielding { .. }
2388            | SuspendReason::Waiting { .. }
2389            | SuspendReason::ExplicitlySuspending { .. } => {
2390                // If there's a thread waiting for this subtask to suspend, this
2391                // is a good time to switch back to it.
2392                if state.switch_item.is_none() {
2393                    state.take_next_switch_item()?;
2394                }
2395
2396                (true, false)
2397            }
2398            SuspendReason::YieldingToSubtask { .. } => (true, true),
2399            SuspendReason::NeedWork => (false, false),
2400        };
2401
2402        let old_next_switch_item = if save_and_restore_next_switch_item {
2403            let item = state.next_switch_item.take();
2404            // Note that we store it in the table here rather than directly in a
2405            // local variable to ensure the fiber is disposed of properly if we
2406            // end up trapping or panicking.
2407            Some(state.push(item)?)
2408        } else {
2409            None
2410        };
2411
2412        let old_guest_thread = if save_and_restore_thread {
2413            self.current_thread()?
2414        } else {
2415            CurrentThread::None
2416        };
2417
2418        let waiting_set = match &reason {
2419            SuspendReason::Waiting { set, .. } => Some(*set),
2420            _ => None,
2421        };
2422
2423        let suspend_reason = &mut self.concurrent_state_mut()?.suspend_reason;
2424        assert!(suspend_reason.is_none());
2425        *suspend_reason = Some(reason);
2426
2427        // We'll panic if we call `Self::with_blocking` when the fiber is being
2428        // disposed, so check for that and exit ASAP if appropriate.
2429        if !self.fiber_async_state_mut().can_block() {
2430            return Err(format_err!("future dropped"));
2431        }
2432
2433        // Count this thread as waiting on `waiting_set`, if any, for as long as
2434        // it's suspended.
2435        if let Some(set) = waiting_set {
2436            self.concurrent_state_mut()?.get_mut(set)?.num_waiting += 1;
2437        }
2438
2439        self.with_blocking(|_, cx| cx.suspend(StoreFiberYield::ReleaseStore))?;
2440
2441        if let Some(set) = waiting_set {
2442            self.concurrent_state_mut()?.get_mut(set)?.stop_waiting()?;
2443        }
2444
2445        if save_and_restore_thread {
2446            self.set_thread(old_guest_thread)?;
2447        }
2448
2449        if let Some(item) = old_next_switch_item {
2450            let state = self.concurrent_state_mut()?;
2451            state.next_switch_item = state.delete(item)?;
2452        }
2453
2454        Ok(())
2455    }
2456
2457    fn wait_for_event(
2458        &mut self,
2459        caller_instance: RuntimeInstance,
2460        waitable: Waitable,
2461    ) -> Result<()> {
2462        let caller = self.current_guest_thread()?;
2463        let state = self.concurrent_state_mut()?;
2464
2465        waitable.trap_if_in_waitable_set(state)?;
2466
2467        let set = state.get_mut(caller.thread)?.sync_call_set;
2468        waitable.join(state, Some(set))?;
2469
2470        self.switch_or_trap_if_may_not_suspend(caller_instance)?;
2471
2472        self.suspend(SuspendReason::Waiting {
2473            set,
2474            thread: caller,
2475        })?;
2476        let state = self.concurrent_state_mut()?;
2477
2478        waitable.join(state, None)
2479    }
2480
2481    /// Cleans up the data structures backing the `guest_thread` specified,
2482    /// removing it from the internal tables of `runtime_instance` as well.
2483    ///
2484    /// This function is used whenever a guest thread has fully exited and
2485    /// completed. This'll clean up the associated `GuestThread` structure and
2486    /// related resources it contains.
2487    ///
2488    /// Other functionality that this implements is:
2489    ///
2490    /// * This will perform conditional cleanup of the `GuestTask` that owns
2491    ///   this thread if `cleanup_task` is `CleanupTask::Yes`.
2492    /// * If there are no more threads in the `GuestTask` that this thread is
2493    ///   associated with, and if the task hasn't produced a result (e.g. it's not
2494    ///   returned or cancelled), then a trap will be raised that a result
2495    ///   wasn't ever produced.
2496    /// * If this task is finished, meaning the top-level thread exited and
2497    ///   additionally it's been returned or cancelled, then this will handle
2498    ///   management of the store's "active interesting tasks" counter.
2499    ///
2500    /// Effectively this is intended to be a "narrow waist" through which many
2501    /// destruction operations are funneled through.
2502    fn cleanup_thread(
2503        &mut self,
2504        guest_thread: QualifiedThreadId,
2505        runtime_instance: RuntimeInstance,
2506        cleanup_task: CleanupTask,
2507    ) -> Result<()> {
2508        let state = self.concurrent_state_mut()?;
2509        // If we never suspended, we never had a chance to deliver a subtask
2510        // status update, if any, to our caller, so we do that here:
2511        state.take_next_switch_item()?;
2512        let thread_data = state.get_mut(guest_thread.thread)?;
2513        let sync_call_set = thread_data.sync_call_set;
2514        if let Some(guest_id) = thread_data.instance_rep {
2515            self.instance_state(runtime_instance)
2516                .thread_handle_table()
2517                .guest_thread_remove(guest_id)?;
2518        }
2519        let state = self.concurrent_state_mut()?;
2520
2521        // Clean up any pending subtasks in the sync_call_set
2522        for waitable in mem::take(&mut state.get_mut(sync_call_set)?.ready) {
2523            if let Some(Event::Subtask {
2524                status: Status::Returned | Status::ReturnCancelled,
2525            }) = waitable.common(self.concurrent_state_mut()?)?.event
2526            {
2527                waitable.delete_from(self)?;
2528            }
2529        }
2530
2531        let state = self.concurrent_state_mut()?;
2532        state.delete(guest_thread.thread)?;
2533        state.delete(sync_call_set)?;
2534        let task = state.get_mut(guest_thread.task)?;
2535        task.threads.remove(&guest_thread.thread);
2536
2537        if task.threads.is_empty() && !task.returned_or_cancelled() {
2538            bail!(Trap::NoAsyncResult);
2539        }
2540        let ready_to_delete = task.ready_to_delete();
2541
2542        if !task.decremented_interesting_task_count && task.exited && task.returned_or_cancelled() {
2543            task.decremented_interesting_task_count = true;
2544
2545            debug_assert!(state.interesting_tasks > 0);
2546            state.interesting_tasks -= 1;
2547            if state.interesting_tasks == 0
2548                && let Some(waker) = state.interesting_tasks_empty_waker.take()
2549            {
2550                waker.wake();
2551            }
2552        }
2553
2554        match cleanup_task {
2555            CleanupTask::Yes => {
2556                if ready_to_delete {
2557                    Waitable::Guest(guest_thread.task).delete_from(self)?;
2558                }
2559            }
2560            CleanupTask::No => {}
2561        }
2562
2563        Ok(())
2564    }
2565
2566    /// Performs cancellation of the `guest_task` specified with the
2567    /// precondition that the task hasn't lowered its parameters.
2568    ///
2569    /// In this situation the task hasn't ever been started meaning it hasn't
2570    /// actually run any wasm code yet. This requires cleaning up metadata such
2571    /// as thread information attached to the task.
2572    ///
2573    /// The main two entrypoints for this function are:
2574    ///
2575    /// * Task cancellation via `subtask.cancel`, the intrinsic.
2576    /// * Dropping a host `call_async` future which needs to cancel the task
2577    ///   because it cannot reference its parameters any more.
2578    fn cancel_guest_subtask_without_lowered_parameters(
2579        &mut self,
2580        caller_instance: RuntimeInstance,
2581        guest_task: TableId<GuestTask>,
2582    ) -> Result<()> {
2583        let concurrent_state = self.concurrent_state_mut()?;
2584        let task = concurrent_state.get_mut(guest_task)?;
2585        assert!(!task.already_lowered_parameters());
2586        // The task is in a `starting` state, meaning it hasn't run at
2587        // all yet.  Here we update its fields to indicate that it is
2588        // ready to delete immediately once `subtask.drop` is called.
2589        task.lower_params = None;
2590        task.lift_result = None;
2591        task.exited = true;
2592        let instance = task.instance;
2593
2594        // Clean up the thread within this task as it's now never going
2595        // to run.
2596        assert_eq!(1, task.threads.len());
2597        let thread = *task.threads.iter().next().unwrap();
2598        self.cleanup_thread(
2599            QualifiedThreadId {
2600                task: guest_task,
2601                thread,
2602            },
2603            caller_instance,
2604            CleanupTask::No,
2605        )?;
2606
2607        // Not yet started; cancel and remove from pending
2608        let pending = &mut self.instance_state(instance).concurrent_state().pending;
2609        let pending_count = pending.len();
2610        pending.retain(|thread, _| thread.task != guest_task);
2611        // If there were no pending threads for this task, we're in an error state
2612        if pending.len() == pending_count {
2613            bail!(Trap::SubtaskCancelAfterTerminal);
2614        }
2615        Ok(())
2616    }
2617
2618    /// Used by `ResourceTables` to record the scope of a borrow to get undone
2619    /// in the future.
2620    pub(crate) fn current_scope(&mut self) -> Result<Option<CurrentScope>> {
2621        if !self.concurrency_support() {
2622            return Ok(self
2623                .current_scope_id_not_concurrent()?
2624                .map(|id| CurrentScope::Id(Scope::Id(id))));
2625        }
2626
2627        Ok(match self.current_thread()? {
2628            CurrentThread::Guest(id) => Some(CurrentScope::Id(Scope::Id(id.task.rep()))),
2629            CurrentThread::Host(id) => Some(CurrentScope::Id(Scope::HostId(id.rep()))),
2630            CurrentThread::DeferredHost(_) => Some(CurrentScope::DeferredHost),
2631            CurrentThread::None => return Ok(None),
2632        })
2633    }
2634
2635    pub(crate) fn queue_task(
2636        &mut self,
2637        task: impl FnOnce(&mut dyn VMStore) -> Result<()> + Send + 'static,
2638    ) -> Result<()> {
2639        self.concurrent_state_mut()?
2640            .push_high_priority(WorkItem::WorkerFunction(AlwaysMut::new(Box::new(task))));
2641        Ok(())
2642    }
2643
2644    /// Used in `poll_until` just prior to trapping due to a "deadlock"
2645    /// condition.
2646    ///
2647    /// This helps us distinguish between a simple deadlock condition (where no
2648    /// work is available for the event loop to do, nor is there any way for new
2649    /// work to be added) and a "cannot block sync task" condition where at
2650    /// least once instance has an outstanding sync-typed task running, in which
2651    /// case we'll trap with a different error message.
2652    fn any_may_not_suspend(&mut self) -> Result<bool> {
2653        // Note that this currently requires a linear search across the whole
2654        // `ConcurrentState::table`.  We _could_ optimize that, but since (1)
2655        // this function is only used when trapping, (2) the only thing you can
2656        // really do with a store that's been poisoned by a trap is drop it, and
2657        // (3) we must do a linear search through the table when dropping a
2658        // store anyway to dispose of fibers, it's reasonable for us to also do
2659        // a linear search here.
2660        Ok(self
2661            .concurrent_state_mut()?
2662            .table
2663            .get_mut()
2664            .iter_mut()
2665            .filter_map(|(_, entry)| {
2666                if let Some(task) = entry.downcast_ref::<GuestTask>() {
2667                    Some(task.instance)
2668                } else {
2669                    None
2670                }
2671            })
2672            .collect::<Vec<_>>()
2673            .into_iter()
2674            .any(|instance| {
2675                self.instance_state(instance)
2676                    .concurrent_state()
2677                    .do_not_suspend
2678            }))
2679    }
2680}
2681
2682enum CleanupTask {
2683    Yes,
2684    No,
2685}
2686
2687impl Instance {
2688    /// Get the next pending event for the specified task and (optional)
2689    /// waitable set, along with the waitable handle if applicable.
2690    fn get_event(
2691        self,
2692        store: &mut StoreOpaque,
2693        guest_task: TableId<GuestTask>,
2694        set: Option<TableId<WaitableSet>>,
2695        cancellable: bool,
2696    ) -> Result<Option<(Event, Option<(Waitable, u32)>)>> {
2697        let state = store.concurrent_state_mut()?;
2698
2699        let task = state.get_mut(guest_task)?;
2700        let event = &mut task.event;
2701        if let Some(ev) = event
2702            && (cancellable || !matches!(ev, Event::Cancelled))
2703        {
2704            log::trace!("deliver event {ev:?} to {guest_task:?}");
2705
2706            if matches!(ev, Event::Cancelled) {
2707                task.cancel_request_delivered = true;
2708            }
2709
2710            let ev = *ev;
2711            *event = None;
2712            return Ok(Some((ev, None)));
2713        }
2714
2715        let set = match set {
2716            Some(set) => set,
2717            None => return Ok(None),
2718        };
2719        let waitable = match state.get_mut(set)?.ready.pop_first() {
2720            Some(v) => v,
2721            None => return Ok(None),
2722        };
2723
2724        let common = waitable.common(state)?;
2725        let handle = match common.handle {
2726            Some(h) => h,
2727            None => bail_bug!("handle not set when delivering event"),
2728        };
2729        let event = match common.event.take() {
2730            Some(e) => e,
2731            None => bail_bug!("event not set when delivering event"),
2732        };
2733
2734        log::trace!(
2735            "deliver event {event:?} to {guest_task:?} for {waitable:?} (handle {handle}); set {set:?}"
2736        );
2737
2738        waitable.on_delivery(store, self, event)?;
2739
2740        Ok(Some((event, Some((waitable, handle)))))
2741    }
2742
2743    /// Handle the `CallbackCode` returned from an async-lifted export or its
2744    /// callback.
2745    ///
2746    /// If this returns `Ok(Some(call))`, then `call` should be run immediately
2747    /// using `handle_guest_call`.
2748    fn handle_callback_code(
2749        self,
2750        store: &mut StoreOpaque,
2751        guest_thread: QualifiedThreadId,
2752        runtime_instance: RuntimeComponentInstanceIndex,
2753        code: u32,
2754    ) -> Result<()> {
2755        let (code, set) = unpack_callback_code(code);
2756
2757        log::trace!("received callback code from {guest_thread:?}: {code} (set: {set})");
2758
2759        let state = store.concurrent_state_mut()?;
2760
2761        state.take_next_switch_item()?;
2762
2763        let get_set = |store: &mut StoreOpaque, handle| -> Result<_> {
2764            let set = store
2765                .instance_state(self.runtime_instance(runtime_instance))
2766                .handle_table()
2767                .waitable_set_rep(handle)?;
2768
2769            Ok(TableId::<WaitableSet>::new(set))
2770        };
2771
2772        match code {
2773            callback_code::EXIT => {
2774                log::trace!("implicit thread {guest_thread:?} completed");
2775                let task = store.concurrent_state_mut()?.get_mut(guest_thread.task)?;
2776                task.exited = true;
2777                task.callback = None;
2778
2779                let runtime_instance = self.runtime_instance(runtime_instance);
2780
2781                // Since we're async-typed, returning control to our caller
2782                // won't help resolve any outstanding sync-typed call which
2783                // might be in progress, so we may need to switch or trap before
2784                // exiting this thread:
2785                store.switch_or_trap_if_may_not_suspend(runtime_instance)?;
2786
2787                store.cleanup_thread(guest_thread, runtime_instance, CleanupTask::Yes)?;
2788            }
2789            callback_code::YIELD => {
2790                // Set `GuestTask::wake_on_cancel` to allow `subtask.cancel` to
2791                // promote this thread if appropriate.
2792                let old = state
2793                    .get_mut(guest_thread.thread)?
2794                    .wake_on_cancel
2795                    .replace(WakeOnCancel::Yielding);
2796                if !old.is_none() {
2797                    bail_bug!("thread unexpectedly had wake_on_cancel set");
2798                }
2799
2800                // If an `Event::Cancelled` is pending in `GuestTask::event`
2801                // when this call is delivered, we'll deliver that; otherwise,
2802                // we'll deliver `Event::None`.  Note that the `Event::None` is
2803                // deliberately not stored in `GuestTask::event` since that's
2804                // shared by all threads of the task and the `Event::None` is
2805                // only for this thread.
2806                let call = GuestCall {
2807                    thread: guest_thread,
2808                    kind: GuestCallKind::DeliverEvent {
2809                        instance: self,
2810                        set: None,
2811                    },
2812                };
2813                // Push this thread onto the "low priority" queue so it runs
2814                // after any other threads have had a chance to run.
2815                state.push_low_priority(WorkItem::GuestCall {
2816                    instance: self.runtime_instance(runtime_instance),
2817                    call,
2818                });
2819            }
2820            callback_code::WAIT => {
2821                let set = get_set(store, set)?;
2822                self.wait_with_callback(store, guest_thread, set)?;
2823            }
2824            _ => bail!(Trap::UnsupportedCallbackCode),
2825        }
2826
2827        Ok(())
2828    }
2829
2830    /// Start waiting on `set` with `guest_thread`, which uses a callback.
2831    ///
2832    /// The thread will receive an event from `set` via a
2833    /// `GuestCallKind::DeliverEvent` work item, either immediately if one is
2834    /// already available, or once one is published.
2835    fn wait_with_callback(
2836        self,
2837        store: &mut StoreOpaque,
2838        guest_thread: QualifiedThreadId,
2839        set: TableId<WaitableSet>,
2840    ) -> Result<()> {
2841        let state = store.concurrent_state_mut()?;
2842        // This thread's wait ends when that work item is handled in
2843        // `handle_guest_call`.
2844        state.get_mut(set)?.num_waiting += 1;
2845
2846        if state.get_mut(guest_thread.task)?.event.is_some()
2847            || !state.get_mut(set)?.ready.is_empty()
2848        {
2849            // An event is immediately available; deliver it ASAP.
2850            let instance = state.get_mut(guest_thread.task)?.instance;
2851            state.push_high_priority(WorkItem::GuestCall {
2852                instance,
2853                call: GuestCall {
2854                    thread: guest_thread,
2855                    kind: GuestCallKind::DeliverEvent {
2856                        instance: self,
2857                        set: Some(set),
2858                    },
2859                },
2860            });
2861            return Ok(());
2862        }
2863
2864        // No event is immediately available, so this thread is about to block.
2865        // Like any other suspension point (e.g. the callback `EXIT` path), if
2866        // a sync-typed call is in progress for this instance we must first
2867        // switch to another ready thread of that instance, or trap if there is
2868        // none.
2869        let instance = store
2870            .concurrent_state_mut()?
2871            .get_mut(guest_thread.task)?
2872            .instance;
2873        store.switch_or_trap_if_may_not_suspend(instance)?;
2874
2875        // Register to be woken up when an event is published for this waitable
2876        // set.
2877        //
2878        // Here we also set `GuestTask::wake_on_cancel` which allows
2879        // `subtask.cancel` to interrupt the wait.
2880        let state = store.concurrent_state_mut()?;
2881        let old = state
2882            .get_mut(guest_thread.thread)?
2883            .wake_on_cancel
2884            .replace(WakeOnCancel::Waiting(set));
2885        if !old.is_none() {
2886            bail_bug!("thread unexpectedly had wake_on_cancel set");
2887        }
2888        let old = state
2889            .get_mut(set)?
2890            .waiting
2891            .insert(guest_thread, WaitMode::Callback(self));
2892        if !old.is_none() {
2893            bail_bug!("set's waiting set already had this thread registered");
2894        }
2895        Ok(())
2896    }
2897
2898    /// Stage the specified guest call as the work item to run next, to be
2899    /// started as soon as backpressure and/or reentrance rules allow.
2900    ///
2901    /// SAFETY: The raw pointer arguments must be valid references to guest
2902    /// functions (with the appropriate signatures) when the closures staged by
2903    /// this function are called.
2904    unsafe fn stage_call<T: 'static>(
2905        self,
2906        mut store: StoreContextMut<T>,
2907        guest_thread: QualifiedThreadId,
2908        callee: SendSyncPtr<VMFuncRef>,
2909        param_count: usize,
2910        result_count: usize,
2911        async_: bool,
2912        callback: Option<SendSyncPtr<VMFuncRef>>,
2913        post_return: Option<SendSyncPtr<VMFuncRef>>,
2914        host_caller: bool,
2915    ) -> Result<()> {
2916        /// Return a closure which will call the specified function in the scope
2917        /// of the specified task.
2918        ///
2919        /// This will use `GuestTask::lower_params` to lower the parameters, but
2920        /// will not lift the result; instead, it returns a
2921        /// `[MaybeUninit<ValRaw>; MAX_FLAT_PARAMS]` from which the result, if
2922        /// any, may be lifted.  Note that an async-lifted export will have
2923        /// returned its result using the `task.return` intrinsic (or not
2924        /// returned a result at all, in the case of `task.cancel`), in which
2925        /// case the "result" of this call will either be a callback code or
2926        /// nothing.
2927        ///
2928        /// SAFETY: `callee` must be a valid `*mut VMFuncRef` at the time when
2929        /// the returned closure is called.
2930        unsafe fn make_call<T: 'static>(
2931            store: StoreContextMut<T>,
2932            guest_thread: QualifiedThreadId,
2933            callee: SendSyncPtr<VMFuncRef>,
2934            param_count: usize,
2935            result_count: usize,
2936        ) -> impl FnOnce(&mut dyn VMStore) -> Result<[MaybeUninit<ValRaw>; MAX_FLAT_PARAMS]>
2937        + Send
2938        + Sync
2939        + 'static
2940        + use<T> {
2941            let token = StoreToken::new(store);
2942            move |store: &mut dyn VMStore| {
2943                let mut storage = [MaybeUninit::uninit(); MAX_FLAT_PARAMS];
2944
2945                store
2946                    .concurrent_state_mut()?
2947                    .get_mut(guest_thread.thread)?
2948                    .state = GuestThreadState::Running;
2949                let task = store.concurrent_state_mut()?.get_mut(guest_thread.task)?;
2950                let lower = match task.lower_params.take() {
2951                    Some(l) => l,
2952                    None => bail_bug!("lower_params missing"),
2953                };
2954
2955                lower(store, &mut storage[..param_count])?;
2956
2957                let mut store = token.as_context_mut(store);
2958
2959                // SAFETY: Per the contract documented in `make_call's`
2960                // documentation, `callee` must be a valid pointer.
2961                unsafe {
2962                    crate::Func::call_unchecked_raw(
2963                        &mut store,
2964                        callee.as_non_null(),
2965                        NonNull::new(
2966                            &mut storage[..param_count.max(result_count)]
2967                                as *mut [MaybeUninit<ValRaw>] as _,
2968                        )
2969                        .unwrap(),
2970                        UncaughtException::Trap,
2971                    )?;
2972                }
2973
2974                Ok(storage)
2975            }
2976        }
2977
2978        // SAFETY: Per the contract described in this function documentation,
2979        // the `callee` pointer which `call` closes over must be valid when
2980        // called by the closure we queue below.
2981        let call = unsafe {
2982            make_call(
2983                store.as_context_mut(),
2984                guest_thread,
2985                callee,
2986                param_count,
2987                result_count,
2988            )
2989        };
2990
2991        let callee_instance = store
2992            .0
2993            .concurrent_state_mut()?
2994            .get_mut(guest_thread.task)?
2995            .instance;
2996
2997        let fun = if callback.is_some() {
2998            assert!(async_);
2999
3000            Box::new(move |store: &mut dyn VMStore| {
3001                self.add_guest_thread_to_instance_table(
3002                    guest_thread.thread,
3003                    store,
3004                    callee_instance.index,
3005                )?;
3006                let old_thread = store.set_thread(guest_thread)?;
3007                log::trace!(
3008                    "stackless call: replaced {old_thread:?} with {guest_thread:?} as current thread"
3009                );
3010
3011                store.enter_instance(callee_instance);
3012
3013                // SAFETY: See the documentation for `make_call` to review the
3014                // contract we must uphold for `call` here.
3015                //
3016                // Per the contract described in the `stage_call`
3017                // documentation, the `callee` pointer which `call` closes
3018                // over must be valid.
3019                let storage = call(store)?;
3020
3021                store.exit_instance(callee_instance)?;
3022
3023                store.set_thread(old_thread)?;
3024                let state = store.concurrent_state_mut()?;
3025                if let Some(t) = old_thread.guest() {
3026                    state.get_mut(t.thread)?.state = GuestThreadState::Running;
3027                }
3028                log::trace!("stackless call: restored {old_thread:?} as current thread");
3029
3030                // SAFETY: `wasmparser` will have validated that the callback
3031                // function returns a `i32` result.
3032                let code = unsafe { storage[0].assume_init() }.get_i32() as u32;
3033
3034                self.handle_callback_code(store, guest_thread, callee_instance.index, code)
3035            }) as Box<dyn FnOnce(&mut dyn VMStore) -> Result<()> + Send + Sync>
3036        } else {
3037            let token = StoreToken::new(store.as_context_mut());
3038            Box::new(move |store: &mut dyn VMStore| {
3039                self.add_guest_thread_to_instance_table(
3040                    guest_thread.thread,
3041                    store,
3042                    callee_instance.index,
3043                )?;
3044                let old_thread = store.set_thread(guest_thread)?;
3045                log::trace!(
3046                    "sync/async-stackful call: replaced {old_thread:?} with {guest_thread:?} as current thread",
3047                );
3048                let flags = self.id().get(store).instance_flags(callee_instance.index);
3049
3050                let callee_async_typed = store
3051                    .concurrent_state_mut()?
3052                    .get_mut(guest_thread.task)?
3053                    .async_typed;
3054
3055                // Unless this is a callback-less (i.e. stackful) async-lifted
3056                // or sync-typed export, we need to record that the instance
3057                // cannot be entered until the call returns.
3058                if !async_ && callee_async_typed {
3059                    store.enter_instance(callee_instance);
3060                }
3061
3062                if !callee_async_typed {
3063                    store.enter_sync_call(callee_instance)?;
3064                }
3065
3066                // SAFETY: See the documentation for `make_call` to review the
3067                // contract we must uphold for `call` here.
3068                //
3069                // Per the contract described in the `stage_call`
3070                // documentation, the `callee` pointer which `call` closes
3071                // over must be valid.
3072                let storage = call(store)?;
3073
3074                if !callee_async_typed {
3075                    store.exit_sync_call(callee_instance)?;
3076                }
3077
3078                if !async_ {
3079                    // This is a sync-lifted export, so now is when we lift the
3080                    // result, optionally call the post-return function, if any,
3081                    // and finally notify any current or future waiters that the
3082                    // subtask has returned.
3083
3084                    if callee_async_typed {
3085                        store.exit_instance(callee_instance)?;
3086                    }
3087
3088                    let lift = {
3089                        let state = store.concurrent_state_mut()?;
3090                        if !state.get_mut(guest_thread.task)?.result.is_none() {
3091                            bail_bug!("task has already produced a result");
3092                        }
3093
3094                        match state.get_mut(guest_thread.task)?.lift_result.take() {
3095                            Some(lift) => lift,
3096                            None => bail_bug!("lift_result field is missing"),
3097                        }
3098                    };
3099
3100                    // SAFETY: `result_count` represents the number of core Wasm
3101                    // results returned, per `wasmparser`.
3102                    let result = (lift.lift)(store, unsafe {
3103                        mem::transmute::<&[MaybeUninit<ValRaw>], &[ValRaw]>(
3104                            &storage[..result_count],
3105                        )
3106                    })?;
3107
3108                    let post_return_arg = match result_count {
3109                        0 => ValRaw::i32(0),
3110                        // SAFETY: `result_count` represents the number of
3111                        // core Wasm results returned, per `wasmparser`.
3112                        1 => unsafe { storage[0].assume_init() },
3113                        _ => unreachable!(),
3114                    };
3115
3116                    unsafe {
3117                        call_post_return(
3118                            token.as_context_mut(store),
3119                            post_return.map(|v| v.as_non_null()),
3120                            post_return_arg,
3121                            flags,
3122                        )?;
3123                    }
3124
3125                    self.task_complete(store, guest_thread.task, result, Status::Returned)?;
3126                }
3127
3128                store.set_thread(old_thread)?;
3129
3130                store
3131                    .concurrent_state_mut()?
3132                    .get_mut(guest_thread.task)?
3133                    .exited = true;
3134
3135                log::trace!(
3136                    "clean up thread; async lifted? {async_} async typed? {callee_async_typed}"
3137                );
3138
3139                if callee_async_typed {
3140                    // If we're async-typed, returning control to our caller
3141                    // won't help resolve any outstanding sync-typed call which
3142                    // might be in progress, so we may need to switch or trap
3143                    // before exiting this thread:
3144                    store.switch_or_trap_if_may_not_suspend(callee_instance)?;
3145                }
3146
3147                // This is a callback-less call, so the implicit thread has now completed
3148                store.cleanup_thread(guest_thread, callee_instance, CleanupTask::Yes)?;
3149                Ok(())
3150            })
3151        };
3152
3153        store.0.concurrent_state_mut()?.push_work_item(
3154            WorkItem::GuestCall {
3155                instance: callee_instance,
3156                call: GuestCall {
3157                    thread: guest_thread,
3158                    kind: GuestCallKind::StartImplicit(fun),
3159                },
3160            },
3161            if host_caller {
3162                Priority::High
3163            } else {
3164                Priority::Switch
3165            },
3166        )?;
3167
3168        Ok(())
3169    }
3170
3171    /// Prepare (but do not start) a guest->guest call.
3172    ///
3173    /// This is called from fused adapter code generated in
3174    /// `wasmtime_environ::fact::trampoline::Compiler`.  `start` and `return_`
3175    /// are synthesized Wasm functions which move the parameters from the caller
3176    /// to the callee and the result from the callee to the caller,
3177    /// respectively.  The adapter will call `Self::start_call` immediately
3178    /// after calling this function.
3179    ///
3180    /// SAFETY: All the pointer arguments must be valid pointers to guest
3181    /// entities (and with the expected signatures for the function references
3182    /// -- see `wasmtime_environ::fact::trampoline::Compiler` for details).
3183    unsafe fn prepare_call<T: 'static>(
3184        self,
3185        mut store: StoreContextMut<T>,
3186        start: NonNull<VMFuncRef>,
3187        return_: NonNull<VMFuncRef>,
3188        caller_instance: RuntimeComponentInstanceIndex,
3189        callee_instance: RuntimeComponentInstanceIndex,
3190        task_return_type: TypeTupleIndex,
3191        callee_async_typed: bool,
3192        memory: *mut VMMemoryDefinition,
3193        string_encoding: StringEncoding,
3194        caller_info: CallerInfo,
3195    ) -> Result<()> {
3196        enum ResultInfo {
3197            Heap { results: u32 },
3198            Stack { result_count: u32 },
3199        }
3200
3201        let result_info = match &caller_info {
3202            CallerInfo::Async {
3203                has_result: true,
3204                params,
3205            } => ResultInfo::Heap {
3206                results: match params.last() {
3207                    Some(r) => r.get_u32(),
3208                    None => bail_bug!("retptr missing"),
3209                },
3210            },
3211            CallerInfo::Async {
3212                has_result: false, ..
3213            } => ResultInfo::Stack { result_count: 0 },
3214            CallerInfo::Sync {
3215                result_count,
3216                params,
3217            } if *result_count > u32::try_from(MAX_FLAT_RESULTS)? => ResultInfo::Heap {
3218                results: match params.last() {
3219                    Some(r) => r.get_u32(),
3220                    None => bail_bug!("arg ptr missing"),
3221                },
3222            },
3223            CallerInfo::Sync { result_count, .. } => ResultInfo::Stack {
3224                result_count: *result_count,
3225            },
3226        };
3227
3228        let sync_caller = matches!(caller_info, CallerInfo::Sync { .. });
3229
3230        // Create a new guest task for the call, closing over the `start` and
3231        // `return_` functions to lift the parameters and lower the result,
3232        // respectively.
3233        let start = SendSyncPtr::new(start);
3234        let return_ = SendSyncPtr::new(return_);
3235        let token = StoreToken::new(store.as_context_mut());
3236        let old_thread = store.0.current_guest_thread()?;
3237
3238        let state = store.0.concurrent_state_mut()?;
3239
3240        debug_assert_eq!(
3241            state.get_mut(old_thread.task)?.instance,
3242            self.runtime_instance(caller_instance)
3243        );
3244
3245        let guest_thread = GuestTask::new(
3246            state,
3247            Box::new(move |store, dst| {
3248                let mut store = token.as_context_mut(store);
3249                assert!(dst.len() <= MAX_FLAT_PARAMS);
3250                // The `+ 1` here accounts for the return pointer, if any:
3251                let mut src = [MaybeUninit::uninit(); MAX_FLAT_PARAMS + 1];
3252                let count = match caller_info {
3253                    // Async callers, if they have a result, use the last
3254                    // parameter as a return pointer so chop that off if
3255                    // relevant here.
3256                    CallerInfo::Async { params, has_result } => {
3257                        let params = &params[..params.len() - usize::from(has_result)];
3258                        for (param, src) in params.iter().zip(&mut src) {
3259                            src.write(*param);
3260                        }
3261                        params.len()
3262                    }
3263
3264                    // Sync callers forward everything directly.
3265                    CallerInfo::Sync { params, .. } => {
3266                        for (param, src) in params.iter().zip(&mut src) {
3267                            src.write(*param);
3268                        }
3269                        params.len()
3270                    }
3271                };
3272                // SAFETY: `start` is a valid `*mut VMFuncRef` from
3273                // `wasmtime-cranelift`-generated fused adapter code.  Based on
3274                // how it was constructed (see
3275                // `wasmtime_environ::fact::trampoline::Compiler::compile_async_start_adapter`
3276                // for details) we know it takes count parameters and returns
3277                // `dst.len()` results.
3278                unsafe {
3279                    crate::Func::call_unchecked_raw(
3280                        &mut store,
3281                        start.as_non_null(),
3282                        NonNull::new(
3283                            &mut src[..count.max(dst.len())] as *mut [MaybeUninit<ValRaw>] as _,
3284                        )
3285                        .unwrap(),
3286                        UncaughtException::Trap,
3287                    )?;
3288                }
3289                dst.copy_from_slice(&src[..dst.len()]);
3290                let task = store.0.current_guest_thread()?.task;
3291                let state = store.0.concurrent_state_mut()?;
3292                Waitable::Guest(task).set_event(
3293                    state,
3294                    Some(Event::Subtask {
3295                        status: Status::Started,
3296                    }),
3297                )?;
3298                Ok(())
3299            }),
3300            LiftResult {
3301                lift: Box::new(move |store, src| {
3302                    // SAFETY: See comment in closure passed as `lower_params`
3303                    // parameter above.
3304                    let mut store = token.as_context_mut(store);
3305                    let mut my_src = src.to_owned(); // TODO: use stack to avoid allocation?
3306                    if let ResultInfo::Heap { results } = &result_info {
3307                        my_src.push(ValRaw::u32(*results));
3308                    }
3309
3310                    // SAFETY: `return_` is a valid `*mut VMFuncRef` from
3311                    // `wasmtime-cranelift`-generated fused adapter code.  Based
3312                    // on how it was constructed (see
3313                    // `wasmtime_environ::fact::trampoline::Compiler::compile_async_return_adapter`
3314                    // for details) we know it takes `src.len()` parameters and
3315                    // returns up to 1 result.
3316                    unsafe {
3317                        crate::Func::call_unchecked_raw(
3318                            &mut store,
3319                            return_.as_non_null(),
3320                            my_src.as_mut_slice().into(),
3321                            UncaughtException::Trap,
3322                        )?;
3323                    }
3324
3325                    let thread = store.0.current_guest_thread()?;
3326                    let state = store.0.concurrent_state_mut()?;
3327                    if sync_caller {
3328                        state.get_mut(thread.task)?.sync_result = SyncResult::Produced(
3329                            if let ResultInfo::Stack { result_count } = &result_info {
3330                                match result_count {
3331                                    0 => None,
3332                                    1 => Some(my_src[0]),
3333                                    _ => unreachable!(),
3334                                }
3335                            } else {
3336                                None
3337                            },
3338                        );
3339                    }
3340                    Ok(Box::new(DummyResult) as Box<dyn Any + Send + Sync>)
3341                }),
3342                ty: task_return_type,
3343                memory: NonNull::new(memory).map(SendSyncPtr::new),
3344                string_encoding,
3345            },
3346            Caller::Guest { thread: old_thread },
3347            None,
3348            self.runtime_instance(callee_instance),
3349            callee_async_typed,
3350            // We don't know whether the callee export was lifted sync or async
3351            // yet, but we'll update this in `start_call`:
3352            false,
3353        )?;
3354
3355        // Make the new thread the current one so that `Self::start_call` knows
3356        // which one to start.
3357        store.0.set_thread(guest_thread)?;
3358        log::trace!("pushed {guest_thread:?} as current thread; old thread was {old_thread:?}");
3359
3360        Ok(())
3361    }
3362
3363    /// Call the specified callback function for an async-lifted export.
3364    ///
3365    /// SAFETY: `function` must be a valid reference to a guest function of the
3366    /// correct signature for a callback.
3367    unsafe fn call_callback<T>(
3368        self,
3369        mut store: StoreContextMut<T>,
3370        function: SendSyncPtr<VMFuncRef>,
3371        event: Event,
3372        handle: u32,
3373    ) -> Result<u32> {
3374        let (ordinal, result) = event.parts();
3375        let params = &mut [
3376            ValRaw::u32(ordinal),
3377            ValRaw::u32(handle),
3378            ValRaw::u32(result),
3379        ];
3380        // SAFETY: `func` is a valid `*mut VMFuncRef` from either
3381        // `wasmtime-cranelift`-generated fused adapter code or
3382        // `component::Options`.  Per `wasmparser` callback signature
3383        // validation, we know it takes three parameters and returns one.
3384        unsafe {
3385            crate::Func::call_unchecked_raw(
3386                &mut store,
3387                function.as_non_null(),
3388                params.as_mut_slice().into(),
3389                UncaughtException::Trap,
3390            )?;
3391        }
3392        Ok(params[0].get_u32())
3393    }
3394
3395    /// Start a guest->guest call previously prepared using
3396    /// `Self::prepare_call`.
3397    ///
3398    /// This is called from fused adapter code generated in
3399    /// `wasmtime_environ::fact::trampoline::Compiler`.  The adapter will call
3400    /// this function immediately after calling `Self::prepare_call`.
3401    ///
3402    /// If `flags` contains `START_FLAG_ASYNC_CALLER` then the caller used an
3403    /// async-lowered import and the call's packed status is written to
3404    /// `storage[0]`. Otherwise this blocks until the callee has produced its
3405    /// result, if any, which is written to `storage[0]`.
3406    ///
3407    /// SAFETY: The `*mut VMFuncRef` arguments must be valid pointers to guest
3408    /// functions with the appropriate signatures for the current guest task.
3409    /// If this is a call to an async-lowered import, the actual call may be
3410    /// deferred and run after this function returns, in which case the pointer
3411    /// arguments must also be valid when the call happens.
3412    unsafe fn start_call<T: 'static>(
3413        self,
3414        mut store: StoreContextMut<T>,
3415        callback: *mut VMFuncRef,
3416        post_return: *mut VMFuncRef,
3417        callee: NonNull<VMFuncRef>,
3418        param_count: u32,
3419        result_count: u32,
3420        flags: u32,
3421        storage: &mut [MaybeUninit<ValRaw>],
3422    ) -> Result<()> {
3423        let token = StoreToken::new(store.as_context_mut());
3424        let async_caller = (flags & START_FLAG_ASYNC_CALLER) != 0;
3425        let guest_thread = store.0.current_guest_thread()?;
3426        let state = store.0.concurrent_state_mut()?;
3427
3428        if !state.event_loop_running {
3429            bail_bug!("Instance::start_call called without a running event loop");
3430        }
3431
3432        let callee = SendSyncPtr::new(callee);
3433        let param_count = usize::try_from(param_count)?;
3434        assert!(param_count <= MAX_FLAT_PARAMS);
3435        let result_count = usize::try_from(result_count)?;
3436        assert!(result_count <= MAX_FLAT_RESULTS);
3437
3438        let task = state.get_mut(guest_thread.task)?;
3439        let callee_async_typed = task.async_typed;
3440        let callee_instance = task.instance;
3441
3442        task.async_lifted = (flags & START_FLAG_ASYNC_CALLEE) != 0;
3443
3444        if let Some(callback) = NonNull::new(callback) {
3445            // We're calling an async-lifted export with a callback, so store
3446            // the callback and related context as part of the task so we can
3447            // call it later when needed.
3448            let callback = SendSyncPtr::new(callback);
3449            task.callback = Some(Box::new(move |store, event, handle| {
3450                let store = token.as_context_mut(store);
3451                unsafe { self.call_callback::<T>(store, callback, event, handle) }
3452            }));
3453        }
3454
3455        let Caller::Guest { thread: caller } = &task.caller else {
3456            // As of this writing, `start_call` is only used for guest->guest
3457            // calls.
3458            bail_bug!("start_call unexpectedly invoked for host->guest call");
3459        };
3460        let caller = *caller;
3461        let caller_instance = state.get_mut(caller.task)?.instance;
3462
3463        // Stage the call as the work item to run next, modulo backpressure, etc.
3464        unsafe {
3465            self.stage_call(
3466                store.as_context_mut(),
3467                guest_thread,
3468                callee,
3469                param_count,
3470                result_count,
3471                (flags & START_FLAG_ASYNC_CALLEE) != 0,
3472                NonNull::new(callback).map(SendSyncPtr::new),
3473                NonNull::new(post_return).map(SendSyncPtr::new),
3474                false,
3475            )?;
3476        }
3477
3478        let old_do_not_suspend = if callee_async_typed {
3479            // If we're starting an async-typed call, it is permitted to suspend
3480            // since that will just return control back to the caller, which may
3481            // be able to avoid blocking if needed.
3482            //
3483            // We'll restore the old value after the call either returns or
3484            // suspends.
3485            let state = store.0.instance_state(callee_instance).concurrent_state();
3486            let old_do_not_suspend = state.do_not_suspend;
3487            state.do_not_suspend = false;
3488            Some(old_do_not_suspend)
3489        } else {
3490            None
3491        };
3492
3493        let state = store.0.concurrent_state_mut()?;
3494
3495        // Use the caller's `GuestThread::sync_call_set` to register interest in
3496        // the subtask...
3497        let guest_waitable = Waitable::Guest(guest_thread.task);
3498        let old_set = guest_waitable.common(state)?.set;
3499        let set = state.get_mut(caller.thread)?.sync_call_set;
3500        guest_waitable.join(state, Some(set))?;
3501
3502        store.0.set_thread(CurrentThread::None)?;
3503
3504        // ... and suspend this fiber temporarily while we wait for it to start.
3505        //
3506        // Note that we _could_ call the callee directly using the current fiber
3507        // rather than suspend this one, but that would make reasoning about the
3508        // event loop more complicated and is probably only worth doing if
3509        // there's a measurable performance benefit.  In addition, it would mean
3510        // blocking the caller if the callee calls a blocking sync-lowered
3511        // import, and as of this writing the spec says we must not do that.
3512        //
3513        // Alternatively, the fused adapter code could be modified to call the
3514        // callee directly without calling a host-provided intrinsic at all (in
3515        // which case it would need to do its own, inline backpressure checks,
3516        // etc.).  Again, we'd want to see a measurable performance benefit
3517        // before committing to such an optimization.  And again, we'd need to
3518        // update the spec to allow that.
3519        let mut yielded = false;
3520        let (status, waitable) = loop {
3521            store.0.suspend(if yielded {
3522                SuspendReason::Waiting {
3523                    set,
3524                    thread: caller,
3525                }
3526            } else {
3527                yielded = true;
3528                SuspendReason::YieldingToSubtask { thread: caller }
3529            })?;
3530
3531            if let Some(old_do_not_suspend) = old_do_not_suspend {
3532                store
3533                    .0
3534                    .instance_state(callee_instance)
3535                    .concurrent_state()
3536                    .do_not_suspend = old_do_not_suspend;
3537            }
3538
3539            let state = store.0.concurrent_state_mut()?;
3540
3541            log::trace!("taking event for {:?}", guest_thread.task);
3542            let event = guest_waitable.take_event(state)?;
3543            let Some(Event::Subtask { status }) = event else {
3544                bail_bug!("subtasks should only get subtask events, got {event:?}")
3545            };
3546
3547            log::trace!("status {status:?} for {:?}", guest_thread.task);
3548
3549            if status == Status::Returned {
3550                // It returned, so we can stop waiting.
3551                break (status, None);
3552            } else if async_caller {
3553                // It hasn't returned yet, but the caller is calling via an
3554                // async-lowered import, so we generate a handle for the task
3555                // waitable and return the status.
3556                let handle = store
3557                    .0
3558                    .instance_state(caller_instance)
3559                    .handle_table()
3560                    .subtask_insert_guest(guest_thread.task.rep())?;
3561                store
3562                    .0
3563                    .concurrent_state_mut()?
3564                    .get_mut(guest_thread.task)?
3565                    .common
3566                    .handle = Some(handle);
3567                break (status, Some(handle));
3568            } else {
3569                // The callee hasn't returned yet, and the caller is calling via
3570                // a sync-lowered import, so we loop and keep waiting until the
3571                // callee returns.
3572                store.0.switch_or_trap_if_may_not_suspend(caller_instance)?;
3573            }
3574        };
3575
3576        guest_waitable.join(store.0.concurrent_state_mut()?, old_set)?;
3577
3578        // Reset the current thread to point to the caller as it resumes control.
3579        store.0.set_thread(caller)?;
3580        store
3581            .0
3582            .concurrent_state_mut()?
3583            .get_mut(caller.thread)?
3584            .state = GuestThreadState::Running;
3585        log::trace!("popped current thread {guest_thread:?}; new thread is {caller:?}");
3586
3587        if async_caller {
3588            let Some(slot) = storage.first_mut() else {
3589                bail_bug!("no storage for async call status");
3590            };
3591            *slot = MaybeUninit::new(ValRaw::u32(status.pack(waitable)));
3592        } else {
3593            // The caller used a sync-lowered import, in which case the result,
3594            // if any, has been stashed in `GuestTask::sync_result`.
3595            let state = store.0.concurrent_state_mut()?;
3596            let task = state.get_mut(guest_thread.task)?;
3597            if let Some(result) = task.sync_result.take()? {
3598                if let Some(result) = result {
3599                    storage[0] = MaybeUninit::new(result);
3600                }
3601
3602                if task.exited && task.ready_to_delete() {
3603                    Waitable::Guest(guest_thread.task).delete_from(store.0)?;
3604                }
3605            }
3606        }
3607
3608        Ok(())
3609    }
3610
3611    /// Poll the specified future once on behalf of a guest->host call using an
3612    /// async-lowered import.
3613    ///
3614    /// If it returns `Ready`, return `Ok(None)`.  Otherwise, if it returns
3615    /// `Pending`, add it to the set of futures to be polled as part of this
3616    /// instance's event loop until it completes, and then return
3617    /// `Ok(Some(handle))` where `handle` is the waitable handle to return.
3618    ///
3619    /// Whether the future returns `Ready` immediately or later, the `lower`
3620    /// function will be used to lower the result, if any, into the guest caller's
3621    /// stack and linear memory. The `lower` function is invoked with the
3622    /// `Option<R>` param `None` if the future is cancelled. The
3623    /// `Option<TableId<HostTask>>` is passed as `Some` if the host task was
3624    /// materialized during execution and allows `lower` to delete the task if
3625    /// needed.
3626    pub(crate) fn first_poll<T: 'static, R: Send + 'static>(
3627        self,
3628        mut store: StoreContextMut<'_, T>,
3629        host_task: EnteredHostTask,
3630        result_may_require_realloc: bool,
3631        future: impl Future<Output = Result<R>> + Send + 'static,
3632        lower: impl FnOnce(StoreContextMut<T>, Option<R>, bool, Option<TableId<HostTask>>) -> Result<()>
3633        + Send
3634        + 'static,
3635    ) -> Result<u32> {
3636        let token = StoreToken::new(store.as_context_mut());
3637
3638        // Create an abortable future which hooks calls to poll and manages call
3639        // context state for the future.
3640        let (join_handle, future) = JoinHandle::run(future);
3641        let mut future = Box::pin(future);
3642
3643        // Finally, poll the future.  We can use a dummy `Waker` here because
3644        // we'll add the future to `ConcurrentState::futures` and poll it
3645        // automatically from the event loop if it doesn't complete immediately
3646        // here.
3647        let poll = tls::set(store.0, || {
3648            future
3649                .as_mut()
3650                .poll(&mut Context::from_waker(&Waker::noop()))
3651        });
3652
3653        match poll {
3654            // It finished immediately; lower the result and delete the task.
3655            Poll::Ready(result) => {
3656                let result = result.transpose()?;
3657                // Check if the host task was materialized so that it can be
3658                // deleted in `lower`.
3659                let task = store.0.current_materialized_host_task()?;
3660                lower(store.as_context_mut(), result, true, task)?;
3661                return Ok(Status::Returned.pack(None));
3662            }
3663
3664            // Future isn't ready yet, so fall through.
3665            Poll::Pending => {}
3666        }
3667
3668        // The future will outlive this call frame, so materialize the deferred
3669        // host task and attach its cancellation handle before publishing it to
3670        // the event loop.
3671        let Some(task) = store.0.materialize_host_task_id()? else {
3672            bail_bug!("current thread is not a host thread")
3673        };
3674        {
3675            let state = &mut store.0.concurrent_state_mut()?.get_mut(task)?.state;
3676            assert!(matches!(state, HostTaskState::CalleeStarted));
3677            *state = HostTaskState::CalleeRunning(join_handle);
3678        }
3679
3680        // It hasn't finished yet; add the future to
3681        // `ConcurrentState::futures` so it will be polled by the event
3682        // loop and allocate a waitable handle to return to the guest.
3683
3684        // Wrap the future in a closure responsible for lowering the result into
3685        // the guest's stack and memory, as well as notifying any waiters that
3686        // the task returned.
3687        let future = Box::pin(async move {
3688            let result = match run_with_host_task_set(task, future).await? {
3689                Some(result) => Some(result?),
3690                None => None,
3691            };
3692            let on_complete = move |store: &mut dyn VMStore| {
3693                // Restore the `current_thread` to be the host so `lower` knows
3694                // how to manipulate borrows and knows which scope of borrows
3695                // to check.
3696                let mut store = token.as_context_mut(store);
3697                let old = store.0.set_thread(task)?;
3698
3699                let status = if result.is_some() {
3700                    Status::Returned
3701                } else {
3702                    Status::ReturnCancelled
3703                };
3704
3705                lower(store.as_context_mut(), result, false, Some(task))?;
3706                let state = store.0.concurrent_state_mut()?;
3707                match &mut state.get_mut(task)?.state {
3708                    // Cancellation is only complete once lowering has finished
3709                    // and the terminal event can be published.
3710                    pending @ HostTaskState::CalleeCancelling => {
3711                        *pending = HostTaskState::CalleeDone { cancelled: true };
3712                    }
3713
3714                    // Otherwise transition this task to the done state.
3715                    other => *other = HostTaskState::CalleeDone { cancelled: false },
3716                }
3717                Waitable::Host(task).set_event(state, Some(Event::Subtask { status }))?;
3718
3719                store.0.set_thread(old)?;
3720                Ok(())
3721            };
3722
3723            tls::get(move |store| {
3724                if result_may_require_realloc {
3725                    // Lowering may call the guest's realloc function, so run
3726                    // it on a worker fiber. This is required because the guest
3727                    // might try to suspend the fiber which is unsound when it's
3728                    // not actually on a fiber.
3729                    store
3730                        .concurrent_state_mut()?
3731                        .push_high_priority(WorkItem::WorkerFunction(AlwaysMut::new(Box::new(
3732                            on_complete,
3733                        ))));
3734                    Ok(())
3735                } else {
3736                    // No guest code can be called while lowering this result,
3737                    // so avoid scheduling a worker fiber.
3738                    on_complete(store)
3739                }
3740            })
3741        });
3742
3743        // Make this task visible to the guest and then record what it
3744        // was made visible as.
3745        let caller = match host_task {
3746            Some(caller) => caller,
3747            None => bail_bug!("host task wasn't created but should have been"),
3748        };
3749        let state = store.0.concurrent_state_mut()?;
3750        state.push_future(future);
3751        let instance = state.get_mut(caller.task)?.instance;
3752        let handle = store
3753            .0
3754            .instance_state(instance)
3755            .handle_table()
3756            .subtask_insert_host(task.rep())?;
3757        store.0.concurrent_state_mut()?.get_mut(task)?.common.handle = Some(handle);
3758        log::trace!("assign {task:?} handle {handle} for {caller:?} instance {instance:?}");
3759
3760        // Restore the currently running thread to this host task's
3761        // caller. Note that the host task isn't deallocated as it's
3762        // within the store and will get deallocated later.
3763        store.0.set_thread(caller)?;
3764        Ok(Status::Started.pack(Some(handle)))
3765    }
3766
3767    /// Implements the `task.return` intrinsic, lifting the result for the
3768    /// current guest task.
3769    pub(crate) fn task_return(
3770        self,
3771        store: &mut dyn VMStore,
3772        ty: TypeTupleIndex,
3773        options: OptionsIndex,
3774        storage: &[ValRaw],
3775    ) -> Result<()> {
3776        let guest_thread = store.current_guest_thread()?;
3777        let state = store.concurrent_state_mut()?;
3778        if !state.get_mut(guest_thread.task)?.async_lifted {
3779            bail!(Trap::TaskReturnOrCancelSyncLifted);
3780        }
3781        let lift = state
3782            .get_mut(guest_thread.task)?
3783            .lift_result
3784            .take()
3785            .ok_or_else(|| Trap::TaskCancelOrReturnTwice)?;
3786        if !state.get_mut(guest_thread.task)?.result.is_none() {
3787            bail_bug!("task result unexpectedly already set");
3788        }
3789
3790        let CanonicalOptions {
3791            string_encoding,
3792            data_model,
3793            ..
3794        } = &self.id().get(store).component().env_component().options[options];
3795
3796        let invalid = ty != lift.ty
3797            || string_encoding != &lift.string_encoding
3798            || match data_model {
3799                CanonicalOptionsDataModel::LinearMemory(opts) => match opts.memory {
3800                    Some(memory) => {
3801                        let expected = lift.memory.map(|v| v.as_ptr()).unwrap_or(ptr::null_mut());
3802                        let actual = self.id().get(store).runtime_memory(memory);
3803                        expected != actual.as_ptr()
3804                    }
3805                    // Memory not specified, meaning it didn't need to be
3806                    // specified per validation, so not invalid.
3807                    None => false,
3808                },
3809                // Always invalid as this isn't supported.
3810                CanonicalOptionsDataModel::Gc { .. } => true,
3811            };
3812
3813        if invalid {
3814            bail!(Trap::TaskReturnInvalid);
3815        }
3816
3817        log::trace!("task.return for {guest_thread:?}");
3818
3819        let result = (lift.lift)(store, storage)?;
3820        self.task_complete(store, guest_thread.task, result, Status::Returned)
3821    }
3822
3823    /// Implements the `task.cancel` intrinsic.
3824    pub(crate) fn task_cancel(self, store: &mut StoreOpaque) -> Result<()> {
3825        let guest_thread = store.current_guest_thread()?;
3826        let state = store.concurrent_state_mut()?;
3827        let task = state.get_mut(guest_thread.task)?;
3828        if !task.async_lifted {
3829            bail!(Trap::TaskReturnOrCancelSyncLifted);
3830        }
3831        if !task.cancel_request_delivered {
3832            bail!(Trap::TaskCancelNotCancelled);
3833        }
3834        _ = task
3835            .lift_result
3836            .take()
3837            .ok_or_else(|| Trap::TaskCancelOrReturnTwice)?;
3838
3839        if !task.result.is_none() {
3840            bail_bug!("task result should not bet set yet");
3841        }
3842
3843        log::trace!("task.cancel for {guest_thread:?}");
3844
3845        self.task_complete(
3846            store,
3847            guest_thread.task,
3848            Box::new(DummyResult),
3849            Status::ReturnCancelled,
3850        )
3851    }
3852
3853    /// Complete the specified guest task (i.e. indicate that it has either
3854    /// returned a (possibly empty) result or cancelled itself).
3855    ///
3856    /// This will return any resource borrows and notify any current or future
3857    /// waiters that the task has completed.
3858    fn task_complete(
3859        self,
3860        store: &mut StoreOpaque,
3861        guest_task: TableId<GuestTask>,
3862        result: Box<dyn Any + Send + Sync>,
3863        status: Status,
3864    ) -> Result<()> {
3865        store
3866            .component_resource_tables(Some(self))?
3867            .validate_scope_exit()?;
3868
3869        let state = store.concurrent_state_mut()?;
3870        let task = state.get_mut(guest_task)?;
3871
3872        // A cancellation request is only delivered to a task which hasn't yet
3873        // resolved, so drop any pending, undelivered `Event::Cancelled` now
3874        // that this task has resolved.
3875        task.event = None;
3876
3877        if let Caller::Host { tx, .. } = &mut task.caller {
3878            if let Some(tx) = tx.take() {
3879                _ = tx.send(result);
3880            }
3881        } else {
3882            task.result = Some(result);
3883            Waitable::Guest(guest_task).set_event(state, Some(Event::Subtask { status }))?;
3884        }
3885
3886        Ok(())
3887    }
3888
3889    /// Implements the `waitable-set.new` intrinsic.
3890    pub(crate) fn waitable_set_new(
3891        self,
3892        store: &mut StoreOpaque,
3893        caller_instance: RuntimeComponentInstanceIndex,
3894    ) -> Result<u32> {
3895        let set = store.concurrent_state_mut()?.push(WaitableSet::default())?;
3896        let handle = store
3897            .instance_state(self.runtime_instance(caller_instance))
3898            .handle_table()
3899            .waitable_set_insert(set.rep())?;
3900        log::trace!("new waitable set {set:?} (handle {handle})");
3901        Ok(handle)
3902    }
3903
3904    /// Implements the `waitable-set.drop` intrinsic.
3905    pub(crate) fn waitable_set_drop(
3906        self,
3907        store: &mut StoreOpaque,
3908        caller_instance: RuntimeComponentInstanceIndex,
3909        set: u32,
3910    ) -> Result<()> {
3911        let rep = store
3912            .instance_state(self.runtime_instance(caller_instance))
3913            .handle_table()
3914            .waitable_set_remove(set)?;
3915
3916        log::trace!("drop waitable set {rep} (handle {set})");
3917
3918        // Note that we're careful to check for waiters _before_ deleting the
3919        // set to avoid dropping any waiters in `WaitMode::Fiber(_)`, which
3920        // would panic.  See `drop-waitable-set-with-waiters.wast` for details.
3921        let set = store
3922            .concurrent_state_mut()?
3923            .get_mut(TableId::<WaitableSet>::new(rep))?;
3924        if set.num_waiting > 0 {
3925            bail!(Trap::WaitableSetDropHasWaiters);
3926        }
3927
3928        store
3929            .concurrent_state_mut()?
3930            .delete(TableId::<WaitableSet>::new(rep))?;
3931
3932        Ok(())
3933    }
3934
3935    /// Implements the `waitable.join` intrinsic.
3936    pub(crate) fn waitable_join(
3937        self,
3938        store: &mut StoreOpaque,
3939        caller_instance: RuntimeComponentInstanceIndex,
3940        waitable_handle: u32,
3941        set_handle: u32,
3942    ) -> Result<()> {
3943        let mut instance = self.id().get_mut(store);
3944        let waitable =
3945            Waitable::from_instance(instance.as_mut(), caller_instance, waitable_handle)?;
3946
3947        let set = if set_handle == 0 {
3948            None
3949        } else {
3950            let set = instance.instance_states().0[caller_instance]
3951                .handle_table()
3952                .waitable_set_rep(set_handle)?;
3953
3954            let state = store.concurrent_state_mut()?;
3955            if let Some(old) = waitable.common(state)?.set
3956                && state.get_mut(old)?.is_sync_call_set
3957            {
3958                bail!(Trap::WaitableSyncAndAsync);
3959            }
3960
3961            Some(TableId::<WaitableSet>::new(set))
3962        };
3963
3964        log::trace!(
3965            "waitable {waitable:?} (handle {waitable_handle}) join set {set:?} (handle {set_handle})",
3966        );
3967
3968        waitable.join(store.concurrent_state_mut()?, set)
3969    }
3970
3971    /// Implements the `subtask.drop` intrinsic.
3972    pub(crate) fn subtask_drop(
3973        self,
3974        store: &mut StoreOpaque,
3975        caller_instance: RuntimeComponentInstanceIndex,
3976        task_id: u32,
3977    ) -> Result<()> {
3978        self.waitable_join(store, caller_instance, task_id, 0)?;
3979
3980        let (rep, is_host) = store
3981            .instance_state(self.runtime_instance(caller_instance))
3982            .handle_table()
3983            .subtask_remove(task_id)?;
3984
3985        let concurrent_state = store.concurrent_state_mut()?;
3986        let (waitable, delete) = if is_host {
3987            let id = TableId::<HostTask>::new(rep);
3988            let task = concurrent_state.get_mut(id)?;
3989            match &task.state {
3990                HostTaskState::CalleeRunning(_) | HostTaskState::CalleeCancelling => {
3991                    bail!(Trap::SubtaskDropNotResolved)
3992                }
3993                HostTaskState::CalleeDone { .. } => {}
3994                HostTaskState::CalleeStarted | HostTaskState::CalleeFinished(_) => {
3995                    bail_bug!("invalid state for callee in `subtask.drop`")
3996                }
3997            }
3998
3999            (Waitable::Host(id), true)
4000        } else {
4001            let id = TableId::<GuestTask>::new(rep);
4002            let task = concurrent_state.get_mut(id)?;
4003            if task.lift_result.is_some() {
4004                bail!(Trap::SubtaskDropNotResolved);
4005            }
4006            (
4007                Waitable::Guest(id),
4008                concurrent_state.get_mut(id)?.ready_to_delete(),
4009            )
4010        };
4011
4012        waitable.common(concurrent_state)?.handle = None;
4013
4014        // If this subtask has an event that means that the terminal status of
4015        // this subtask wasn't yet received so it can't be dropped yet.
4016        if waitable.take_event(concurrent_state)?.is_some() {
4017            bail!(Trap::SubtaskDropNotResolved);
4018        }
4019
4020        if delete {
4021            waitable.delete_from(store)?;
4022        }
4023
4024        log::trace!("subtask_drop {waitable:?} (handle {task_id})");
4025        Ok(())
4026    }
4027
4028    /// Implements the `waitable-set.wait` intrinsic.
4029    pub(crate) fn waitable_set_wait(
4030        self,
4031        store: &mut StoreOpaque,
4032        options: OptionsIndex,
4033        set: u32,
4034        payload: u32,
4035    ) -> Result<u32> {
4036        let &CanonicalOptions {
4037            instance: caller_instance,
4038            ..
4039        } = &self.id().get(store).component().env_component().options[options];
4040        let caller = self.runtime_instance(caller_instance);
4041        let rep = store
4042            .instance_state(self.runtime_instance(caller_instance))
4043            .handle_table()
4044            .waitable_set_rep(set)?;
4045
4046        self.waitable_check(
4047            store,
4048            caller,
4049            WaitableCheck::Wait,
4050            WaitableCheckParams {
4051                set: TableId::new(rep),
4052                options,
4053                payload,
4054            },
4055        )
4056    }
4057
4058    /// Implements the `waitable-set.poll` intrinsic.
4059    pub(crate) fn waitable_set_poll(
4060        self,
4061        store: &mut StoreOpaque,
4062        options: OptionsIndex,
4063        set: u32,
4064        payload: u32,
4065    ) -> Result<u32> {
4066        let &CanonicalOptions {
4067            instance: caller_instance,
4068            ..
4069        } = &self.id().get(store).component().env_component().options[options];
4070        let caller = self.runtime_instance(caller_instance);
4071        let rep = store
4072            .instance_state(caller)
4073            .handle_table()
4074            .waitable_set_rep(set)?;
4075
4076        self.waitable_check(
4077            store,
4078            caller,
4079            WaitableCheck::Poll,
4080            WaitableCheckParams {
4081                set: TableId::new(rep),
4082                options,
4083                payload,
4084            },
4085        )
4086    }
4087
4088    /// Implements the `thread.index` intrinsic.
4089    pub(crate) fn thread_index(&self, store: &mut dyn VMStore) -> Result<u32> {
4090        let thread_id = store.current_guest_thread()?.thread;
4091        match store
4092            .concurrent_state_mut()?
4093            .get_mut(thread_id)?
4094            .instance_rep
4095        {
4096            Some(r) => Ok(r),
4097            None => bail_bug!("thread should have instance_rep by now"),
4098        }
4099    }
4100
4101    /// Implements the `thread.new-indirect` intrinsic.
4102    pub(crate) fn thread_new_indirect<T: 'static>(
4103        self,
4104        mut store: StoreContextMut<T>,
4105        runtime_instance: RuntimeComponentInstanceIndex,
4106        start_func_ty: ModuleInternedTypeIndex,
4107        start_func_table_idx: RuntimeTableIndex,
4108        start_func_idx: u32,
4109        context: i32,
4110    ) -> Result<u32> {
4111        log::trace!("creating new thread");
4112
4113        let (instance, registry) = self.id().get_mut_and_registry(store.0);
4114        let Some(start_func_ty) = instance.component().signatures().shared_type(start_func_ty)
4115        else {
4116            bail_bug!("thread.new-indirect start function type should be registered");
4117        };
4118        let callee = instance
4119            .index_runtime_func_table(registry, start_func_table_idx, start_func_idx as u64)?
4120            .ok_or_else(|| Trap::ThreadNewIndirectUninitialized)?;
4121        if callee.type_index(store.0) != start_func_ty {
4122            bail!(Trap::ThreadNewIndirectInvalidType);
4123        }
4124
4125        let token = StoreToken::new(store.as_context_mut());
4126        let start_func = Box::new(
4127            move |store: &mut dyn VMStore, guest_thread: QualifiedThreadId| -> Result<()> {
4128                let old_thread = store.set_thread(guest_thread)?;
4129                log::trace!(
4130                    "thread start: replaced {old_thread:?} with {guest_thread:?} as current thread"
4131                );
4132
4133                let mut store = token.as_context_mut(store);
4134                let mut params = [ValRaw::i32(context)];
4135                // Use call_unchecked rather than call or call_async, as we don't want to run the function
4136                // on a separate fiber if we're running in an async store.
4137                unsafe { callee.call_unchecked(store.as_context_mut(), &mut params)? };
4138
4139                store.0.set_thread(old_thread)?;
4140
4141                let runtime_instance = self.runtime_instance(runtime_instance);
4142
4143                // We're not returning to any caller, so we may need to switch
4144                // or trap if the instance has an outstanding sync-typed call:
4145                store
4146                    .0
4147                    .switch_or_trap_if_may_not_suspend(runtime_instance)?;
4148
4149                store
4150                    .0
4151                    .cleanup_thread(guest_thread, runtime_instance, CleanupTask::Yes)?;
4152
4153                log::trace!("explicit thread {guest_thread:?} completed");
4154                let state = store.0.concurrent_state_mut()?;
4155                if let Some(t) = old_thread.guest() {
4156                    state.get_mut(t.thread)?.state = GuestThreadState::Running;
4157                }
4158                log::trace!("thread start: restored {old_thread:?} as current thread");
4159
4160                Ok(())
4161            },
4162        );
4163
4164        let current_thread = store.0.current_guest_thread()?;
4165        let state = store.0.concurrent_state_mut()?;
4166        let parent_task = current_thread.task;
4167
4168        let new_thread = GuestThread::new_explicit(state, parent_task, start_func)?;
4169        let thread_id = state.push(new_thread)?;
4170        state.get_mut(parent_task)?.threads.insert(thread_id);
4171
4172        log::trace!("new thread with id {thread_id:?} created");
4173
4174        self.add_guest_thread_to_instance_table(thread_id, store.0, runtime_instance)
4175    }
4176
4177    pub(crate) fn resume_thread(
4178        self,
4179        store: &mut StoreOpaque,
4180        runtime_instance: RuntimeComponentInstanceIndex,
4181        thread_idx: u32,
4182        how: ResumeThread,
4183    ) -> Result<bool> {
4184        let thread_id =
4185            GuestThread::from_instance(self.id().get_mut(store), runtime_instance, thread_idx)?;
4186        let state = store.concurrent_state_mut()?;
4187        let guest_thread = QualifiedThreadId::qualify(state, thread_id)?;
4188
4189        if store.current_guest_thread()? == guest_thread {
4190            bail!(Trap::CannotResumeThread);
4191        }
4192
4193        let priority = match how {
4194            ResumeThread::Resume => Priority::Switch,
4195            ResumeThread::ResumeLater => Priority::Low,
4196
4197            // Promotion is a noop unless the thread is ready, so only search
4198            // through pending work items to see if something matches.
4199            ResumeThread::Promote => {
4200                let instance = self.runtime_instance(runtime_instance);
4201                let do_not_enter = store
4202                    .instance_state(instance)
4203                    .concurrent_state()
4204                    .do_not_enter;
4205                return store.concurrent_state_mut()?.promote_work_item_matching(
4206                    |item: &WorkItem| match item {
4207                        WorkItem::ResumeThread { thread, .. }
4208                        | WorkItem::ResumeFiber { thread, .. } => *thread == guest_thread,
4209
4210                        WorkItem::GuestCall {
4211                            call: GuestCall { thread, kind },
4212                            ..
4213                        } => {
4214                            *thread == guest_thread
4215                                && match kind {
4216                                    // Events can only be delivered if this
4217                                    // instance's exclusive lock isn't held, so
4218                                    // gate on that.
4219                                    GuestCallKind::DeliverEvent { .. } => !do_not_enter,
4220
4221                                    // A not-yet-started thread made ready via
4222                                    // `thread.resume-later`.
4223                                    GuestCallKind::StartExplicit(_) => true,
4224
4225                                    // Technically shouldn't be possible since
4226                                    // if this hasn't started nothing could get
4227                                    // its id...
4228                                    GuestCallKind::StartImplicit(_) => false,
4229                                }
4230                        }
4231
4232                        WorkItem::PushFuture(_) | WorkItem::WorkerFunction(_) => false,
4233                    },
4234                );
4235            }
4236        };
4237
4238        let state = store.concurrent_state_mut()?;
4239        let thread = state.get_mut(guest_thread.thread)?;
4240
4241        // When resuming a thread it must be in a suspended state otherwise
4242        // this operation is a trap.
4243        match &thread.state {
4244            GuestThreadState::NotStartedExplicit(_) | GuestThreadState::Suspended(_) => {}
4245            _ => bail!(Trap::CannotResumeThread),
4246        }
4247
4248        match mem::replace(&mut thread.state, GuestThreadState::Running) {
4249            GuestThreadState::NotStartedExplicit(start_func) => {
4250                log::trace!("starting thread {guest_thread:?}");
4251                let guest_call = WorkItem::GuestCall {
4252                    instance: self.runtime_instance(runtime_instance),
4253                    call: GuestCall {
4254                        thread: guest_thread,
4255                        kind: GuestCallKind::StartExplicit(Box::new(move |store| {
4256                            start_func(store, guest_thread)
4257                        })),
4258                    },
4259                };
4260                store
4261                    .concurrent_state_mut()?
4262                    .push_work_item(guest_call, priority)?;
4263            }
4264            GuestThreadState::Suspended(fiber) => {
4265                log::trace!("resuming thread {thread_id:?} that was suspended");
4266                store.concurrent_state_mut()?.push_work_item(
4267                    WorkItem::ResumeFiber {
4268                        instance: self.runtime_instance(runtime_instance),
4269                        thread: guest_thread,
4270                        fiber,
4271                    },
4272                    priority,
4273                )?;
4274            }
4275            other => {
4276                thread.state = other;
4277                bail_bug!("thread state checked to be resumable above");
4278            }
4279        }
4280        Ok(true)
4281    }
4282
4283    fn add_guest_thread_to_instance_table(
4284        self,
4285        thread_id: TableId<GuestThread>,
4286        store: &mut StoreOpaque,
4287        runtime_instance: RuntimeComponentInstanceIndex,
4288    ) -> Result<u32> {
4289        let guest_id = store
4290            .instance_state(self.runtime_instance(runtime_instance))
4291            .thread_handle_table()
4292            .guest_thread_insert(thread_id.rep())?;
4293        store
4294            .concurrent_state_mut()?
4295            .get_mut(thread_id)?
4296            .instance_rep = Some(guest_id);
4297        Ok(guest_id)
4298    }
4299
4300    /// Helper function for the `thread.yield`, thread.suspend`,
4301    /// `thread.suspend-then-resume`, `thread.suspend-then-promote`,
4302    /// `thread.yield-then-resume`, and `thread.yield-then-promote` intrinsics.
4303    pub(crate) fn suspension_intrinsic(
4304        self,
4305        store: &mut StoreOpaque,
4306        caller: RuntimeComponentInstanceIndex,
4307        yielding: bool,
4308        to_thread: SuspensionTarget,
4309    ) -> Result<WaitResult> {
4310        let check_suspend = match to_thread {
4311            SuspensionTarget::Promote(thread) => {
4312                !self.resume_thread(store, caller, thread, ResumeThread::Promote)?
4313            }
4314            SuspensionTarget::Resume(thread) => {
4315                if !self.resume_thread(store, caller, thread, ResumeThread::Resume)? {
4316                    bail_bug!(
4317                        "`resume_thread` should only ever return false \
4318                         when `ResumeThread::Promote` is passed to it"
4319                    );
4320                }
4321                false
4322            }
4323            SuspensionTarget::None => true,
4324        };
4325
4326        if check_suspend && !store.switch_if_may_not_suspend(self.runtime_instance(caller))? {
4327            return if yielding {
4328                Ok(WaitResult::Completed)
4329            } else {
4330                Err(Trap::CannotBlockSyncTask.into())
4331            };
4332        }
4333
4334        let guest_thread = store.current_guest_thread()?;
4335
4336        let reason = if yielding {
4337            SuspendReason::Yielding {
4338                thread: guest_thread,
4339            }
4340        } else {
4341            SuspendReason::ExplicitlySuspending {
4342                thread: guest_thread,
4343            }
4344        };
4345
4346        store.suspend(reason)?;
4347
4348        Ok(WaitResult::Completed)
4349    }
4350
4351    /// Helper function for the `waitable-set.wait` and `waitable-set.poll` intrinsics.
4352    fn waitable_check(
4353        self,
4354        store: &mut StoreOpaque,
4355        caller: RuntimeInstance,
4356        check: WaitableCheck,
4357        params: WaitableCheckParams,
4358    ) -> Result<u32> {
4359        let guest_thread = store.current_guest_thread()?;
4360
4361        log::trace!("waitable check for {guest_thread:?}; set {:?}", params.set);
4362
4363        // If we're waiting, and there are no events immediately available,
4364        // suspend the fiber until that changes.
4365        match &check {
4366            WaitableCheck::Wait => {
4367                let set = params.set;
4368
4369                // Note that this is a loop since, by the time this fiber is
4370                // resumed after being woken up, another thread may have already
4371                // taken the event (e.g. via `waitable-set.poll`), in which case
4372                // we keep waiting.
4373                loop {
4374                    let state = store.concurrent_state_mut()?;
4375                    let task = state.get_mut(guest_thread.task)?;
4376                    if !(task.event.is_none() || matches!(task.event, Some(Event::Cancelled)))
4377                        || !state.get_mut(set)?.ready.is_empty()
4378                    {
4379                        break;
4380                    }
4381
4382                    store.switch_or_trap_if_may_not_suspend(caller)?;
4383
4384                    store.suspend(SuspendReason::Waiting {
4385                        set,
4386                        thread: guest_thread,
4387                    })?;
4388                }
4389            }
4390            WaitableCheck::Poll => {}
4391        }
4392
4393        log::trace!(
4394            "waitable check for {guest_thread:?}; set {:?}, part two",
4395            params.set
4396        );
4397
4398        // Deliver any pending events to the guest and return.
4399        let event = self.get_event(store, guest_thread.task, Some(params.set), false)?;
4400
4401        let (ordinal, handle, result) = match &check {
4402            WaitableCheck::Wait => {
4403                let (event, waitable) = match event {
4404                    Some(p) => p,
4405                    None => bail_bug!("event expected to be present"),
4406                };
4407                let handle = waitable.map(|(_, v)| v).unwrap_or(0);
4408                let (ordinal, result) = event.parts();
4409                (ordinal, handle, result)
4410            }
4411            WaitableCheck::Poll => {
4412                if let Some((event, waitable)) = event {
4413                    let handle = waitable.map(|(_, v)| v).unwrap_or(0);
4414                    let (ordinal, result) = event.parts();
4415                    (ordinal, handle, result)
4416                } else {
4417                    log::trace!(
4418                        "no events ready to deliver via waitable-set.poll to {:?}; set {:?}",
4419                        guest_thread.task,
4420                        params.set
4421                    );
4422                    let (ordinal, result) = Event::None.parts();
4423                    (ordinal, 0, result)
4424                }
4425            }
4426        };
4427        let memory = self.options_memory_mut(store, params.options);
4428        let ptr = crate::component::func::validate_inbounds_dynamic(
4429            &CanonicalAbiInfo::POINTER_PAIR,
4430            memory,
4431            &ValRaw::u32(params.payload),
4432        )?;
4433        memory[ptr + 0..][..4].copy_from_slice(&handle.to_le_bytes());
4434        memory[ptr + 4..][..4].copy_from_slice(&result.to_le_bytes());
4435        Ok(ordinal)
4436    }
4437
4438    /// Implements the `subtask.cancel` intrinsic.
4439    pub(crate) fn subtask_cancel(
4440        self,
4441        store: &mut StoreOpaque,
4442        caller_instance: RuntimeComponentInstanceIndex,
4443        async_: bool,
4444        task_id: u32,
4445    ) -> Result<u32> {
4446        let (rep, is_host) = store
4447            .instance_state(self.runtime_instance(caller_instance))
4448            .handle_table()
4449            .subtask_rep(task_id)?;
4450        let waitable = if is_host {
4451            Waitable::Host(TableId::<HostTask>::new(rep))
4452        } else {
4453            Waitable::Guest(TableId::<GuestTask>::new(rep))
4454        };
4455        let concurrent_state = store.concurrent_state_mut()?;
4456
4457        log::trace!("subtask_cancel {waitable:?} (handle {task_id}; async {async_})");
4458
4459        waitable.trap_if_in_waitable_set(concurrent_state)?;
4460
4461        let needs_block;
4462        if let Waitable::Host(host_task) = waitable {
4463            let state = &mut concurrent_state.get_mut(host_task)?.state;
4464            match state {
4465                // If the callee is still running, signal an abort is requested.
4466                //
4467                // After cancelling this falls through to block waiting for the
4468                // host task to actually finish assuming that `async_` is false.
4469                // This blocking behavior resolves the race of `handle.abort()`
4470                // with the task actually getting cancelled or finishing.
4471                HostTaskState::CalleeRunning(handle) => {
4472                    handle.abort();
4473                    *state = HostTaskState::CalleeCancelling;
4474                    needs_block = true;
4475                }
4476
4477                // Cancellation was already requested, so fail as the task can't
4478                // be cancelled twice.
4479                HostTaskState::CalleeCancelling | HostTaskState::CalleeDone { cancelled: true } => {
4480                    bail!(Trap::SubtaskCancelAfterTerminal);
4481                }
4482                HostTaskState::CalleeDone { cancelled: false } => {
4483                    // The callee is already done so there's no need to
4484                    // block further for an event.
4485                    *state = HostTaskState::CalleeDone { cancelled: true };
4486                    needs_block = false;
4487                }
4488
4489                // These states should not be possible for a subtask that's
4490                // visible from the guest, so trap here.
4491                HostTaskState::CalleeStarted | HostTaskState::CalleeFinished(_) => {
4492                    bail_bug!("invalid states for host callee")
4493                }
4494            }
4495        } else {
4496            let guest_task = TableId::<GuestTask>::new(rep);
4497            let task = concurrent_state.get_mut(guest_task)?;
4498            if !task.already_lowered_parameters() {
4499                store.cancel_guest_subtask_without_lowered_parameters(
4500                    self.runtime_instance(caller_instance),
4501                    guest_task,
4502                )?;
4503                return Ok(Status::StartCancelled as u32);
4504            } else if !task.returned_or_cancelled() {
4505                // Started, but not yet returned or cancelled; send the
4506                // `CANCELLED` event
4507                //
4508                // Note that this might overwrite an event that was set earlier
4509                // (e.g. `Event::None` if the task is yielding, or
4510                // `Event::Cancelled` if it was already cancelled), but that's
4511                // okay -- this should supersede the previous state.
4512                task.event = Some(Event::Cancelled);
4513                let runtime_instance = task.instance;
4514                for thread in task.threads.clone() {
4515                    let thread = QualifiedThreadId {
4516                        task: guest_task,
4517                        thread,
4518                    };
4519                    let concurrent_state = store.concurrent_state_mut()?;
4520                    let thread_mut = concurrent_state.get_mut(thread.thread)?;
4521
4522                    let yield_ = |store: &mut StoreOpaque| {
4523                        // While we're yielding, temporarily set
4524                        // `do_not_suspend` to false on the subtask's instance
4525                        // since we'll be getting control back if it does
4526                        // suspend.
4527                        let state = store.instance_state(runtime_instance).concurrent_state();
4528                        let old_do_not_suspend = state.do_not_suspend;
4529                        state.do_not_suspend = false;
4530
4531                        let caller = store.current_guest_thread()?;
4532
4533                        // Temporarily add the waitable to the caller's
4534                        // `sync_call_set` to ensure that (1) it isn't already
4535                        // part of a different set and (2) it can't be added to
4536                        // a different set while we yield to the subtask.
4537                        let state = store.concurrent_state_mut()?;
4538                        let set = state.get_mut(caller.thread)?.sync_call_set;
4539                        waitable.join(state, Some(set))?;
4540
4541                        store.suspend(SuspendReason::YieldingToSubtask { thread: caller })?;
4542
4543                        let state = store.concurrent_state_mut()?;
4544                        waitable.join(state, None)?;
4545
4546                        store
4547                            .instance_state(runtime_instance)
4548                            .concurrent_state()
4549                            .do_not_suspend = old_do_not_suspend;
4550
4551                        Ok::<(), crate::Error>(())
4552                    };
4553
4554                    match thread_mut.wake_on_cancel.take() {
4555                        WakeOnCancel::Waiting(set) => {
4556                            // The thread is in a cancellable wait, so wake it up:
4557                            let item = match concurrent_state.get_mut(set)?.waiting.remove(&thread)
4558                            {
4559                                Some(WaitMode::Callback(instance)) => WorkItem::GuestCall {
4560                                    instance: runtime_instance,
4561                                    call: GuestCall {
4562                                        thread,
4563                                        kind: GuestCallKind::DeliverEvent {
4564                                            instance,
4565                                            set: Some(set),
4566                                        },
4567                                    },
4568                                },
4569                                other => bail_bug!(
4570                                    "expected `Some(WaitMode::Callback(_))`; got `{other:?}`"
4571                                ),
4572                            };
4573                            concurrent_state.set_switch_item(item)?;
4574
4575                            yield_(store)?;
4576
4577                            break;
4578                        }
4579                        WakeOnCancel::Yielding => {
4580                            if concurrent_state.promote_thread_work_item(thread)? {
4581                                yield_(store)?;
4582                                break;
4583                            } else if store
4584                                .instance_state(runtime_instance)
4585                                .concurrent_state()
4586                                .pending
4587                                .contains_key(&thread)
4588                            {
4589                                // The thread's work item was parked in
4590                                // `pending` because its instance can't be
4591                                // entered right now, so the cancellation
4592                                // can't be delivered yet. It'll be delivered
4593                                // via `GuestTask::event` once the instance
4594                                // becomes enterable and the work item runs.
4595                                // Until then the thread is still yielding, so
4596                                // restore `wake_on_cancel` and move on to the
4597                                // next thread.
4598                                store
4599                                    .concurrent_state_mut()?
4600                                    .get_mut(thread.thread)?
4601                                    .wake_on_cancel = WakeOnCancel::Yielding;
4602                            } else {
4603                                bail_bug!("thread with `WakeOnCancel::Yielding` not promotable");
4604                            }
4605                        }
4606                        WakeOnCancel::None => {}
4607                    }
4608                }
4609
4610                // Guest tasks need to block if they have not yet returned or
4611                // cancelled, even as a result of the event delivery above.
4612                needs_block = !store
4613                    .concurrent_state_mut()?
4614                    .get_mut(guest_task)?
4615                    .returned_or_cancelled()
4616            } else {
4617                needs_block = false;
4618            }
4619        };
4620
4621        // If we need to block waiting on the terminal status of this subtask
4622        // then return immediately in `async` mode, or otherwise wait for the
4623        // event to get signaled through the store.
4624        if needs_block {
4625            if async_ {
4626                return Ok(BLOCKED);
4627            }
4628
4629            // Save and later restore `next_switch_item` during a sync cancel so
4630            // we don't try to switch to it while blocking.
4631            let old_next_switch_item = {
4632                let state = store.concurrent_state_mut()?;
4633                let item = state.next_switch_item.take();
4634                // Note that we store it in the table here rather than directly
4635                // in a local variable to ensure the fiber is disposed of
4636                // properly if we end up trapping or panicking.
4637                state.push(item)?
4638            };
4639
4640            // Wait for this waitable to get signaled with its terminal
4641            // status. Once that's done fall through to the shared code.
4642            store.wait_for_event(self.runtime_instance(caller_instance), waitable)?;
4643
4644            let state = store.concurrent_state_mut()?;
4645            state.next_switch_item = state.delete(old_next_switch_item)?;
4646
4647            // .. fall through to determine what event's in store for us.
4648        }
4649
4650        let event = waitable.take_event(store.concurrent_state_mut()?)?;
4651        if let Some(Event::Subtask {
4652            status: status @ (Status::Returned | Status::ReturnCancelled),
4653        }) = event
4654        {
4655            Ok(status as u32)
4656        } else {
4657            bail!(Trap::SubtaskCancelAfterTerminal);
4658        }
4659    }
4660}
4661
4662/// Trait representing component model ABI async intrinsics and fused adapter
4663/// helper functions.
4664///
4665/// SAFETY (callers): Most of the methods in this trait accept raw pointers,
4666/// which must be valid for at least the duration of the call (and possibly for
4667/// as long as the relevant guest task exists, in the case of `*mut VMFuncRef`
4668/// pointers used for async calls).
4669pub trait VMComponentAsyncStore {
4670    /// A helper function for fused adapter modules involving calls where the
4671    /// one of the caller or callee is async.
4672    ///
4673    /// This helper is not used when the caller and callee both use the sync
4674    /// ABI, only when at least one is async is this used.
4675    unsafe fn prepare_call(
4676        &mut self,
4677        instance: Instance,
4678        memory: *mut VMMemoryDefinition,
4679        start: NonNull<VMFuncRef>,
4680        return_: NonNull<VMFuncRef>,
4681        caller_instance: RuntimeComponentInstanceIndex,
4682        callee_instance: RuntimeComponentInstanceIndex,
4683        task_return_type: TypeTupleIndex,
4684        callee_async: bool,
4685        string_encoding: StringEncoding,
4686        result_count: u32,
4687        storage: *mut ValRaw,
4688        storage_len: usize,
4689    ) -> Result<()>;
4690
4691    /// A helper function for fused adapter modules to start a call previously
4692    /// set up with `prepare_call`.
4693    unsafe fn start_call(
4694        &mut self,
4695        instance: Instance,
4696        callback: *mut VMFuncRef,
4697        post_return: *mut VMFuncRef,
4698        callee: NonNull<VMFuncRef>,
4699        param_count: u32,
4700        result_count: u32,
4701        flags: u32,
4702        storage: *mut MaybeUninit<ValRaw>,
4703        storage_len: usize,
4704    ) -> Result<()>;
4705
4706    /// The `future.write` intrinsic.
4707    fn future_write(
4708        &mut self,
4709        instance: Instance,
4710        caller: RuntimeComponentInstanceIndex,
4711        ty: TypeFutureTableIndex,
4712        options: OptionsIndex,
4713        future: u32,
4714        address: u32,
4715    ) -> Result<u32>;
4716
4717    /// The `future.read` intrinsic.
4718    fn future_read(
4719        &mut self,
4720        instance: Instance,
4721        caller: RuntimeComponentInstanceIndex,
4722        ty: TypeFutureTableIndex,
4723        options: OptionsIndex,
4724        future: u32,
4725        address: u32,
4726    ) -> Result<u32>;
4727
4728    /// The `future.drop-writable` intrinsic.
4729    fn future_drop_writable(
4730        &mut self,
4731        instance: Instance,
4732        ty: TypeFutureTableIndex,
4733        writer: u32,
4734    ) -> Result<()>;
4735
4736    /// The `stream.write` intrinsic.
4737    fn stream_write(
4738        &mut self,
4739        instance: Instance,
4740        caller: RuntimeComponentInstanceIndex,
4741        ty: TypeStreamTableIndex,
4742        options: OptionsIndex,
4743        stream: u32,
4744        address: u32,
4745        count: u32,
4746    ) -> Result<u32>;
4747
4748    /// The `stream.read` intrinsic.
4749    fn stream_read(
4750        &mut self,
4751        instance: Instance,
4752        caller: RuntimeComponentInstanceIndex,
4753        ty: TypeStreamTableIndex,
4754        options: OptionsIndex,
4755        stream: u32,
4756        address: u32,
4757        count: u32,
4758    ) -> Result<u32>;
4759
4760    /// The "fast-path" implementation of the `stream.write` intrinsic for
4761    /// "flat" (i.e. memcpy-able) payloads.
4762    fn flat_stream_write(
4763        &mut self,
4764        instance: Instance,
4765        caller: RuntimeComponentInstanceIndex,
4766        ty: TypeStreamTableIndex,
4767        options: OptionsIndex,
4768        payload_size: u32,
4769        payload_align: u32,
4770        stream: u32,
4771        address: u32,
4772        count: u32,
4773    ) -> Result<u32>;
4774
4775    /// The "fast-path" implementation of the `stream.read` intrinsic for "flat"
4776    /// (i.e. memcpy-able) payloads.
4777    fn flat_stream_read(
4778        &mut self,
4779        instance: Instance,
4780        caller: RuntimeComponentInstanceIndex,
4781        ty: TypeStreamTableIndex,
4782        options: OptionsIndex,
4783        payload_size: u32,
4784        payload_align: u32,
4785        stream: u32,
4786        address: u32,
4787        count: u32,
4788    ) -> Result<u32>;
4789
4790    /// The `stream.drop-writable` intrinsic.
4791    fn stream_drop_writable(
4792        &mut self,
4793        instance: Instance,
4794        ty: TypeStreamTableIndex,
4795        writer: u32,
4796    ) -> Result<()>;
4797
4798    /// The `error-context.debug-message` intrinsic.
4799    fn error_context_debug_message(
4800        &mut self,
4801        instance: Instance,
4802        ty: TypeComponentLocalErrorContextTableIndex,
4803        options: OptionsIndex,
4804        err_ctx_handle: u32,
4805        debug_msg_address: u32,
4806    ) -> Result<()>;
4807
4808    /// The `thread.new-indirect` intrinsic
4809    fn thread_new_indirect(
4810        &mut self,
4811        instance: Instance,
4812        caller: RuntimeComponentInstanceIndex,
4813        func_ty_idx: ModuleInternedTypeIndex,
4814        start_func_table_idx: RuntimeTableIndex,
4815        start_func_idx: u32,
4816        context: i32,
4817    ) -> Result<u32>;
4818}
4819
4820/// SAFETY: See trait docs.
4821impl<T: 'static> VMComponentAsyncStore for StoreInner<T> {
4822    unsafe fn prepare_call(
4823        &mut self,
4824        instance: Instance,
4825        memory: *mut VMMemoryDefinition,
4826        start: NonNull<VMFuncRef>,
4827        return_: NonNull<VMFuncRef>,
4828        caller_instance: RuntimeComponentInstanceIndex,
4829        callee_instance: RuntimeComponentInstanceIndex,
4830        task_return_type: TypeTupleIndex,
4831        callee_async: bool,
4832        string_encoding: StringEncoding,
4833        result_count_or_max_if_async: u32,
4834        storage: *mut ValRaw,
4835        storage_len: usize,
4836    ) -> Result<()> {
4837        // SAFETY: The `wasmtime_cranelift`-generated code that calls
4838        // this method will have ensured that `storage` is a valid
4839        // pointer containing at least `storage_len` items.
4840        let params = unsafe { core::slice::from_raw_parts(storage, storage_len) }.to_vec();
4841
4842        unsafe {
4843            instance.prepare_call(
4844                StoreContextMut(self),
4845                start,
4846                return_,
4847                caller_instance,
4848                callee_instance,
4849                task_return_type,
4850                callee_async,
4851                memory,
4852                string_encoding,
4853                match result_count_or_max_if_async {
4854                    PREPARE_ASYNC_NO_RESULT => CallerInfo::Async {
4855                        params,
4856                        has_result: false,
4857                    },
4858                    PREPARE_ASYNC_WITH_RESULT => CallerInfo::Async {
4859                        params,
4860                        has_result: true,
4861                    },
4862                    result_count => CallerInfo::Sync {
4863                        params,
4864                        result_count,
4865                    },
4866                },
4867            )
4868        }
4869    }
4870
4871    unsafe fn start_call(
4872        &mut self,
4873        instance: Instance,
4874        callback: *mut VMFuncRef,
4875        post_return: *mut VMFuncRef,
4876        callee: NonNull<VMFuncRef>,
4877        param_count: u32,
4878        result_count: u32,
4879        flags: u32,
4880        storage: *mut MaybeUninit<ValRaw>,
4881        storage_len: usize,
4882    ) -> Result<()> {
4883        unsafe {
4884            instance.start_call(
4885                StoreContextMut(self),
4886                callback,
4887                post_return,
4888                callee,
4889                param_count,
4890                result_count,
4891                flags,
4892                // SAFETY: The `wasmtime_cranelift`-generated code that calls
4893                // this method will have ensured that `storage` is a valid
4894                // pointer containing at least `storage_len` items.
4895                core::slice::from_raw_parts_mut(storage, storage_len),
4896            )
4897        }
4898    }
4899
4900    fn future_write(
4901        &mut self,
4902        instance: Instance,
4903        caller: RuntimeComponentInstanceIndex,
4904        ty: TypeFutureTableIndex,
4905        options: OptionsIndex,
4906        future: u32,
4907        address: u32,
4908    ) -> Result<u32> {
4909        instance
4910            .guest_write(
4911                StoreContextMut(self),
4912                caller,
4913                TransmitIndex::Future(ty),
4914                options,
4915                None,
4916                future,
4917                address,
4918                1,
4919            )
4920            .map(|result| result.encode())
4921    }
4922
4923    fn future_read(
4924        &mut self,
4925        instance: Instance,
4926        caller: RuntimeComponentInstanceIndex,
4927        ty: TypeFutureTableIndex,
4928        options: OptionsIndex,
4929        future: u32,
4930        address: u32,
4931    ) -> Result<u32> {
4932        instance
4933            .guest_read(
4934                StoreContextMut(self),
4935                caller,
4936                TransmitIndex::Future(ty),
4937                options,
4938                None,
4939                future,
4940                address,
4941                1,
4942            )
4943            .map(|result| result.encode())
4944    }
4945
4946    fn stream_write(
4947        &mut self,
4948        instance: Instance,
4949        caller: RuntimeComponentInstanceIndex,
4950        ty: TypeStreamTableIndex,
4951        options: OptionsIndex,
4952        stream: u32,
4953        address: u32,
4954        count: u32,
4955    ) -> Result<u32> {
4956        instance
4957            .guest_write(
4958                StoreContextMut(self),
4959                caller,
4960                TransmitIndex::Stream(ty),
4961                options,
4962                None,
4963                stream,
4964                address,
4965                count,
4966            )
4967            .map(|result| result.encode())
4968    }
4969
4970    fn stream_read(
4971        &mut self,
4972        instance: Instance,
4973        caller: RuntimeComponentInstanceIndex,
4974        ty: TypeStreamTableIndex,
4975        options: OptionsIndex,
4976        stream: u32,
4977        address: u32,
4978        count: u32,
4979    ) -> Result<u32> {
4980        instance
4981            .guest_read(
4982                StoreContextMut(self),
4983                caller,
4984                TransmitIndex::Stream(ty),
4985                options,
4986                None,
4987                stream,
4988                address,
4989                count,
4990            )
4991            .map(|result| result.encode())
4992    }
4993
4994    fn future_drop_writable(
4995        &mut self,
4996        instance: Instance,
4997        ty: TypeFutureTableIndex,
4998        writer: u32,
4999    ) -> Result<()> {
5000        instance.guest_drop_writable(self, TransmitIndex::Future(ty), writer)
5001    }
5002
5003    fn flat_stream_write(
5004        &mut self,
5005        instance: Instance,
5006        caller: RuntimeComponentInstanceIndex,
5007        ty: TypeStreamTableIndex,
5008        options: OptionsIndex,
5009        payload_size: u32,
5010        payload_align: u32,
5011        stream: u32,
5012        address: u32,
5013        count: u32,
5014    ) -> Result<u32> {
5015        instance
5016            .guest_write(
5017                StoreContextMut(self),
5018                caller,
5019                TransmitIndex::Stream(ty),
5020                options,
5021                Some(FlatAbi {
5022                    size: payload_size,
5023                    align: payload_align,
5024                }),
5025                stream,
5026                address,
5027                count,
5028            )
5029            .map(|result| result.encode())
5030    }
5031
5032    fn flat_stream_read(
5033        &mut self,
5034        instance: Instance,
5035        caller: RuntimeComponentInstanceIndex,
5036        ty: TypeStreamTableIndex,
5037        options: OptionsIndex,
5038        payload_size: u32,
5039        payload_align: u32,
5040        stream: u32,
5041        address: u32,
5042        count: u32,
5043    ) -> Result<u32> {
5044        instance
5045            .guest_read(
5046                StoreContextMut(self),
5047                caller,
5048                TransmitIndex::Stream(ty),
5049                options,
5050                Some(FlatAbi {
5051                    size: payload_size,
5052                    align: payload_align,
5053                }),
5054                stream,
5055                address,
5056                count,
5057            )
5058            .map(|result| result.encode())
5059    }
5060
5061    fn stream_drop_writable(
5062        &mut self,
5063        instance: Instance,
5064        ty: TypeStreamTableIndex,
5065        writer: u32,
5066    ) -> Result<()> {
5067        instance.guest_drop_writable(self, TransmitIndex::Stream(ty), writer)
5068    }
5069
5070    fn error_context_debug_message(
5071        &mut self,
5072        instance: Instance,
5073        ty: TypeComponentLocalErrorContextTableIndex,
5074        options: OptionsIndex,
5075        err_ctx_handle: u32,
5076        debug_msg_address: u32,
5077    ) -> Result<()> {
5078        instance.error_context_debug_message(
5079            StoreContextMut(self),
5080            ty,
5081            options,
5082            err_ctx_handle,
5083            debug_msg_address,
5084        )
5085    }
5086
5087    fn thread_new_indirect(
5088        &mut self,
5089        instance: Instance,
5090        caller: RuntimeComponentInstanceIndex,
5091        func_ty_idx: ModuleInternedTypeIndex,
5092        start_func_table_idx: RuntimeTableIndex,
5093        start_func_idx: u32,
5094        context: i32,
5095    ) -> Result<u32> {
5096        instance.thread_new_indirect(
5097            StoreContextMut(self),
5098            caller,
5099            func_ty_idx,
5100            start_func_table_idx,
5101            start_func_idx,
5102            context,
5103        )
5104    }
5105}
5106
5107type HostTaskFuture = Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>;
5108
5109/// Runs the given future with the current thread set to `task` each time it is
5110/// polled.
5111async fn run_with_host_task_set<F>(task: TableId<HostTask>, future: F) -> Result<F::Output>
5112where
5113    F: Future,
5114{
5115    let mut future = pin!(future);
5116    future::poll_fn(|cx| {
5117        let old_thread = match tls::get(|store| store.set_thread(task)) {
5118            Ok(thread) => thread,
5119            Err(error) => return Poll::Ready(Err(error)),
5120        };
5121        let result = future.as_mut().poll(cx);
5122        match tls::get(|store| store.set_thread(old_thread)) {
5123            Ok(_) => result.map(Ok),
5124            Err(error) => Poll::Ready(Err(error)),
5125        }
5126    })
5127    .await
5128}
5129
5130/// Represents the state of a pending host task.
5131///
5132/// This is used to represent tasks when the guest calls into the host.
5133pub(crate) struct HostTask {
5134    common: WaitableCommon,
5135
5136    /// State of borrows/etc the host needs to track. Used when the guest passes
5137    /// borrows to the host, for example.
5138    call_context: CallContext,
5139
5140    state: HostTaskState,
5141
5142    group: TaskGroupId,
5143}
5144
5145enum HostTaskState {
5146    /// A host task has been created and it's considered "started".
5147    ///
5148    /// The host task has yet to enter `first_poll` or `poll_and_block` which
5149    /// is where this will get updated further.
5150    CalleeStarted,
5151
5152    /// State used for tasks in `first_poll` meaning that the guest did an async
5153    /// lower of a host async function which is blocked. The specified handle is
5154    /// linked to the future in the main `FuturesUnordered` of a store which is
5155    /// used to cancel it if the guest requests cancellation.
5156    CalleeRunning(JoinHandle),
5157
5158    /// Cancellation was requested, but the completion worker has not finished
5159    /// lowering the result and publishing the terminal event. The task must
5160    /// remain in the table until that worker completes.
5161    CalleeCancelling,
5162
5163    /// Terminal state used for tasks in `poll_and_block` to store the result of
5164    /// their computation. Note that this state is not used for tasks in
5165    /// `first_poll`.
5166    CalleeFinished(LiftedResult),
5167
5168    /// Terminal state for host tasks meaning that the task was cancelled or the
5169    /// result was taken.
5170    CalleeDone { cancelled: bool },
5171}
5172
5173impl HostTask {
5174    fn new(
5175        concurrent_state: &mut ConcurrentState,
5176        state: HostTaskState,
5177        caller: QualifiedThreadId,
5178    ) -> Result<Self> {
5179        let group = concurrent_state.get_mut(caller.task)?.group;
5180        concurrent_state.increment_group_ref_count(group)?;
5181
5182        Ok(Self {
5183            common: WaitableCommon::default(),
5184            call_context: CallContext::default(),
5185            state,
5186            group,
5187        })
5188    }
5189}
5190
5191impl TableDebug for HostTask {
5192    fn type_name() -> &'static str {
5193        "HostTask"
5194    }
5195}
5196
5197type CallbackFn = Box<dyn Fn(&mut dyn VMStore, Event, u32) -> Result<u32> + Send + Sync + 'static>;
5198
5199/// Represents the caller of a given guest task.
5200enum Caller {
5201    /// The host called the guest task.
5202    Host {
5203        /// If present, may be used to deliver the result.
5204        tx: Option<oneshot::Sender<LiftedResult>>,
5205        /// If true, there's a host future that must be dropped before the task
5206        /// can be deleted.
5207        host_future_present: bool,
5208        /// The host task which called into the guest, or `None` for a call from
5209        /// the top-level host. The task may belong to an entirely unrelated
5210        /// top-level component instance than the one the host called into.
5211        caller: Option<TableId<HostTask>>,
5212    },
5213    /// Another guest thread called the guest task
5214    Guest {
5215        /// The id of the caller
5216        thread: QualifiedThreadId,
5217    },
5218}
5219
5220/// Represents a closure and related canonical ABI parameters required to
5221/// validate a `task.return` call at runtime and lift the result.
5222struct LiftResult {
5223    lift: RawLift,
5224    ty: TypeTupleIndex,
5225    memory: Option<SendSyncPtr<VMMemoryDefinition>>,
5226    string_encoding: StringEncoding,
5227}
5228
5229/// The table ID for a guest thread, qualified by the task to which it belongs.
5230///
5231/// This exists to minimize table lookups and the necessity to pass stores around mutably
5232/// for the common case of identifying the task to which a thread belongs.
5233#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq)]
5234pub(crate) struct QualifiedThreadId {
5235    task: TableId<GuestTask>,
5236    thread: TableId<GuestThread>,
5237}
5238
5239impl QualifiedThreadId {
5240    fn qualify(
5241        state: &mut ConcurrentState,
5242        thread: TableId<GuestThread>,
5243    ) -> Result<QualifiedThreadId> {
5244        Ok(QualifiedThreadId {
5245            task: state.get_mut(thread)?.parent_task,
5246            thread,
5247        })
5248    }
5249}
5250
5251impl fmt::Debug for QualifiedThreadId {
5252    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
5253        f.debug_tuple("QualifiedThreadId")
5254            .field(&self.task.rep())
5255            .field(&self.thread.rep())
5256            .finish()
5257    }
5258}
5259
5260enum GuestThreadState {
5261    NotStartedImplicit,
5262    NotStartedExplicit(
5263        Box<dyn FnOnce(&mut dyn VMStore, QualifiedThreadId) -> Result<()> + Send + Sync>,
5264    ),
5265    Running,
5266    Suspended(StoreFiber<'static>),
5267    Ready {
5268        fiber: StoreFiber<'static>,
5269    },
5270    Completed,
5271}
5272
5273impl fmt::Debug for GuestThreadState {
5274    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
5275        match self {
5276            Self::NotStartedImplicit => f.debug_tuple("NotStartedImplicit").finish(),
5277            Self::NotStartedExplicit(_) => f.debug_tuple("NotStartedExplicit").finish(),
5278            Self::Running => f.debug_tuple("Running").finish(),
5279            Self::Suspended(_) => f.debug_tuple("Suspended").finish(),
5280            Self::Ready { .. } => f.debug_struct("Ready").finish(),
5281            Self::Completed => f.debug_tuple("Completed").finish(),
5282        }
5283    }
5284}
5285
5286#[derive(Copy, Clone, PartialEq, Eq, Debug)]
5287enum WakeOnCancel {
5288    None,
5289    Waiting(TableId<WaitableSet>),
5290    Yielding,
5291}
5292
5293impl WakeOnCancel {
5294    fn is_none(self) -> bool {
5295        matches!(self, WakeOnCancel::None)
5296    }
5297
5298    fn replace(&mut self, other: WakeOnCancel) -> Self {
5299        let old = *self;
5300        *self = other;
5301        old
5302    }
5303
5304    fn take(&mut self) -> Self {
5305        self.replace(WakeOnCancel::None)
5306    }
5307}
5308
5309pub struct GuestThread {
5310    /// Context-local state used to implement the `context.{get,set}`
5311    /// intrinsics.
5312    context: [u32; NUM_COMPONENT_CONTEXT_SLOTS],
5313    /// The owning guest task.
5314    parent_task: TableId<GuestTask>,
5315    /// If non-`None`, indicates that the thread is currently either waiting a
5316    /// waitable set or yielding but may be cancelled and woken immediately.
5317    wake_on_cancel: WakeOnCancel,
5318    /// The execution state of this guest thread
5319    state: GuestThreadState,
5320    /// The index of this thread in the component instance's handle table.
5321    /// This must always be `Some` after initialization.
5322    instance_rep: Option<u32>,
5323    /// Scratch waitable set used to watch subtasks during synchronous calls.
5324    sync_call_set: TableId<WaitableSet>,
5325    /// The old value of `do_not_suspend` prior to the sync-typed task for which
5326    /// this thread was created, if relevant.
5327    old_do_not_suspend: Option<bool>,
5328}
5329
5330impl GuestThread {
5331    /// Retrieve the `GuestThread` corresponding to the specified guest-visible
5332    /// handle.
5333    fn from_instance(
5334        state: Pin<&mut ComponentInstance>,
5335        caller_instance: RuntimeComponentInstanceIndex,
5336        guest_thread: u32,
5337    ) -> Result<TableId<Self>> {
5338        let rep = state.instance_states().0[caller_instance]
5339            .thread_handle_table()
5340            .guest_thread_rep(guest_thread)?;
5341        Ok(TableId::new(rep))
5342    }
5343
5344    fn new_implicit(state: &mut ConcurrentState, parent_task: TableId<GuestTask>) -> Result<Self> {
5345        let sync_call_set = state.push(WaitableSet {
5346            is_sync_call_set: true,
5347            ..WaitableSet::default()
5348        })?;
5349        Ok(Self {
5350            context: [0; NUM_COMPONENT_CONTEXT_SLOTS],
5351            parent_task,
5352            wake_on_cancel: WakeOnCancel::None,
5353            state: GuestThreadState::NotStartedImplicit,
5354            instance_rep: None,
5355            sync_call_set,
5356            old_do_not_suspend: None,
5357        })
5358    }
5359
5360    fn new_explicit(
5361        state: &mut ConcurrentState,
5362        parent_task: TableId<GuestTask>,
5363        start_func: Box<
5364            dyn FnOnce(&mut dyn VMStore, QualifiedThreadId) -> Result<()> + Send + Sync,
5365        >,
5366    ) -> Result<Self> {
5367        let sync_call_set = state.push(WaitableSet {
5368            is_sync_call_set: true,
5369            ..WaitableSet::default()
5370        })?;
5371        Ok(Self {
5372            context: [0; NUM_COMPONENT_CONTEXT_SLOTS],
5373            parent_task,
5374            wake_on_cancel: WakeOnCancel::None,
5375            state: GuestThreadState::NotStartedExplicit(start_func),
5376            instance_rep: None,
5377            sync_call_set,
5378            old_do_not_suspend: None,
5379        })
5380    }
5381}
5382
5383impl TableDebug for GuestThread {
5384    fn type_name() -> &'static str {
5385        "GuestThread"
5386    }
5387}
5388
5389enum SyncResult {
5390    NotProduced,
5391    Produced(Option<ValRaw>),
5392    Taken,
5393}
5394
5395impl SyncResult {
5396    fn take(&mut self) -> Result<Option<Option<ValRaw>>> {
5397        Ok(match mem::replace(self, SyncResult::Taken) {
5398            SyncResult::NotProduced => None,
5399            SyncResult::Produced(val) => Some(val),
5400            SyncResult::Taken => {
5401                bail_bug!("attempted to take a synchronous result that was already taken")
5402            }
5403        })
5404    }
5405}
5406
5407#[derive(Debug)]
5408enum HostFutureState {
5409    NotApplicable,
5410    Live,
5411    Dropped,
5412}
5413
5414/// Represents a pending guest task.
5415pub(crate) struct GuestTask {
5416    /// See `WaitableCommon`
5417    common: WaitableCommon,
5418    /// Closure to lower the parameters passed to this task.
5419    lower_params: Option<RawLower>,
5420    /// See `LiftResult`
5421    lift_result: Option<LiftResult>,
5422    /// A place to stash the type-erased lifted result if it can't be delivered
5423    /// immediately.
5424    result: Option<LiftedResult>,
5425    /// Closure to call the callback function for an async-lifted export, if
5426    /// provided.
5427    callback: Option<CallbackFn>,
5428    /// See `Caller`
5429    caller: Caller,
5430    /// Borrow state for this task.
5431    ///
5432    /// Keeps track of `borrow<T>` received to this task to ensure that
5433    /// everything is dropped by the time it exits.
5434    call_context: CallContext,
5435    /// A place to stash the lowered result for a sync-to-async call until it
5436    /// can be returned to the caller.
5437    sync_result: SyncResult,
5438    /// Whether or not the task has been cancelled (i.e. whether the
5439    /// cancellation request has been delivered to the task, and thus whether
5440    /// the task is permitted to call `task.cancel`).
5441    cancel_request_delivered: bool,
5442    /// Whether or not we've sent a `Status::Starting` event to any current or
5443    /// future waiters for this waitable.
5444    starting_sent: bool,
5445    /// The runtime instance to which the exported function for this guest task
5446    /// belongs.
5447    ///
5448    /// Note that the task may do a sync->sync call via a fused adapter which
5449    /// results in that task executing code in a different instance, and it may
5450    /// call host functions and intrinsics from that other instance.
5451    instance: RuntimeInstance,
5452    /// If present, a pending `Event::Cancelled` to be delivered to this task.
5453    event: Option<Event>,
5454    /// Whether or not the task has exited.
5455    exited: bool,
5456    /// Threads belonging to this task
5457    threads: HashSet<TableId<GuestThread>>,
5458    /// The state of the host future that represents an async task, which must
5459    /// be dropped before we can delete the task.
5460    host_future_state: HostFutureState,
5461    /// Indicates whether this task was created for a call to an async-typed
5462    /// export.
5463    async_typed: bool,
5464    /// Indicates whether this task was created for a call to an async-lifted
5465    /// export.
5466    async_lifted: bool,
5467
5468    decremented_interesting_task_count: bool,
5469
5470    group: TaskGroupId,
5471}
5472
5473impl GuestTask {
5474    fn already_lowered_parameters(&self) -> bool {
5475        // We reset `lower_params` after we lower the parameters
5476        self.lower_params.is_none()
5477    }
5478
5479    fn returned_or_cancelled(&self) -> bool {
5480        // We reset `lift_result` after we return or exit
5481        self.lift_result.is_none()
5482    }
5483
5484    fn ready_to_delete(&self) -> bool {
5485        let threads_completed = self.threads.is_empty();
5486        let has_sync_result = matches!(self.sync_result, SyncResult::Produced(_));
5487        let pending_completion_event = matches!(
5488            self.common.event,
5489            Some(Event::Subtask {
5490                status: Status::Returned | Status::ReturnCancelled
5491            })
5492        );
5493        let ready = threads_completed
5494            && !has_sync_result
5495            && !pending_completion_event
5496            && !matches!(self.host_future_state, HostFutureState::Live);
5497        log::trace!(
5498            "ready to delete? {ready} (threads_completed: {}, has_sync_result: {}, pending_completion_event: {}, host_future_state: {:?})",
5499            threads_completed,
5500            has_sync_result,
5501            pending_completion_event,
5502            self.host_future_state
5503        );
5504        ready
5505    }
5506
5507    fn new(
5508        state: &mut ConcurrentState,
5509        lower_params: RawLower,
5510        lift_result: LiftResult,
5511        caller: Caller,
5512        callback: Option<CallbackFn>,
5513        instance: RuntimeInstance,
5514        async_typed: bool,
5515        async_lifted: bool,
5516    ) -> Result<QualifiedThreadId> {
5517        let host_future_state = match &caller {
5518            Caller::Guest { .. } => HostFutureState::NotApplicable,
5519            Caller::Host {
5520                host_future_present,
5521                ..
5522            } => {
5523                if *host_future_present {
5524                    HostFutureState::Live
5525                } else {
5526                    HostFutureState::NotApplicable
5527                }
5528            }
5529        };
5530
5531        let group = match caller {
5532            Caller::Guest { thread } => {
5533                let group = state.get_mut(thread.task)?.group;
5534                state.increment_group_ref_count(group)?;
5535                group
5536            }
5537            Caller::Host { .. } => state.make_task_group()?,
5538        };
5539
5540        let task = state.push(Self {
5541            common: WaitableCommon::default(),
5542            lower_params: Some(lower_params),
5543            lift_result: Some(lift_result),
5544            result: None,
5545            callback,
5546            caller,
5547            call_context: CallContext::default(),
5548            sync_result: SyncResult::NotProduced,
5549            cancel_request_delivered: false,
5550            starting_sent: false,
5551            instance,
5552            event: None,
5553            exited: false,
5554            threads: HashSet::new(),
5555            host_future_state,
5556            async_typed,
5557            async_lifted,
5558            decremented_interesting_task_count: false,
5559            group,
5560        })?;
5561        let new_thread = GuestThread::new_implicit(state, task)?;
5562        let thread = state.push(new_thread)?;
5563        state.get_mut(task)?.threads.insert(thread);
5564        state.interesting_tasks += 1;
5565        let thread = QualifiedThreadId { task, thread };
5566        log::trace!("new implicit thread {thread:?} for instance {instance:?}");
5567        Ok(thread)
5568    }
5569}
5570
5571impl TableDebug for GuestTask {
5572    fn type_name() -> &'static str {
5573        "GuestTask"
5574    }
5575}
5576
5577/// Represents state common to all kinds of waitables.
5578#[derive(Default)]
5579struct WaitableCommon {
5580    /// The currently pending event for this waitable, if any.
5581    event: Option<Event>,
5582    /// The set to which this waitable belongs, if any.
5583    set: Option<TableId<WaitableSet>>,
5584    /// The handle with which the guest refers to this waitable, if any.
5585    handle: Option<u32>,
5586}
5587
5588/// Represents a Component Model Async `waitable`.
5589#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq)]
5590enum Waitable {
5591    /// A host task
5592    Host(TableId<HostTask>),
5593    /// A guest task
5594    Guest(TableId<GuestTask>),
5595    /// The read or write end of a stream or future
5596    Transmit(TableId<TransmitHandle>),
5597}
5598
5599impl Waitable {
5600    /// Retrieve the `Waitable` corresponding to the specified guest-visible
5601    /// handle.
5602    fn from_instance(
5603        state: Pin<&mut ComponentInstance>,
5604        caller_instance: RuntimeComponentInstanceIndex,
5605        waitable: u32,
5606    ) -> Result<Self> {
5607        use crate::runtime::vm::component::Waitable;
5608
5609        let (waitable, kind) = state.instance_states().0[caller_instance]
5610            .handle_table()
5611            .waitable_rep(waitable)?;
5612
5613        Ok(match kind {
5614            Waitable::Subtask { is_host: true } => Self::Host(TableId::new(waitable)),
5615            Waitable::Subtask { is_host: false } => Self::Guest(TableId::new(waitable)),
5616            Waitable::Stream | Waitable::Future => Self::Transmit(TableId::new(waitable)),
5617        })
5618    }
5619
5620    /// Retrieve the host-visible identifier for this `Waitable`.
5621    fn rep(&self) -> u32 {
5622        match self {
5623            Self::Host(id) => id.rep(),
5624            Self::Guest(id) => id.rep(),
5625            Self::Transmit(id) => id.rep(),
5626        }
5627    }
5628
5629    /// Move this `Waitable` to the specified set (when `set` is `Some(_)`) or
5630    /// remove it from any set it may currently belong to (when `set` is
5631    /// `None`).
5632    fn join(&self, state: &mut ConcurrentState, set: Option<TableId<WaitableSet>>) -> Result<()> {
5633        log::trace!("waitable {self:?} join set {set:?}");
5634
5635        let old = mem::replace(&mut self.common(state)?.set, set);
5636
5637        if let Some(old) = old {
5638            match *self {
5639                Waitable::Host(id) => state.remove_child(id, old),
5640                Waitable::Guest(id) => state.remove_child(id, old),
5641                Waitable::Transmit(id) => state.remove_child(id, old),
5642            }?;
5643
5644            state.get_mut(old)?.ready.remove(self);
5645        }
5646
5647        if let Some(set) = set {
5648            match *self {
5649                Waitable::Host(id) => state.add_child(id, set),
5650                Waitable::Guest(id) => state.add_child(id, set),
5651                Waitable::Transmit(id) => state.add_child(id, set),
5652            }?;
5653
5654            if self.common(state)?.event.is_some() {
5655                self.mark_ready(state)?;
5656            }
5657        }
5658
5659        Ok(())
5660    }
5661
5662    /// Retrieve mutable access to the `WaitableCommon` for this `Waitable`.
5663    fn common<'a>(&self, state: &'a mut ConcurrentState) -> Result<&'a mut WaitableCommon> {
5664        Ok(match self {
5665            Self::Host(id) => &mut state.get_mut(*id)?.common,
5666            Self::Guest(id) => &mut state.get_mut(*id)?.common,
5667            Self::Transmit(id) => &mut state.get_mut(*id)?.common,
5668        })
5669    }
5670
5671    /// Trap if this waitable is currently a member of a waitable set.
5672    ///
5673    /// A synchronous stream/future/subtask operation may end up blocking on
5674    /// this waitable, so it is not allowed to run while the waitable is also
5675    /// being watched by a waitable set.
5676    fn trap_if_in_waitable_set(&self, state: &mut ConcurrentState) -> Result<()> {
5677        if self.common(state)?.set.is_some() {
5678            bail!(Trap::WaitableSyncAndAsync);
5679        }
5680        Ok(())
5681    }
5682
5683    /// Set or clear the pending event for this waitable and either deliver it
5684    /// to the first waiter, if any, or mark it as ready to be delivered to the
5685    /// next waiter that arrives.
5686    fn set_event(&self, state: &mut ConcurrentState, event: Option<Event>) -> Result<()> {
5687        log::trace!("set event for {self:?}: {event:?}");
5688        self.common(state)?.event = event;
5689        self.mark_ready(state)
5690    }
5691
5692    /// Take the pending event from this waitable, leaving `None` in its place.
5693    fn take_event(&self, state: &mut ConcurrentState) -> Result<Option<Event>> {
5694        let common = self.common(state)?;
5695        let event = common.event.take();
5696        if let Some(set) = self.common(state)?.set {
5697            state.get_mut(set)?.ready.remove(self);
5698        }
5699
5700        Ok(event)
5701    }
5702
5703    /// Deliver the current event for this waitable to the first waiter, if any,
5704    /// or else mark it as ready to be delivered to the next waiter that
5705    /// arrives.
5706    fn mark_ready(&self, state: &mut ConcurrentState) -> Result<()> {
5707        if let Some(set) = self.common(state)?.set {
5708            state.get_mut(set)?.ready.insert(*self);
5709            state.wake_waiter(set)?;
5710        }
5711        Ok(())
5712    }
5713
5714    /// Remove this waitable from the store's rep table.
5715    fn delete_from(&self, store: &mut StoreOpaque) -> Result<()> {
5716        match self {
5717            Self::Host(task) => {
5718                log::trace!("delete host task {task:?}");
5719                let state = store.concurrent_state_mut()?;
5720                let task = state.delete(*task)?;
5721
5722                state.decrement_group_ref_count(task.group)?;
5723            }
5724            Self::Guest(task) => {
5725                log::trace!("delete guest task {task:?}");
5726                let state = store.concurrent_state_mut()?;
5727                let task = state.delete(*task)?;
5728
5729                state.decrement_group_ref_count(task.group)?;
5730
5731                // When a guest task is created it increments the
5732                // `ConcurrentState::interesting_tasks` counter, and that needs
5733                // to be paired with a decrement. There are a few situations in
5734                // which the decrement needs to happen which don't all funnel
5735                // through here, so in lieu of that at least try to catch issues
5736                // where we forgot to do a decrement.
5737                debug_assert!(task.decremented_interesting_task_count);
5738            }
5739            Self::Transmit(task) => {
5740                store.concurrent_state_mut()?.delete(*task)?;
5741            }
5742        }
5743
5744        Ok(())
5745    }
5746}
5747
5748impl fmt::Debug for Waitable {
5749    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
5750        match self {
5751            Self::Host(id) => write!(f, "{id:?}"),
5752            Self::Guest(id) => write!(f, "{id:?}"),
5753            Self::Transmit(id) => write!(f, "{id:?}"),
5754        }
5755    }
5756}
5757
5758/// Represents a Component Model Async `waitable-set`.
5759#[derive(Default)]
5760struct WaitableSet {
5761    /// Which waitables in this set have pending events, if any.
5762    ready: BTreeSet<Waitable>,
5763    /// Which guest threads are currently waiting on this set, if any.
5764    waiting: BTreeMap<QualifiedThreadId, WaitMode>,
5765    /// How many guest threads are waiting on this set, mirroring the spec's
5766    /// `WaitableSet.num_waiting`.
5767    num_waiting: usize,
5768    /// Whether this set is a synthetic, internal one meant for handling
5769    /// synchronous calls.
5770    is_sync_call_set: bool,
5771}
5772
5773impl WaitableSet {
5774    /// Records that a thread counted in `num_waiting` is no longer waiting.
5775    fn stop_waiting(&mut self) -> Result<()> {
5776        self.num_waiting = match self.num_waiting.checked_sub(1) {
5777            Some(n) => n,
5778            None => bail_bug!("waiter not accounted for in waitable set"),
5779        };
5780        Ok(())
5781    }
5782}
5783
5784impl TableDebug for WaitableSet {
5785    fn type_name() -> &'static str {
5786        "WaitableSet"
5787    }
5788}
5789
5790/// Type-erased closure to lower the parameters for a guest task.
5791type RawLower =
5792    Box<dyn FnOnce(&mut dyn VMStore, &mut [MaybeUninit<ValRaw>]) -> Result<()> + Send + Sync>;
5793
5794/// Type-erased closure to lift the result for a guest task.
5795type RawLift = Box<
5796    dyn FnOnce(&mut dyn VMStore, &[ValRaw]) -> Result<Box<dyn Any + Send + Sync>> + Send + Sync,
5797>;
5798
5799/// Type erased result of a guest task which may be downcast to the expected
5800/// type by a host caller (or simply ignored in the case of a guest caller; see
5801/// `DummyResult`).
5802type LiftedResult = Box<dyn Any + Send + Sync>;
5803
5804/// Used to return a result from a `LiftFn` when the actual result has already
5805/// been lowered to a guest task's stack and linear memory.
5806struct DummyResult;
5807
5808/// Represents the Component Model Async state of a (sub-)component instance.
5809#[derive(Default)]
5810pub struct ConcurrentInstanceState {
5811    /// Whether backpressure is set for this instance (enabled if >0)
5812    backpressure: u16,
5813    /// Whether this instance can be entered
5814    do_not_enter: bool,
5815    /// Whether this instance may suspend (i.e. whether this instance is
5816    /// currently running a sync-typed function).
5817    do_not_suspend: bool,
5818    /// Pending calls for this instance which require `Self::backpressure` to be
5819    /// zero and/or `Self::do_not_enter` to be false before they can proceed.
5820    pending: BTreeMap<QualifiedThreadId, GuestCallKind>,
5821}
5822
5823impl ConcurrentInstanceState {
5824    pub fn pending_is_empty(&self) -> bool {
5825        self.pending.is_empty()
5826    }
5827}
5828
5829#[derive(Debug, Copy, Clone)]
5830pub(crate) enum CurrentThread {
5831    /// The currently running thread is a guest, identified here with its
5832    /// task/thread id combo.
5833    Guest(QualifiedThreadId),
5834    /// The currently running thread is a host task.
5835    Host(TableId<HostTask>),
5836    /// The currently running thread is a host call whose task has not yet been
5837    /// materialized. The contained ID identifies its guest caller.
5838    DeferredHost(QualifiedThreadId),
5839    /// There is no currently running thread because we are in the main event
5840    /// loop or concurrency is disabled.
5841    None,
5842}
5843
5844impl CurrentThread {
5845    fn guest(&self) -> Option<&QualifiedThreadId> {
5846        match self {
5847            Self::Guest(id) => Some(id),
5848            _ => None,
5849        }
5850    }
5851
5852    fn guest_task(&self) -> Option<TableId<GuestTask>> {
5853        match self {
5854            Self::Guest(id) => Some(id.task),
5855            _ => None,
5856        }
5857    }
5858
5859    fn is_none(&self) -> bool {
5860        matches!(self, Self::None)
5861    }
5862}
5863
5864impl From<QualifiedThreadId> for CurrentThread {
5865    fn from(id: QualifiedThreadId) -> Self {
5866        Self::Guest(id)
5867    }
5868}
5869
5870impl From<TableId<HostTask>> for CurrentThread {
5871    fn from(id: TableId<HostTask>) -> Self {
5872        Self::Host(id)
5873    }
5874}
5875
5876enum Priority {
5877    Switch,
5878    High,
5879    Low,
5880}
5881
5882/// Represents the Component Model Async state of a store.
5883pub struct ConcurrentState {
5884    /// The currently running thread, if any.
5885    ///
5886    /// Note that we lazily materialize threads on-demand and this field is not
5887    /// necessarily up-to-date. The `StoreOpaque::current_thread` method should
5888    /// be preferred over directly accessing this field.
5889    unforced_current_thread: CurrentThread,
5890
5891    /// Borrow state for the deferred host call, if any.
5892    ///
5893    /// This is `Some` if and only if [`Self::unforced_current_thread`] is
5894    /// [`CurrentThread::DeferredHost`]. Materializing the host task moves this
5895    /// context into that task.
5896    deferred_host_call_context: Option<CallContext>,
5897
5898    /// The set of pending host and background tasks, if any.
5899    ///
5900    /// See `ComponentInstance::poll_until` for where we temporarily take this
5901    /// out, poll it, then put it back to avoid any mutable aliasing hazards.
5902    futures: AlwaysMut<Option<FuturesUnordered<HostTaskFuture>>>,
5903    /// The table of waitables, waitable sets, etc.
5904    table: AlwaysMut<ResourceTable>,
5905    /// The next item to switch to if any.
5906    ///
5907    /// This takes precedence over items in the `high_priority` queue below and
5908    /// should be used in cases such as subtask calls and thread resume/promote
5909    /// operations where we must switch to a specific thread at the next turn of
5910    /// the event loop regardless of what happens to be present in the
5911    /// `high_priority` queue.
5912    switch_item: Option<WorkItem>,
5913    /// The item to set `switch_item` to when the current thread suspends.
5914    ///
5915    /// This is used when ever an async-lowered import or `subtask.cancel` is
5916    /// called in order to track the thread to switch back to once the current
5917    /// subtask suspends, if any.
5918    next_switch_item: Option<WorkItem>,
5919    /// The "high priority" work queue for this store's event loop.
5920    high_priority: VecDeque<WorkItem>,
5921    /// The "low priority" work queue for this store's event loop.
5922    low_priority: VecDeque<WorkItem>,
5923    /// A place to stash the reason a fiber is suspending so that the code which
5924    /// resumed it will know under what conditions the fiber should be resumed
5925    /// again.
5926    suspend_reason: Option<SuspendReason>,
5927    /// A cached fiber which is waiting for work to do.
5928    ///
5929    /// This helps us avoid creating a new fiber for each `GuestCall` work item.
5930    worker: Option<StoreFiber<'static>>,
5931    /// A place to stash the work item for which we're resuming a worker fiber.
5932    worker_item: Option<WorkerItem>,
5933
5934    /// Reference counts for all component error contexts
5935    ///
5936    /// NOTE: it is possible the global ref count to be *greater* than the sum of
5937    /// (sub)component ref counts as tracked by `error_context_tables`, for
5938    /// example when the host holds one or more references to error contexts.
5939    ///
5940    /// The key of this primary map is often referred to as the "rep" (i.e. host-side
5941    /// component-wide representation) of the index into concurrent state for a given
5942    /// stored `ErrorContext`.
5943    ///
5944    /// Stated another way, `TypeComponentGlobalErrorContextTableIndex` is essentially the same
5945    /// as a `TableId<ErrorContextState>`.
5946    global_error_context_ref_counts:
5947        BTreeMap<TypeComponentGlobalErrorContextTableIndex, GlobalErrorContextRefCount>,
5948
5949    /// The number of "interesting tasks" currently executing in the store.
5950    ///
5951    /// This tracks the concept of a component instance lifetime as defined in
5952    /// https://github.com/WebAssembly/component-model/pull/643. Specifically
5953    /// all tasks currently increment this counter which then gets decremented
5954    /// when they exit. In the future some tasks might not increment this
5955    /// counter, but for now all do.
5956    ///
5957    /// This is used to implement `Accessor::poll_no_interesting_tasks` to
5958    /// inform the embedder when all tasks have completed. This is then
5959    /// used in wasmtime-wasi-http, for example, to know when an instance is
5960    /// idle.
5961    interesting_tasks: usize,
5962
5963    /// Single waker to notify when `interesting_tasks` reaches 0.
5964    ///
5965    /// Used in the implementation of `Accessor::poll_no_interesting_tasks`.
5966    interesting_tasks_empty_waker: Option<Waker>,
5967
5968    /// Single waker to notify when a component instance goes from
5969    /// not-concurrently-callable to concurrently-callable.
5970    ///
5971    /// Used in the implementation of `Accessor::poll_ready_for_concurrent_call`.
5972    ready_for_concurrent_call_waker: Option<Waker>,
5973
5974    /// Whether the `StoreContextMut::poll_until` event loop is running.
5975    event_loop_running: bool,
5976
5977    /// Stack of switch items to push to and pop from when entering and exiting
5978    /// sync-to-sync calls.
5979    saved_next_switch_items: Vec<Option<WorkItem>>,
5980
5981    /// See [TaskGroupHook].
5982    #[cfg(feature = "task-group-hook")]
5983    task_group_hook: Option<Box<dyn TaskGroupHook>>,
5984}
5985
5986impl Default for ConcurrentState {
5987    fn default() -> Self {
5988        Self {
5989            unforced_current_thread: CurrentThread::None,
5990            deferred_host_call_context: None,
5991            table: AlwaysMut::new(ResourceTable::new()),
5992            futures: AlwaysMut::new(Some(FuturesUnordered::new())),
5993            switch_item: None,
5994            next_switch_item: None,
5995            high_priority: VecDeque::new(),
5996            low_priority: VecDeque::new(),
5997            suspend_reason: None,
5998            worker: None,
5999            worker_item: None,
6000            global_error_context_ref_counts: BTreeMap::new(),
6001            interesting_tasks: 0,
6002            interesting_tasks_empty_waker: None,
6003            ready_for_concurrent_call_waker: None,
6004            event_loop_running: false,
6005            saved_next_switch_items: Vec::new(),
6006            #[cfg(feature = "task-group-hook")]
6007            task_group_hook: None,
6008        }
6009    }
6010}
6011
6012impl ConcurrentState {
6013    /// Take ownership of any fibers and futures owned by this object.
6014    ///
6015    /// This should be used when disposing of the `Store` containing this object
6016    /// in order to gracefully resolve any and all fibers using
6017    /// `StoreFiber::dispose`.  This is necessary to avoid possible
6018    /// use-after-free bugs due to fibers which may still have access to the
6019    /// `Store`.
6020    ///
6021    /// Additionally, the futures collected with this function should be dropped
6022    /// within a `tls::set` call, which will ensure than any futures closing
6023    /// over an `&Accessor` will have access to the store when dropped, allowing
6024    /// e.g. `WithAccessor[AndValue]` instances to be disposed of without
6025    /// panicking.
6026    ///
6027    /// Note that this will leave the object in an inconsistent and unusable
6028    /// state, so it should only be used just prior to dropping it.
6029    pub(crate) fn take_fibers_and_futures(
6030        &mut self,
6031        fibers: &mut Vec<StoreFiber<'static>>,
6032        futures: &mut Vec<FuturesUnordered<HostTaskFuture>>,
6033    ) {
6034        let mut items = Vec::new();
6035        for (_, entry) in self.table.get_mut().iter_mut() {
6036            if let Some(set) = entry.downcast_mut::<WaitableSet>() {
6037                for mode in mem::take(&mut set.waiting).into_values() {
6038                    match mode {
6039                        WaitMode::Fiber(fiber) => {
6040                            fibers.push(fiber);
6041                        }
6042                        WaitMode::Callback(_) => {}
6043                    }
6044                }
6045            } else if let Some(thread) = entry.downcast_mut::<GuestThread>() {
6046                if let GuestThreadState::Suspended(fiber) | GuestThreadState::Ready { fiber, .. } =
6047                    mem::replace(&mut thread.state, GuestThreadState::Completed)
6048                {
6049                    fibers.push(fiber);
6050                }
6051            } else if let Some(item) = entry.downcast_mut::<Option<WorkItem>>() {
6052                if let Some(item) = item.take() {
6053                    items.push(item);
6054                }
6055            }
6056        }
6057
6058        if let Some(fiber) = self.worker.take() {
6059            fibers.push(fiber);
6060        }
6061
6062        let mut handle_item = |item| match item {
6063            WorkItem::ResumeFiber { fiber, .. } => {
6064                fibers.push(fiber);
6065            }
6066            WorkItem::PushFuture(future) => {
6067                self.futures
6068                    .get_mut()
6069                    .as_mut()
6070                    .unwrap()
6071                    .push(future.into_inner());
6072            }
6073            WorkItem::ResumeThread { .. }
6074            | WorkItem::GuestCall { .. }
6075            | WorkItem::WorkerFunction(_) => {}
6076        };
6077
6078        for item in items {
6079            handle_item(item);
6080        }
6081        if let Some(item) = self.switch_item.take() {
6082            handle_item(item);
6083        }
6084        if let Some(item) = self.next_switch_item.take() {
6085            handle_item(item);
6086        }
6087        for item in mem::take(&mut self.high_priority) {
6088            handle_item(item);
6089        }
6090        for item in mem::take(&mut self.low_priority) {
6091            handle_item(item);
6092        }
6093        for item in mem::take(&mut self.saved_next_switch_items)
6094            .into_iter()
6095            .filter_map(|v| v)
6096        {
6097            handle_item(item);
6098        }
6099
6100        if let Some(them) = self.futures.get_mut().take() {
6101            futures.push(them);
6102        }
6103    }
6104
6105    #[cfg(feature = "gc")]
6106    pub(crate) fn trace_fiber_roots(
6107        &mut self,
6108        modules: &ModuleRegistry,
6109        unwind: &dyn Unwind,
6110        gc_roots_list: &mut GcRootsList,
6111    ) {
6112        let ConcurrentState {
6113            table,
6114            worker,
6115            switch_item,
6116            next_switch_item,
6117            high_priority,
6118            low_priority,
6119            saved_next_switch_items,
6120
6121            // TODO(cm-gc): This field contains `ValRaw`s, but they are never GC
6122            // references because the component model doesn't support GC yet. We
6123            // will need to trace these somehow when it does.
6124            futures: _,
6125
6126            // These fields do not contain GC references.
6127            worker_item: _,
6128            unforced_current_thread: _,
6129            deferred_host_call_context: _,
6130            suspend_reason: _,
6131            global_error_context_ref_counts: _,
6132            interesting_tasks: _,
6133            interesting_tasks_empty_waker: _,
6134            ready_for_concurrent_call_waker: _,
6135            event_loop_running: _,
6136            #[cfg(feature = "task-group-hook")]
6137                task_group_hook: _,
6138        } = self;
6139
6140        for (_, entry) in table.get_mut().iter_mut() {
6141            if let Some(set) = entry.downcast_mut::<WaitableSet>() {
6142                for mode in set.waiting.values_mut() {
6143                    match mode {
6144                        WaitMode::Fiber(fiber) => {
6145                            fiber.trace_gc_roots(modules, unwind, gc_roots_list);
6146                        }
6147                        WaitMode::Callback(_) => {}
6148                    }
6149                }
6150            } else if let Some(thread) = entry.downcast_mut::<GuestThread>() {
6151                if let GuestThreadState::Suspended(fiber) | GuestThreadState::Ready { fiber, .. } =
6152                    &mut thread.state
6153                {
6154                    fiber.trace_gc_roots(modules, unwind, gc_roots_list);
6155                }
6156            } else if let Some(Some(WorkItem::ResumeFiber { fiber, .. })) =
6157                entry.downcast_mut::<Option<WorkItem>>()
6158            {
6159                fiber.trace_gc_roots(modules, unwind, gc_roots_list);
6160            }
6161        }
6162
6163        if let Some(fiber) = worker {
6164            fiber.trace_gc_roots(modules, unwind, gc_roots_list);
6165        }
6166
6167        let mut handle_item = |item: &mut WorkItem| match item {
6168            WorkItem::ResumeFiber { fiber, .. } => {
6169                fiber.trace_gc_roots(modules, unwind, gc_roots_list);
6170            }
6171            WorkItem::PushFuture(_future) => {
6172                // TODO(cm-gc): once futures can contain GC roots, we will need
6173                // to trace them.
6174            }
6175            WorkItem::ResumeThread { .. }
6176            | WorkItem::GuestCall { .. }
6177            | WorkItem::WorkerFunction(_) => {}
6178        };
6179
6180        if let Some(item) = switch_item {
6181            handle_item(item);
6182        }
6183        if let Some(item) = next_switch_item {
6184            handle_item(item);
6185        }
6186        for item in high_priority {
6187            handle_item(item);
6188        }
6189        for item in low_priority {
6190            handle_item(item);
6191        }
6192        for item in saved_next_switch_items
6193            .iter_mut()
6194            .filter_map(|v| v.as_mut())
6195        {
6196            handle_item(item);
6197        }
6198    }
6199
6200    fn push<V: Send + Sync + 'static>(
6201        &mut self,
6202        value: V,
6203    ) -> Result<TableId<V>, ResourceTableError> {
6204        self.table.get_mut().push(value).map(TableId::from)
6205    }
6206
6207    fn get_mut<V: 'static>(&mut self, id: TableId<V>) -> Result<&mut V, ResourceTableError> {
6208        self.table.get_mut().get_mut(&Resource::from(id))
6209    }
6210
6211    pub fn add_child<T: 'static, U: 'static>(
6212        &mut self,
6213        child: TableId<T>,
6214        parent: TableId<U>,
6215    ) -> Result<(), ResourceTableError> {
6216        self.table
6217            .get_mut()
6218            .add_child(Resource::from(child), Resource::from(parent))
6219    }
6220
6221    pub fn remove_child<T: 'static, U: 'static>(
6222        &mut self,
6223        child: TableId<T>,
6224        parent: TableId<U>,
6225    ) -> Result<(), ResourceTableError> {
6226        self.table
6227            .get_mut()
6228            .remove_child(Resource::from(child), Resource::from(parent))
6229    }
6230
6231    fn delete<V: 'static>(&mut self, id: TableId<V>) -> Result<V, ResourceTableError> {
6232        self.table.get_mut().delete(Resource::from(id))
6233    }
6234
6235    fn push_future(&mut self, future: HostTaskFuture) {
6236        // Note that we can't directly push to `ConcurrentState::futures` here
6237        // since this may be called from a future that's being polled inside
6238        // `Self::poll_until`, which temporarily removes the `FuturesUnordered`
6239        // so it has exclusive access while polling it.  Therefore, we push a
6240        // work item to the "high priority" queue, which will actually push to
6241        // `ConcurrentState::futures` later.
6242        self.push_high_priority(WorkItem::PushFuture(AlwaysMut::new(future)));
6243    }
6244
6245    fn set_switch_item(&mut self, item: WorkItem) -> Result<()> {
6246        log::trace!("set switch item: {item:?}");
6247
6248        if self.switch_item.is_some() {
6249            bail_bug!("switch item already set");
6250        }
6251
6252        self.switch_item = Some(item);
6253
6254        Ok(())
6255    }
6256
6257    fn take_next_switch_item(&mut self) -> Result<()> {
6258        if let Some(item) = self.next_switch_item.take() {
6259            self.set_switch_item(item)?;
6260        }
6261        Ok(())
6262    }
6263
6264    /// Wake the first thread waiting on `set`, if any, so that it may receive
6265    /// one of the set's pending events.
6266    fn wake_waiter(&mut self, set: TableId<WaitableSet>) -> Result<()> {
6267        let Some((thread, mode)) = self.get_mut(set)?.waiting.pop_first() else {
6268            return Ok(());
6269        };
6270        let wake_on_cancel = self.get_mut(thread.thread)?.wake_on_cancel.take();
6271        assert!(wake_on_cancel.is_none() || wake_on_cancel == WakeOnCancel::Waiting(set));
6272
6273        let instance = self.get_mut(thread.task)?.instance;
6274        let item = match mode {
6275            WaitMode::Fiber(fiber) => WorkItem::ResumeFiber {
6276                instance,
6277                thread,
6278                fiber,
6279            },
6280            WaitMode::Callback(callback_instance) => WorkItem::GuestCall {
6281                instance,
6282                call: GuestCall {
6283                    thread,
6284                    kind: GuestCallKind::DeliverEvent {
6285                        instance: callback_instance,
6286                        set: Some(set),
6287                    },
6288                },
6289            },
6290        };
6291        self.push_high_priority(item);
6292        Ok(())
6293    }
6294
6295    fn push_high_priority(&mut self, item: WorkItem) {
6296        log::trace!("push high priority: {item:?}");
6297        self.high_priority.push_front(item);
6298    }
6299
6300    fn push_low_priority(&mut self, item: WorkItem) {
6301        log::trace!("push low priority: {item:?}");
6302        self.low_priority.push_front(item);
6303    }
6304
6305    fn push_work_item(&mut self, item: WorkItem, priority: Priority) -> Result<()> {
6306        match priority {
6307            Priority::Switch => self.set_switch_item(item)?,
6308            Priority::High => self.push_high_priority(item),
6309            Priority::Low => self.push_low_priority(item),
6310        }
6311
6312        Ok(())
6313    }
6314
6315    fn promote_instance_local_thread_work_item(
6316        &mut self,
6317        current_instance: RuntimeInstance,
6318    ) -> Result<bool> {
6319        log::trace!("promote thread work items for {current_instance:?}");
6320
6321        self.promote_work_item_matching(|item: &WorkItem| {
6322            let result = match item {
6323                WorkItem::ResumeThread { instance, .. }
6324                | WorkItem::ResumeFiber { instance, .. }
6325                | WorkItem::GuestCall { instance, .. } => *instance == current_instance,
6326                _ => false,
6327            };
6328
6329            log::trace!("candidate {item:?}: {result}");
6330            result
6331        })
6332    }
6333
6334    fn promote_thread_work_item(&mut self, thread: QualifiedThreadId) -> Result<bool> {
6335        self.promote_work_item_matching(|item: &WorkItem| match item {
6336            WorkItem::ResumeThread {
6337                thread: item_thread,
6338                ..
6339            }
6340            | WorkItem::GuestCall {
6341                call:
6342                    GuestCall {
6343                        thread: item_thread,
6344                        ..
6345                    },
6346                ..
6347            } => *item_thread == thread,
6348            _ => false,
6349        })
6350    }
6351
6352    fn promote_work_item_matching<F>(&mut self, mut predicate: F) -> Result<bool>
6353    where
6354        F: FnMut(&WorkItem) -> bool,
6355    {
6356        // Note the use of `.rev()` below to preserve ordering given that items
6357        // are popped from the back of the `VecDeque`s by `poll_until` and
6358        // pushed to the front by `push_{high,low}_priority`.
6359
6360        for item in mem::take(&mut self.high_priority).into_iter().rev() {
6361            if self.switch_item.is_none() && predicate(&item) {
6362                self.set_switch_item(item)?;
6363            } else {
6364                self.push_high_priority(item);
6365            }
6366        }
6367
6368        if self.switch_item.is_none() {
6369            for item in mem::take(&mut self.low_priority).into_iter().rev() {
6370                if self.switch_item.is_none() && predicate(&item) {
6371                    self.set_switch_item(item)?;
6372                } else {
6373                    self.push_low_priority(item);
6374                }
6375            }
6376        }
6377
6378        Ok(self.switch_item.is_some())
6379    }
6380
6381    /// Used by `ResourceTables` to acquire the current `CallContext` for the
6382    /// specified task.
6383    pub fn call_context(&mut self, task: Scope) -> Result<&mut CallContext> {
6384        match task {
6385            Scope::HostId(task) => {
6386                let task: TableId<HostTask> = TableId::new(task);
6387                Ok(&mut self.get_mut(task)?.call_context)
6388            }
6389            Scope::Id(task) => {
6390                let task: TableId<GuestTask> = TableId::new(task);
6391                Ok(&mut self.get_mut(task)?.call_context)
6392            }
6393        }
6394    }
6395
6396    pub(crate) fn deferred_host_call_context(&mut self) -> Option<&mut CallContext> {
6397        self.deferred_host_call_context.as_mut()
6398    }
6399
6400    fn futures_mut(&mut self) -> Result<&mut FuturesUnordered<HostTaskFuture>> {
6401        match self.futures.get_mut().as_mut() {
6402            Some(f) => Ok(f),
6403            None => bail_bug!("futures field of concurrent state is currently taken"),
6404        }
6405    }
6406
6407    pub(crate) fn table(&mut self) -> &mut ResourceTable {
6408        self.table.get_mut()
6409    }
6410
6411    fn debug_assert_deferred_host_invariant(&self) {
6412        debug_assert_eq!(
6413            self.deferred_host_call_context.is_some(),
6414            matches!(self.unforced_current_thread, CurrentThread::DeferredHost(_)),
6415            "a deferred host thread and call context must exist together",
6416        );
6417    }
6418
6419    fn materialize_host_task(&mut self) -> Result<CurrentThread> {
6420        self.debug_assert_deferred_host_invariant();
6421        let caller = match self.unforced_current_thread {
6422            CurrentThread::DeferredHost(caller) => caller,
6423            thread => return Ok(thread),
6424        };
6425
6426        // Push first so allocation failure leaves the deferred state intact.
6427        let task = HostTask::new(self, HostTaskState::CalleeStarted, caller)?;
6428        let task = self.push(task)?;
6429        let call_context = self
6430            .deferred_host_call_context
6431            .take()
6432            .expect("deferred host call context should be present");
6433        self.get_mut(task)
6434            .expect("newly inserted host task should be present")
6435            .call_context = call_context;
6436        self.unforced_current_thread = CurrentThread::Host(task);
6437        self.debug_assert_deferred_host_invariant();
6438        log::trace!("new host task materialized {task:?}");
6439        Ok(CurrentThread::Host(task))
6440    }
6441
6442    fn materialize_current_host_task_id(&mut self) -> Result<Option<TableId<HostTask>>> {
6443        match self.materialize_host_task()? {
6444            CurrentThread::Host(id) => Ok(Some(id)),
6445            CurrentThread::None => Ok(None),
6446            CurrentThread::Guest(_) => {
6447                bail_bug!("tried to materialize a host task id from a guest thread")
6448            }
6449            CurrentThread::DeferredHost(_) => {
6450                bail_bug!(
6451                    "current thread is a deferred host thread which should have been materialized"
6452                )
6453            }
6454        }
6455    }
6456
6457    pub(crate) fn materialize_current_scope(&mut self) -> Result<Scope> {
6458        match self.materialize_host_task()? {
6459            CurrentThread::Host(id) => Ok(Scope::HostId(id.rep())),
6460            _ => bail_bug!("current scope is not a deferred host scope"),
6461        }
6462    }
6463}
6464
6465/// Provide a type hint to compiler about the shape of a parameter lower
6466/// closure.
6467fn for_any_lower<
6468    F: FnOnce(&mut dyn VMStore, &mut [MaybeUninit<ValRaw>]) -> Result<()> + Send + Sync,
6469>(
6470    fun: F,
6471) -> F {
6472    fun
6473}
6474
6475/// Provide a type hint to compiler about the shape of a result lift closure.
6476fn for_any_lift<
6477    F: FnOnce(&mut dyn VMStore, &[ValRaw]) -> Result<Box<dyn Any + Send + Sync>> + Send + Sync,
6478>(
6479    fun: F,
6480) -> F {
6481    fun
6482}
6483
6484fn check_ambient_store(id: StoreId) {
6485    let message = "\
6486        `Future`s which depend on asynchronous component tasks, streams, or \
6487        futures to complete may only be polled from the event loop of the \
6488        store to which they belong.  Please use \
6489        `StoreContextMut::{run_concurrent,spawn}` to poll or await them.\
6490    ";
6491    tls::try_get(|store| {
6492        let matched = match store {
6493            tls::TryGet::Some(store) => store.id() == id,
6494            tls::TryGet::Taken | tls::TryGet::None => false,
6495        };
6496
6497        if !matched {
6498            panic!("{message}")
6499        }
6500    });
6501}
6502
6503fn unpack_callback_code(code: u32) -> (u32, u32) {
6504    (code & 0xF, code >> 4)
6505}
6506
6507/// Helper struct for packaging parameters to be passed to
6508/// `ComponentInstance::waitable_check` for calls to `waitable-set.wait` or
6509/// `waitable-set.poll`.
6510struct WaitableCheckParams {
6511    set: TableId<WaitableSet>,
6512    options: OptionsIndex,
6513    payload: u32,
6514}
6515
6516/// Indicates whether `ComponentInstance::waitable_check` is being called for
6517/// `waitable-set.wait` or `waitable-set.poll`.
6518enum WaitableCheck {
6519    Wait,
6520    Poll,
6521}
6522
6523/// Represents a guest task called from the host, prepared using `prepare_call`.
6524pub(crate) struct PreparedCall<R> {
6525    /// The guest export to be called
6526    handle: Func,
6527    /// The guest thread created by `prepare_call`
6528    thread: QualifiedThreadId,
6529    /// The number of lowered core Wasm parameters to pass to the call.
6530    param_count: usize,
6531    /// The `oneshot::Receiver` to which the result of the call will be
6532    /// delivered when it is available.
6533    rx: oneshot::Receiver<LiftedResult>,
6534    /// The instance that this call is prepared for.
6535    runtime_instance: RuntimeInstance,
6536    _phantom: PhantomData<R>,
6537}
6538
6539impl<R> PreparedCall<R> {
6540    /// Get a copy of the `TaskId` for this `PreparedCall`.
6541    pub(crate) fn task_id(&self) -> TaskId {
6542        TaskId {
6543            task: self.thread.task,
6544            runtime_instance: self.runtime_instance,
6545        }
6546    }
6547}
6548
6549/// Represents a task created by `prepare_call`.
6550pub(crate) struct TaskId {
6551    task: TableId<GuestTask>,
6552    runtime_instance: RuntimeInstance,
6553}
6554
6555impl TaskId {
6556    /// The host future for an async task was dropped. If the parameters have not been lowered yet,
6557    /// it is no longer valid to do so, as the lowering closure would see a dangling pointer. In this case,
6558    /// we delete the task eagerly. Otherwise, there may be running threads, or ones that are suspended
6559    /// and can be resumed by other tasks for this component, so we mark the future as dropped
6560    /// and delete the task when all threads are done.
6561    pub(crate) fn host_future_dropped(&self, store: &mut StoreOpaque) -> Result<()> {
6562        let task = store.concurrent_state_mut()?.get_mut(self.task)?;
6563        let delete = if !task.already_lowered_parameters() {
6564            store.cancel_guest_subtask_without_lowered_parameters(
6565                self.runtime_instance,
6566                self.task,
6567            )?;
6568            true
6569        } else {
6570            task.host_future_state = HostFutureState::Dropped;
6571            task.ready_to_delete()
6572        };
6573        if delete {
6574            Waitable::Guest(self.task).delete_from(store)?
6575        }
6576        Ok(())
6577    }
6578}
6579
6580/// Prepare a call to the specified exported Wasm function, providing functions
6581/// for lowering the parameters and lifting the result.
6582///
6583/// To enqueue the returned `PreparedCall` in the `ComponentInstance`'s event
6584/// loop, use `stage_call`.
6585pub(crate) fn prepare_call<T, R>(
6586    mut store: StoreContextMut<T>,
6587    handle: Func,
6588    param_count: usize,
6589    host_future_present: bool,
6590    lower_params: impl FnOnce(StoreContextMut<T>, &mut [MaybeUninit<ValRaw>]) -> Result<()>
6591    + Send
6592    + Sync
6593    + 'static,
6594    lift_result: impl FnOnce(&mut StoreOpaque, &[ValRaw]) -> Result<Box<dyn Any + Send + Sync>>
6595    + Send
6596    + Sync
6597    + 'static,
6598) -> Result<PreparedCall<R>> {
6599    if !store.0.may_enter() {
6600        bail!(Trap::CannotEnterComponent);
6601    }
6602
6603    let (options, _flags, ty, raw_options) = handle.abi_info(store.0);
6604
6605    let instance = handle.instance().id().get(store.0);
6606    let options = &instance.component().env_component().options[options];
6607    let ty = &instance.component().types()[ty];
6608    let async_typed = ty.async_;
6609    let async_lifted = raw_options.async_;
6610    let task_return_type = ty.results;
6611    let component_instance = raw_options.instance;
6612    let callback = options.callback.map(|i| instance.runtime_callback(i));
6613    let memory = options
6614        .memory()
6615        .map(|i| instance.runtime_memory(i))
6616        .map(SendSyncPtr::new);
6617    let string_encoding = options.string_encoding;
6618    let token = StoreToken::new(store.as_context_mut());
6619    let caller = store.0.materialize_host_task_id()?;
6620    let state = store.0.concurrent_state_mut()?;
6621
6622    let (tx, rx) = oneshot::channel();
6623
6624    let instance = handle.instance().runtime_instance(component_instance);
6625    let thread = GuestTask::new(
6626        state,
6627        Box::new(for_any_lower(move |store, params| {
6628            lower_params(token.as_context_mut(store), params)
6629        })),
6630        LiftResult {
6631            lift: Box::new(for_any_lift(move |store, result| {
6632                lift_result(store, result)
6633            })),
6634            ty: task_return_type,
6635            memory,
6636            string_encoding,
6637        },
6638        Caller::Host {
6639            tx: Some(tx),
6640            host_future_present,
6641            caller,
6642        },
6643        callback.map(|callback| {
6644            let callback = SendSyncPtr::new(callback);
6645            let instance = handle.instance();
6646            Box::new(move |store: &mut dyn VMStore, event, handle| {
6647                let store = token.as_context_mut(store);
6648                // SAFETY: Per the contract of `prepare_call`, the callback
6649                // will remain valid at least as long is this task exists.
6650                unsafe { instance.call_callback(store, callback, event, handle) }
6651            }) as CallbackFn
6652        }),
6653        instance,
6654        async_typed,
6655        async_lifted,
6656    )?;
6657
6658    Ok(PreparedCall {
6659        handle,
6660        thread,
6661        param_count,
6662        runtime_instance: instance,
6663        rx,
6664        _phantom: PhantomData,
6665    })
6666}
6667
6668pub(crate) struct StagedCall<R> {
6669    store: StoreId,
6670    rx: oneshot::Receiver<LiftedResult>,
6671    _marker: PhantomData<fn() -> R>,
6672    group: TaskGroupId,
6673}
6674
6675impl<R> StagedCall<R> {
6676    /// Queue a call previously prepared using `prepare_call` to be run as part of
6677    /// the associated `ComponentInstance`'s event loop.
6678    ///
6679    /// The returned future will resolve to the result once it is available, but
6680    /// must only be polled via the instance's event loop. See
6681    /// `StoreContextMut::run_concurrent` for details.
6682    pub(crate) fn new<T: 'static>(
6683        mut store: StoreContextMut<T>,
6684        prepared: PreparedCall<R>,
6685    ) -> Result<StagedCall<R>> {
6686        let PreparedCall {
6687            handle,
6688            thread,
6689            param_count,
6690            rx,
6691            ..
6692        } = prepared;
6693
6694        stage_call0(store.as_context_mut(), handle, thread, param_count)?;
6695
6696        Ok(StagedCall {
6697            store: store.0.id(),
6698            rx,
6699            _marker: PhantomData,
6700            group: store.0.concurrent_state_mut()?.get_mut(thread.task)?.group,
6701        })
6702    }
6703}
6704
6705impl<R> Future for StagedCall<R>
6706where
6707    R: 'static,
6708{
6709    type Output = Result<R>;
6710
6711    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
6712        check_ambient_store(self.store);
6713        Pin::new(&mut self.rx).poll(cx).map(|result| match result {
6714            Ok(r) => match r.downcast() {
6715                Ok(r) => Ok(*r),
6716                Err(_) => bail_bug!("wrong type of value produced"),
6717            },
6718            Err(oneshot::Canceled) => bail_bug!("channel erroneously dropped"),
6719        })
6720    }
6721}
6722
6723/// Queue a call previously prepared using `prepare_call` to be run as part of
6724/// the associated `ComponentInstance`'s event loop.
6725fn stage_call0<T: 'static>(
6726    store: StoreContextMut<T>,
6727    handle: Func,
6728    guest_thread: QualifiedThreadId,
6729    param_count: usize,
6730) -> Result<()> {
6731    let (_options, _, _ty, raw_options) = handle.abi_info(store.0);
6732    let is_concurrent = raw_options.async_;
6733    let callback = raw_options.callback;
6734    let instance = handle.instance();
6735    let callee = handle.lifted_core_func(store.0);
6736    let post_return = raw_options
6737        .post_return
6738        .map(|i| instance.id().get(store.0).runtime_post_return(i));
6739    let callback = callback.map(|i| {
6740        let instance = instance.id().get(store.0);
6741        SendSyncPtr::new(instance.runtime_callback(i))
6742    });
6743
6744    log::trace!("queueing call {guest_thread:?}");
6745
6746    // SAFETY: `callee`, `callback`, and `post_return` are valid pointers
6747    // (with signatures appropriate for this call) and will remain valid as
6748    // long as this instance is valid.
6749    unsafe {
6750        instance.stage_call(
6751            store,
6752            guest_thread,
6753            SendSyncPtr::new(callee),
6754            param_count,
6755            1,
6756            is_concurrent,
6757            callback,
6758            post_return.map(SendSyncPtr::new),
6759            true,
6760        )
6761    }
6762}
6763
6764#[cfg(all(test, feature = "cranelift", feature = "wat"))]
6765mod tests {
6766    use super::*;
6767    use crate::component::{Component, Linker};
6768    use crate::store::AsStoreOpaque;
6769    use crate::{Config, Engine};
6770
6771    fn host_subtask(
6772        state: HostTaskState,
6773        event: Option<Event>,
6774    ) -> Result<(Store<()>, Instance, TableId<HostTask>, u32)> {
6775        let mut config = Config::new();
6776        config.wasm_component_model_async(true);
6777        let engine = Engine::new(&config)?;
6778        let component = Component::new(&engine, "(component)")?;
6779        let mut store = Store::new(&engine, ());
6780        let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
6781        let store_opaque = store.as_store_opaque();
6782        let concurrent_state = store_opaque.concurrent_state_mut()?;
6783        // These tests create an isolated host task without a guest caller, so
6784        // give it its own task group instead of inheriting the caller's group.
6785        let group = concurrent_state.make_task_group()?;
6786        let task = concurrent_state.push(HostTask {
6787            common: WaitableCommon::default(),
6788            call_context: CallContext::default(),
6789            state,
6790            group,
6791        })?;
6792        let handle = store_opaque
6793            .instance_state(instance.runtime_instance(RuntimeComponentInstanceIndex::from_u32(0)))
6794            .handle_table()
6795            .subtask_insert_host(task.rep())?;
6796        let common = &mut store_opaque.concurrent_state_mut()?.get_mut(task)?.common;
6797        common.handle = Some(handle);
6798        common.event = event;
6799        Ok((store, instance, task, handle))
6800    }
6801
6802    #[test]
6803    fn host_subtask_drop_during_cancellation() -> Result<()> {
6804        for abort_completed in [false, true] {
6805            let (handle, future) = JoinHandle::run(future::pending::<()>());
6806            let mut future = pin!(future);
6807            let (mut store, instance, task, handle) =
6808                host_subtask(HostTaskState::CalleeRunning(handle), None)?;
6809            let store = store.as_store_opaque();
6810            let caller = RuntimeComponentInstanceIndex::from_u32(0);
6811            assert_eq!(
6812                instance.subtask_cancel(store, caller, true, handle)?,
6813                BLOCKED
6814            );
6815            if abort_completed {
6816                // Even after the abort resolves, the completion worker still
6817                // needs the task entry to lower the result and publish an event.
6818                assert!(matches!(
6819                    future
6820                        .as_mut()
6821                        .poll(&mut Context::from_waker(Waker::noop())),
6822                    Poll::Ready(None),
6823                ));
6824            }
6825            for async_ in [false, true] {
6826                let err = instance
6827                    .subtask_cancel(store, caller, async_, handle)
6828                    .unwrap_err();
6829                assert_eq!(err.downcast::<Trap>()?, Trap::SubtaskCancelAfterTerminal);
6830            }
6831            let err = instance.subtask_drop(store, caller, handle).unwrap_err();
6832            assert_eq!(err.downcast::<Trap>()?, Trap::SubtaskDropNotResolved);
6833            assert!(store.concurrent_state_mut()?.get_mut(task).is_ok());
6834        }
6835        Ok(())
6836    }
6837
6838    #[test]
6839    fn host_subtask_cancel_after_completion() -> Result<()> {
6840        for async_ in [false, true] {
6841            let (mut store, instance, task, handle) = host_subtask(
6842                HostTaskState::CalleeDone { cancelled: false },
6843                Some(Event::Subtask {
6844                    status: Status::Returned,
6845                }),
6846            )?;
6847            let store = store.as_store_opaque();
6848            let caller = RuntimeComponentInstanceIndex::from_u32(0);
6849            assert_eq!(
6850                instance.subtask_cancel(store, caller, async_, handle)?,
6851                Status::Returned as u32,
6852            );
6853            let err = instance
6854                .subtask_cancel(store, caller, async_, handle)
6855                .unwrap_err();
6856            assert_eq!(err.downcast::<Trap>()?, Trap::SubtaskCancelAfterTerminal);
6857            instance.subtask_drop(store, caller, handle)?;
6858            assert!(store.concurrent_state_mut()?.get_mut(task).is_err());
6859        }
6860        Ok(())
6861    }
6862
6863    #[test]
6864    fn host_subtask_drop_requires_terminal_event_delivery() -> Result<()> {
6865        for (cancelled, status) in [
6866            (false, Status::Returned),
6867            (true, Status::Returned),
6868            (true, Status::ReturnCancelled),
6869        ] {
6870            for delivered in [false, true] {
6871                let event = if delivered {
6872                    None
6873                } else {
6874                    Some(Event::Subtask { status })
6875                };
6876                let (mut store, instance, task, handle) =
6877                    host_subtask(HostTaskState::CalleeDone { cancelled }, event)?;
6878                let store = store.as_store_opaque();
6879                let result = instance.subtask_drop(
6880                    store,
6881                    RuntimeComponentInstanceIndex::from_u32(0),
6882                    handle,
6883                );
6884                if delivered {
6885                    result?;
6886                    assert!(store.concurrent_state_mut()?.get_mut(task).is_err());
6887                } else {
6888                    let err = result.unwrap_err();
6889                    assert_eq!(err.downcast::<Trap>()?, Trap::SubtaskDropNotResolved);
6890                    assert!(store.concurrent_state_mut()?.get_mut(task).is_ok());
6891                }
6892            }
6893        }
6894        Ok(())
6895    }
6896}