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