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