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