wasmtime/runtime/store/

gc.rs

//! GC-related methods for stores.

use crate::RootSet;
use crate::error::{Context, ensure};
use crate::hash_map::HashMap;
use crate::module::ModuleRegistry;
use crate::store::{
    Asyncness, AutoAssertNoGc, InstanceId, StoreOpaque, StoreResourceLimiter, yield_now,
};
use crate::type_registry::RegisteredType;
use crate::vm::{
    self, Backtrace, Frame, GcRootsList, GcStore, InstanceAllocationRequest, NopHasher,
    SendSyncPtr, StoreGcHostAllocTypes, TraceInfo, VMGcRef,
};
use crate::{
    Engine, ExnRef, GcHeapOutOfMemory, Result, Rooted, Store, StoreContextMut, ThrownException,
    bail,
};
use core::fmt;
use core::mem::ManuallyDrop;
use core::num::NonZeroU32;
use core::ops::{Deref, DerefMut};
use core::ptr::NonNull;
use wasmtime_environ::DefinedTagIndex;
use wasmtime_environ::packed_option::ReservedValue;

#[derive(Default)]
pub(crate) struct StoreGcData {
    gc_roots: RootSet,
    gc_roots_list: GcRootsList,
    // Types for which the embedder has created an allocator for.
    gc_host_alloc_types: StoreGcHostAllocTypes,
    /// Pending exception, if any. This is also a GC root, because it
    /// needs to be rooted somewhere between the time that a pending
    /// exception is set and the time that the handling code takes the
    /// exception object. We use this rooting strategy rather than a
    /// root in an `Err` branch of a `Result` on the host side because
    /// it is less error-prone with respect to rooting behavior. See
    /// `throw()`, `take_pending_exception()`,
    /// `peek_pending_exception()`, `has_pending_exception()`, and
    /// `catch()`.
    ///
    /// Also note that the underlying reference here is a `VMExnRef`, a
    /// refinement of `VMGcRef`, but rooting APIs right now make it difficult to
    /// work with that directly so this is stored as `VMGcRef` instead.
    pending_exception: Option<VMGcRef>,

    /// A store-local cache of engine-level subtype checks.
    ///
    /// Dynamic subtype checks (e.g. for `ref.cast` instructions) consult the
    /// engine's type registry, which is protected by a read-write lock.
    /// Testing locally shows that parallel readers take a significant
    /// performance hit when using this lock (as measured with `wasmtime
    /// serve`). Thus this cache is intended to serve as a fast path where the
    /// lock need not be hit at all.
    ///
    /// Note that the size of this cache is currently bounded to a fixed size.
    /// Once this cache is full then all future consultations which aren't in
    /// the cache go upstream to the engine itself (slow). In the future
    /// this'll either be removed entirely (see #13484) or will have some sort
    /// of LRU-like behavior.
    ///
    /// FIXME(#13484) this field is a temporary workaround for a "true
    /// solution" where supertypes are stored inline in a compiled-code-visible
    /// location which means that a libcall isn't needed at all (nor
    /// synchronization) to determine subtype/supertype relationships.
    subtype_check_cache: HashMap<u64, bool, NopHasher>,
}

impl<T> Store<T> {
    /// Perform garbage collection.
    ///
    /// Note that it is not required to actively call this function. GC will
    /// automatically happen according to various internal heuristics. This is
    /// provided if fine-grained control over the GC is desired.
    ///
    /// If you are calling this method after an attempted allocation failed, you
    /// may pass in the [`GcHeapOutOfMemory`][crate::GcHeapOutOfMemory] error.
    /// When you do so, this method will attempt to create enough space in the
    /// GC heap for that allocation, so that it will succeed on the next
    /// attempt.
    ///
    /// # Errors
    ///
    /// This method will fail if an [async limiter is
    /// configured](Store::limiter_async) in which case [`Store::gc_async`] must
    /// be used instead.
    pub fn gc(&mut self, why: Option<&crate::GcHeapOutOfMemory<()>>) -> Result<()> {
        StoreContextMut(&mut self.inner).gc(why)
    }

    /// Returns the current capacity of the GC heap in bytes, or 0 if the GC
    /// heap has not been initialized yet.
    pub fn gc_heap_capacity(&self) -> usize {
        self.inner.gc_heap_capacity()
    }

    /// Set an exception as the currently pending exception, and
    /// return an error that propagates the throw.
    ///
    /// This method takes an exception object and stores it in the
    /// `Store` as the currently pending exception. This is a special
    /// rooted slot that holds the exception as long as it is
    /// propagating. This method then returns a `ThrownException`
    /// error, which is a special type that indicates a pending
    /// exception exists. When this type propagates as an error
    /// returned from a Wasm-to-host call, the pending exception is
    /// thrown within the Wasm context, and either caught or
    /// propagated further to the host-to-Wasm call boundary. If an
    /// exception is thrown out of Wasm (or across Wasm from a
    /// hostcall) back to the host-to-Wasm call boundary, *that*
    /// invocation returns a `ThrownException`, and the pending
    /// exception slot is again set. In other words, the
    /// `ThrownException` error type should propagate upward exactly
    /// and only when a pending exception is set.
    ///
    /// To take the pending exception, use [`Self::take_pending_exception`].
    ///
    /// This method is parameterized over `R` for convenience, but
    /// will always return an `Err`.
    ///
    /// If there is already a pending exception in the store then the previous
    /// one will be overwritten.
    ///
    /// # Errors
    ///
    /// This method will return an error if `exception` is unrooted. Otherwise
    /// this method will always return `ThrownException`.
    pub fn throw<R>(&mut self, exception: Rooted<ExnRef>) -> Result<R> {
        self.inner.throw_impl(exception)
    }

    /// Take the currently pending exception, if any, and return it,
    /// removing it from the "pending exception" slot.
    ///
    /// If there is no pending exception, returns `None`.
    ///
    /// Note: the returned exception is a LIFO root (see
    /// [`crate::Rooted`]), rooted in the current handle scope. Take
    /// care to ensure that it is re-rooted or otherwise does not
    /// escape this scope! It is usually best to allow an exception
    /// object to be rooted in the store's "pending exception" slot
    /// until the final consumer has taken it, rather than root it and
    /// pass it up the callstack in some other way.
    ///
    /// This method is useful to implement ad-hoc exception plumbing
    /// in various ways, but for the most idiomatic handling, see
    /// [`StoreContextMut::throw`].
    pub fn take_pending_exception(&mut self) -> Option<Rooted<ExnRef>> {
        self.inner.take_pending_exception_rooted()
    }
}

impl<'a, T> StoreContextMut<'a, T> {
    /// Perform garbage collection.
    ///
    /// Same as [`Store::gc`].
    pub fn gc(&mut self, why: Option<&GcHeapOutOfMemory<()>>) -> Result<()> {
        let (mut limiter, store) = self.0.validate_sync_resource_limiter_and_store_opaque()?;
        vm::assert_ready(store.gc(
            limiter.as_mut(),
            None,
            why.map(|e| e.bytes_needed()),
            Asyncness::No,
        ))?;
        Ok(())
    }

    /// Set an exception as the currently pending exception, and
    /// return an error that propagates the throw.
    ///
    /// See [`Store::throw`] for more details.
    #[cfg(feature = "gc")]
    pub fn throw<R>(&mut self, exception: Rooted<ExnRef>) -> Result<R> {
        self.0.inner.throw_impl(exception)
    }

    /// Take the currently pending exception, if any, and return it,
    /// removing it from the "pending exception" slot.
    ///
    /// See [`Store::take_pending_exception`] for more details.
    #[cfg(feature = "gc")]
    pub fn take_pending_exception(&mut self) -> Option<Rooted<ExnRef>> {
        self.0.inner.take_pending_exception_rooted()
    }
}

#[derive(Debug)]
struct GcHeapGrowthFailed;

impl fmt::Display for GcHeapGrowthFailed {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("GC heap growth failed")
    }
}

impl core::error::Error for GcHeapGrowthFailed {}

impl StoreOpaque {
    /// Perform any growth or GC needed to allocate `bytes_needed` bytes.
    ///
    /// Note that even when this function returns it is not guaranteed
    /// that a GC allocation of size `bytes_needed` will succeed. Growing the GC
    /// heap could fail, and then performing a collection could succeed but
    /// might not free up enough space. Therefore, callers should not assume
    /// that a retried allocation will always succeed.
    ///
    /// The `root` argument passed in is considered a root for this GC operation
    /// and its new value is returned as well.
    pub(crate) async fn gc(
        &mut self,
        limiter: Option<&mut StoreResourceLimiter<'_>>,
        root: Option<VMGcRef>,
        bytes_needed: Option<u64>,
        asyncness: Asyncness,
    ) -> Result<Option<VMGcRef>> {
        let mut scope = crate::OpaqueRootScope::new(self);
        scope.trim_gc_liveness_flags(true);
        let store_id = scope.id();
        let root = root.map(|r| scope.gc_roots_mut().push_lifo_root(store_id, r));

        scope
            .collect_and_maybe_grow_gc_heap(limiter, bytes_needed, asyncness)
            .await?;

        Ok(root.map(|r| {
            let r = r
                .get_gc_ref(&scope)
                .expect("still in scope")
                .unchecked_copy();
            scope.clone_gc_ref(&r)
        }))
    }

    // This lives on the Store because it must simultaneously borrow
    // `gc_store` and `gc_roots`, and is invoked from other modules to
    // which we do not want to expose the raw fields for piecewise
    // borrows.
    pub(crate) fn trim_gc_liveness_flags(&mut self, eager: bool) {
        if let Some(gc_store) = self.gc_store.as_mut() {
            self.gc_data.gc_roots.trim_liveness_flags(gc_store, eager);
        }
    }

    /// Helper invoked as part of `gc`, whose purpose is to GC and
    /// maybe grow for a pending allocation of a given size.
    async fn collect_and_maybe_grow_gc_heap(
        &mut self,
        limiter: Option<&mut StoreResourceLimiter<'_>>,
        bytes_needed: Option<u64>,
        asyncness: Asyncness,
    ) -> Result<()> {
        log::trace!("collect_and_maybe_grow_gc_heap(bytes_needed = {bytes_needed:#x?})");
        self.do_gc(asyncness).await?;
        if let Some(n) = bytes_needed
            && n > u64::try_from(self.gc_heap_capacity())?.saturating_sub(
                self.gc_store.as_ref().map_or(0, |gc| {
                    u64::try_from(gc.last_post_gc_allocated_bytes.unwrap_or(0)).unwrap()
                }),
            )
        {
            if let Err(e) = self.grow_gc_heap(limiter, n, asyncness).await {
                if e.is::<GcHeapGrowthFailed>() {
                    log::trace!("ignoring GC heap growth failure: {e}");
                } else {
                    return Err(e);
                }
            }
        }
        Ok(())
    }

    /// Attempt to grow the GC heap by `bytes_needed` bytes.
    ///
    /// Returns an error if growing the GC heap fails.
    pub(crate) async fn grow_gc_heap(
        &mut self,
        limiter: Option<&mut StoreResourceLimiter<'_>>,
        bytes_needed: u64,
        asyncness: Asyncness,
    ) -> Result<()> {
        log::trace!("Attempting to grow the GC heap by at least {bytes_needed:#x} bytes");

        if bytes_needed == 0 {
            return Ok(());
        }

        // If the GC heap needs a collection before growth (e.g. the copying
        // collector's active space is the second half), do a GC first.
        if self
            .gc_store
            .as_ref()
            .map_or(false, |gc| gc.gc_heap.needs_gc_before_next_growth())
        {
            self.do_gc(asyncness).await?;
            debug_assert!(
                !self
                    .gc_store
                    .as_ref()
                    .map_or(false, |gc| gc.gc_heap.needs_gc_before_next_growth()),
                "needs_gc_before_next_growth should return false after a GC"
            );
        }

        let page_size = self.engine().tunables().gc_heap_memory_type().page_size();

        // Take the GC heap's underlying memory out of the GC heap, attempt to
        // grow it, then replace it.
        let mut heap = TakenGcHeap::new(self);

        let current_size_in_bytes = u64::try_from(heap.memory.byte_size())?;
        let current_size_in_pages = current_size_in_bytes / page_size;

        // Aim to double the heap size, amortizing the cost of growth.
        let doubled_size_in_pages = current_size_in_pages.saturating_mul(2);
        assert!(doubled_size_in_pages >= current_size_in_pages);
        let delta_pages_for_doubling = doubled_size_in_pages - current_size_in_pages;

        // When doubling our size, saturate at the maximum memory size in pages.
        //
        // TODO: we should consult the instance allocator for its configured
        // maximum memory size, if any, rather than assuming the index
        // type's maximum size.
        let max_size_in_bytes = 1 << 32;
        let max_size_in_pages = max_size_in_bytes / page_size;
        let delta_to_max_size_in_pages = max_size_in_pages - current_size_in_pages;
        let delta_pages_for_alloc = delta_pages_for_doubling.min(delta_to_max_size_in_pages);

        // But always make sure we are attempting to grow at least as many pages
        // as needed by the requested allocation. This must happen *after* the
        // max-size saturation, so that if we are at the max already, we do not
        // succeed in growing by zero delta pages, and then return successfully
        // to our caller, who would be assuming that there is now capacity for
        // their allocation.
        let pages_needed = bytes_needed.div_ceil(page_size);
        assert!(pages_needed > 0);
        let delta_pages_for_alloc = delta_pages_for_alloc.max(pages_needed);
        assert!(delta_pages_for_alloc > 0);

        // Safety: we pair growing the GC heap with updating its associated
        // `VMMemoryDefinition` in the `VMStoreContext` immediately
        // afterwards.
        unsafe {
            heap.memory
                .grow(delta_pages_for_alloc, limiter)
                .await
                .context(GcHeapGrowthFailed)?
                .ok_or(GcHeapGrowthFailed)?;
        }
        *heap.store.vm_store_context.gc_heap.get_mut() = heap.memory.vmmemory();

        let new_size_in_bytes = u64::try_from(heap.memory.byte_size())?;
        assert!(new_size_in_bytes > current_size_in_bytes);
        heap.delta_bytes_grown = new_size_in_bytes - current_size_in_bytes;
        let delta_bytes_for_alloc = delta_pages_for_alloc.checked_mul(page_size).unwrap();
        assert!(
            heap.delta_bytes_grown >= delta_bytes_for_alloc,
            "{} should be greater than or equal to {delta_bytes_for_alloc}",
            heap.delta_bytes_grown,
        );
        log::trace!(
            "  -> grew GC heap by {:#x} bytes: new size is {new_size_in_bytes:#x} bytes",
            heap.delta_bytes_grown
        );
        return Ok(());

        struct TakenGcHeap<'a> {
            store: &'a mut StoreOpaque,
            memory: ManuallyDrop<vm::Memory>,
            delta_bytes_grown: u64,
        }

        impl<'a> TakenGcHeap<'a> {
            fn new(store: &'a mut StoreOpaque) -> TakenGcHeap<'a> {
                TakenGcHeap {
                    memory: ManuallyDrop::new(store.unwrap_gc_store_mut().gc_heap.take_memory()),
                    store,
                    delta_bytes_grown: 0,
                }
            }
        }

        impl Drop for TakenGcHeap<'_> {
            fn drop(&mut self) {
                // SAFETY: this `Drop` guard ensures that this has exclusive
                // ownership of fields and is thus safe to take `self.memory`.
                // Additionally for `replace_memory` the memory was previously
                // taken when this was created so it should be safe to place
                // back inside the GC heap.
                unsafe {
                    self.store.unwrap_gc_store_mut().gc_heap.replace_memory(
                        ManuallyDrop::take(&mut self.memory),
                        self.delta_bytes_grown,
                    );
                }
            }
        }
    }

    fn replace_gc_zeal_alloc_counter(
        &mut self,
        new_value: Option<NonZeroU32>,
    ) -> Option<NonZeroU32> {
        if let Some(gc_store) = &mut self.gc_store {
            gc_store.replace_gc_zeal_alloc_counter(new_value)
        } else {
            None
        }
    }

    /// Attempt an allocation, if it fails due to GC OOM, apply the
    /// grow-or-collect heuristic and retry.
    ///
    /// The heuristic is:
    /// - If the last post-collection heap usage is less than half the current
    ///   capacity, collect first, then retry. If that still fails, grow and
    ///   retry one final time.
    /// - Otherwise, grow first and retry.
    pub(crate) async fn retry_after_gc_async<T, U>(
        &mut self,
        mut limiter: Option<&mut StoreResourceLimiter<'_>>,
        value: T,
        asyncness: Asyncness,
        alloc_func: impl Fn(&mut Self, T) -> Result<U>,
    ) -> Result<U>
    where
        T: Send + Sync + 'static,
    {
        self.ensure_gc_store(limiter.as_deref_mut()).await?;

        match alloc_func(self, value) {
            Ok(x) => Ok(x),
            Err(e) => match e.downcast::<crate::GcHeapOutOfMemory<T>>() {
                Ok(oom) => {
                    log::trace!("Got GC heap OOM: {oom}");

                    let (value, oom) = oom.take_inner();
                    let bytes_needed = oom.bytes_needed();

                    let mut store = WithoutGcZealAllocCounter::new(self);

                    let gc_heap_capacity = store
                        .gc_store
                        .as_ref()
                        .map_or(0, |gc_store| gc_store.gc_heap_capacity());
                    let last_gc_heap_usage = store.gc_store.as_ref().map_or(0, |gc_store| {
                        gc_store.last_post_gc_allocated_bytes.unwrap_or(0)
                    });

                    if should_collect_first(bytes_needed, gc_heap_capacity, last_gc_heap_usage) {
                        log::trace!(
                            "Collecting first, then retrying; growing GC heap if collecting didn't \
                             free up enough space, then retrying again"
                        );
                        store
                            .gc(limiter.as_deref_mut(), None, None, asyncness)
                            .await?;

                        match alloc_func(&mut store, value) {
                            Ok(x) => Ok(x),
                            Err(e) => match e.downcast::<crate::GcHeapOutOfMemory<T>>() {
                                Ok(oom2) => {
                                    // Collection wasn't enough; grow and try
                                    // one final time.
                                    let (value, _) = oom2.take_inner();
                                    // Ignore error; we'll get one from
                                    // `alloc_func` below if growth failed and
                                    // failure to grow was fatal.
                                    let _ =
                                        store.grow_gc_heap(limiter, bytes_needed, asyncness).await;

                                    alloc_func(&mut store, value)
                                }
                                Err(e) => Err(e),
                            },
                        }
                    } else {
                        log::trace!(
                            "Grow GC heap first, collecting if growth failed, then retrying"
                        );

                        if let Err(e) = store
                            .grow_gc_heap(limiter.as_deref_mut(), bytes_needed.max(1), asyncness)
                            .await
                        {
                            log::trace!("growing GC heap failed: {e}");
                            store.gc(limiter, None, None, asyncness).await?;
                        }

                        alloc_func(&mut store, value)
                    }
                }
                Err(e) => Err(e),
            },
        }
    }

    /// Set a pending exception.
    ///
    /// The `exnref` is cloned internally and held on this store to be fetched
    /// later by an unwind. This method does *not* set up an unwind request on
    /// the TLS call state; that must be done separately.
    ///
    /// GC barriers are not required by the caller of this function.
    pub(crate) fn set_pending_exception(&mut self, exnref: &VMGcRef) -> crate::Error {
        debug_assert!(exnref.is_exnref(&*self.unwrap_gc_store_mut().gc_heap));
        let gc_store = self.gc_store.as_mut().unwrap();
        match gc_store.write_gc_ref(&mut self.gc_data.pending_exception, Some(exnref)) {
            Ok(()) => ThrownException.into(),
            Err(e) => e,
        }
    }

    /// Takes the pending exception from this store, if any, and exposes it to
    /// WebAssembly, returning the raw representation.
    pub(crate) fn expose_pending_exception_to_wasm(&mut self) -> Option<NonZeroU32> {
        let exnref = self.gc_data.pending_exception.take()?;
        let gc_store = self.unwrap_gc_store_mut();
        debug_assert!(exnref.is_exnref(&*gc_store.gc_heap));
        Some(gc_store.expose_gc_ref_to_wasm(exnref).unwrap())
    }

    /// Takes the pending exception of the store, yielding ownership of its
    /// reference to the `Rooted` that's returned.
    fn take_pending_exception_rooted(&mut self) -> Option<Rooted<ExnRef>> {
        let vmexnref = self.gc_data.pending_exception.take()?;
        debug_assert!(vmexnref.is_exnref(&*self.unwrap_gc_store().gc_heap));
        let mut nogc = AutoAssertNoGc::new(self);
        Some(Rooted::new(&mut nogc, vmexnref))
    }

    /// Returns the (instance,tag) pair that the pending exception in this
    /// store, if any, references.
    pub(crate) fn pending_exception_tag_and_instance(
        &mut self,
    ) -> Option<(InstanceId, DefinedTagIndex)> {
        let pending_exnref = self.gc_data.pending_exception.as_ref()?.unchecked_copy();
        debug_assert!(pending_exnref.is_exnref(&*self.unwrap_gc_store_mut().gc_heap));
        let mut store = AutoAssertNoGc::new(self);

        // Note that if the GC heap is corrupt this will return an error, and in
        // such as situation we return `None` here pretending that there's no
        // pending exception. This defers the GC heap corruption to get detected
        // later. This method is primarily called right now to determine if
        // there's a handler for an exception, and by returning `None` here this
        // turns into just any old embedder error.
        pending_exnref.into_exnref_unchecked().tag(&mut store).ok()
    }

    /// Get an owned rooted reference to the pending exception,
    /// without taking it off the store.
    #[cfg(feature = "debug")]
    pub(crate) fn pending_exception_owned_rooted(
        &mut self,
    ) -> Result<Option<crate::OwnedRooted<ExnRef>>, crate::OutOfMemory> {
        let pending = match &self.gc_data.pending_exception {
            Some(r) => r,
            None => return Ok(None),
        };
        let cloned = self.gc_store.as_mut().unwrap().clone_gc_ref(pending);
        let mut nogc = AutoAssertNoGc::new(self);
        Ok(Some(crate::OwnedRooted::new(&mut nogc, cloned)?))
    }

    /// Stores `exception` within the store to later get thrown.
    ///
    /// Delegates to `self.set_pending_exception` after accessing the internal
    /// exception pointer.
    fn throw_impl<R>(&mut self, exception: Rooted<ExnRef>) -> Result<R> {
        let exception = exception.try_gc_ref(self)?.unchecked_copy();
        Err(self.set_pending_exception(&exception))
    }

    /// Helper method to require that a `GcStore` was previously allocated for
    /// this store, failing if it has not yet been allocated.
    ///
    /// Note that this should only be used in a context where allocation of a
    /// `GcStore` is sure to have already happened prior, otherwise this may
    /// return a confusing error to embedders which is a bug in Wasmtime.
    ///
    /// Some situations where it's safe to call this method:
    ///
    /// * There's already a non-null and non-i31 `VMGcRef` in scope. By existing
    ///   this shows proof that the `GcStore` was previously allocated.
    /// * During instantiation and instance's `needs_gc_heap` flag will be
    ///   handled and instantiation will automatically create a GC store.
    #[inline]
    pub(crate) fn require_gc_store(&self) -> Result<&GcStore> {
        match &self.gc_store {
            Some(gc_store) => Ok(gc_store),
            None => bail!("GC heap not initialized yet"),
        }
    }

    /// Same as [`Self::require_gc_store`], but mutable.
    #[inline]
    pub(crate) fn require_gc_store_mut(&mut self) -> Result<&mut GcStore> {
        match &mut self.gc_store {
            Some(gc_store) => Ok(gc_store),
            None => bail!("GC heap not initialized yet"),
        }
    }

    /// Returns the current capacity of the GC heap in bytes, or 0 if the GC
    /// heap has not been initialized yet.
    pub(crate) fn gc_heap_capacity(&self) -> usize {
        match self.gc_store.as_ref() {
            Some(gc_store) => gc_store.gc_heap_capacity(),
            None => 0,
        }
    }

    async fn do_gc(&mut self, asyncness: Asyncness) -> Result<()> {
        // If the GC heap hasn't been initialized, there is nothing to collect.
        if self.gc_store.is_none() {
            return Ok(());
        }

        if log::log_enabled!(log::Level::Trace) {
            let gc_store = self.gc_store.as_ref().unwrap();
            let capacity = gc_store.gc_heap_capacity();
            let live_set_size = gc_store.last_post_gc_allocated_bytes.unwrap_or(0);
            let utilization = live_set_size as f64 / capacity as f64 * 100.0;
            log::trace!(
                "============ Begin GC ===========\n\
                 \t          GC heap capacity = {capacity:#010x} bytes\n\
                 \tlast post-GC live-set size = {live_set_size:#010x} bytes\n\
                 \t       GC heap utilization = {utilization:.02}%",
            );
        }

        // Take the GC roots out of `self` so we can borrow it mutably but still
        // call mutable methods on `self`.
        let mut roots = core::mem::take(&mut self.gc_data.gc_roots_list);

        self.trace_roots(&mut roots, asyncness).await;
        self.gc_store
            .as_mut()
            .unwrap()
            .gc(
                asyncness,
                unsafe { roots.iter() },
                &self.modules,
                &self.gc_data.gc_host_alloc_types,
                // TODO: Once `Config` has an optional `AsyncFn` field for
                // yielding to the current async runtime
                // (e.g. `tokio::task::yield_now`), use that if set; otherwise
                // fall back to the runtime-agnostic code.
                yield_now,
            )
            .await?;

        // Restore the GC roots for the next GC.
        roots.clear();
        self.gc_data.gc_roots_list = roots;

        if log::log_enabled!(log::Level::Trace) {
            let gc_store = self.gc_store.as_ref().unwrap();
            let capacity = gc_store.gc_heap_capacity();
            let live_set_size = gc_store.last_post_gc_allocated_bytes.unwrap_or(0);
            let utilization = live_set_size as f64 / capacity as f64 * 100.0;
            log::trace!(
                "============ End GC ===========\n\
                 \t     GC heap capacity = {capacity:#010x} bytes\n\
                 \tpost-GC live-set size = {live_set_size:#010x} bytes\n\
                 \t  GC heap utilization = {utilization:.02}%",
            );
        }
        Ok(())
    }

    async fn trace_roots(&mut self, gc_roots_list: &mut GcRootsList, asyncness: Asyncness) {
        log::trace!("Begin trace GC roots");

        // We shouldn't have any leftover, stale GC roots.
        assert!(gc_roots_list.is_empty());

        self.trace_wasm_stack_roots(gc_roots_list);
        if asyncness != Asyncness::No {
            self.yield_now().await;
        }

        #[cfg(feature = "stack-switching")]
        {
            self.trace_wasm_continuation_roots(gc_roots_list);
            if asyncness != Asyncness::No {
                self.yield_now().await;
            }
        }

        self.trace_vmctx_roots(gc_roots_list);
        if asyncness != Asyncness::No {
            self.yield_now().await;
        }

        self.trace_instance_roots(gc_roots_list);
        if asyncness != Asyncness::No {
            self.yield_now().await;
        }

        self.trace_user_roots(gc_roots_list);
        if asyncness != Asyncness::No {
            self.yield_now().await;
        }

        self.trace_pending_exception_roots(gc_roots_list);

        log::trace!("End trace GC roots")
    }

    pub(crate) fn trace_wasm_stack_frame(
        modules: &ModuleRegistry,
        gc_roots_list: &mut GcRootsList,
        frame: Frame,
    ) {
        let pc = frame.pc();
        debug_assert!(pc != 0, "we should always get a valid PC for Wasm frames");

        let fp = frame.fp() as *mut usize;
        debug_assert!(
            !fp.is_null(),
            "we should always get a valid frame pointer for Wasm frames"
        );

        let (store_code, offset) = modules
            .store_code_by_pc(pc)
            .expect("should have store code for Wasm frame");
        let offset = u32::try_from(offset).unwrap();

        let stack_map =
            wasmtime_environ::StackMap::lookup(offset, store_code.code_memory().stack_map_data());

        if let Some(stack_map) = stack_map {
            log::trace!(
                "We have a stack map that maps {} bytes in this Wasm frame",
                stack_map.frame_size()
            );

            let sp = unsafe { stack_map.sp(fp) };
            for stack_slot in unsafe { stack_map.live_gc_refs(sp) } {
                unsafe {
                    Self::trace_wasm_stack_slot(gc_roots_list, stack_slot);
                }
            }
        }

        #[cfg(feature = "debug")]
        if let Some(frame_table) = store_code.code_memory().frame_table() {
            for stack_slot in crate::debug::gc_refs_in_frame(frame_table, offset, fp) {
                unsafe {
                    Self::trace_wasm_stack_slot(gc_roots_list, stack_slot);
                }
            }
        }
    }

    unsafe fn trace_wasm_stack_slot(gc_roots_list: &mut GcRootsList, stack_slot: *mut u32) {
        let raw: u32 = unsafe { core::ptr::read(stack_slot) };
        log::trace!("Stack slot @ {stack_slot:p} = {raw:#x}");

        let gc_ref = vm::VMGcRef::from_raw_u32(raw);
        if gc_ref.is_some() {
            unsafe {
                gc_roots_list
                    .add_wasm_stack_root(SendSyncPtr::new(NonNull::new(stack_slot).unwrap()));
            }
        }
    }

    fn trace_wasm_stack_roots(&mut self, gc_roots_list: &mut GcRootsList) {
        log::trace!("Begin trace GC roots :: Wasm stack");

        Backtrace::trace(self, |frame| {
            Self::trace_wasm_stack_frame(self.modules(), gc_roots_list, frame);
            core::ops::ControlFlow::Continue(())
        });

        #[cfg(feature = "component-model-async")]
        if self.concurrency_support() {
            let unwind = self.unwinder();
            let StoreOpaque {
                modules,
                store_data,
                ..
            } = self;
            store_data
                .components
                .task_state_mut()
                .concurrent_state_mut()
                .trace_fiber_roots(modules, unwind, gc_roots_list);
        }

        log::trace!("End trace GC roots :: Wasm stack");
    }

    #[cfg(feature = "stack-switching")]
    fn trace_wasm_continuation_roots(&mut self, gc_roots_list: &mut GcRootsList) {
        use crate::vm::VMStackState;

        log::trace!("Begin trace GC roots :: continuations");

        for continuation in &self.continuations {
            let state = continuation.common_stack_information.state;

            // FIXME(frank-emrich) In general, it is not enough to just trace
            // through the stacks of continuations; we also need to look through
            // their `cont.bind` arguments. However, we don't currently have
            // enough RTTI information to check if any of the values in the
            // buffers used by `cont.bind` are GC values. As a workaround, note
            // that we currently disallow cont.bind-ing GC values altogether.
            // This way, it is okay not to check them here.
            match state {
                VMStackState::Suspended => {
                    Backtrace::trace_suspended_continuation(self, continuation.deref(), |frame| {
                        Self::trace_wasm_stack_frame(self.modules(), gc_roots_list, frame);
                        core::ops::ControlFlow::Continue(())
                    });
                }
                VMStackState::Running => {
                    // Handled by `trace_wasm_stack_roots`.
                }
                VMStackState::Parent => {
                    // We don't know whether our child is suspended or running, but in
                    // either case things should be handled correctly when traversing
                    // further along in the chain, nothing required at this point.
                }
                VMStackState::Fresh | VMStackState::Returned | VMStackState::Trapped => {
                    // Fresh/terminal continuations have no live GC values on their stack.
                }
            }
        }

        log::trace!("End trace GC roots :: continuations");
    }

    fn trace_vmctx_roots(&mut self, gc_roots_list: &mut GcRootsList) {
        log::trace!("Begin trace GC roots :: vmctx");
        self.for_each_global(|store, global| global.trace_root(store, gc_roots_list));
        self.for_each_table(|store, table| table.trace_roots(store, gc_roots_list));
        log::trace!("End trace GC roots :: vmctx");
    }

    fn trace_instance_roots(&mut self, gc_roots_list: &mut GcRootsList) {
        log::trace!("Begin trace GC roots :: instance");
        for (_id, instance) in &mut self.instances {
            // SAFETY: the instance's GC roots will remain valid for the
            // duration of this GC cycle.
            unsafe {
                instance
                    .handle
                    .get_mut()
                    .trace_element_segment_roots(gc_roots_list);
            }
        }
        log::trace!("End trace GC roots :: instance");
    }

    fn trace_user_roots(&mut self, gc_roots_list: &mut GcRootsList) {
        log::trace!("Begin trace GC roots :: user");
        self.gc_data.gc_roots.trace_roots(gc_roots_list);
        log::trace!("End trace GC roots :: user");
    }

    fn trace_pending_exception_roots(&mut self, gc_roots_list: &mut GcRootsList) {
        log::trace!("Begin trace GC roots :: pending exception");
        if let Some(pending_exception) = self.gc_data.pending_exception.as_mut() {
            unsafe {
                gc_roots_list.add_vmgcref_root(pending_exception.into(), "Pending exception");
            }
        }
        log::trace!("End trace GC roots :: pending exception");
    }

    /// Insert a host-allocated GC type into this store.
    ///
    /// This makes it suitable for the embedder to allocate instances of this
    /// type in this store, and we don't have to worry about the type being
    /// reclaimed (since it is possible that none of the Wasm modules in this
    /// store are holding it alive).
    ///
    /// # Panics
    ///
    /// Panics if `ty` was not registered with this store's engine. The types
    /// here are keyed by their `VMSharedTypeIndex`, which only means anything
    /// within the engine that assigned it.
    pub(crate) fn insert_gc_host_alloc_type(&mut self, ty: RegisteredType) {
        assert!(
            Engine::same(self.engine(), ty.engine()),
            "type used with wrong engine"
        );
        let trace_info = ty.layout().map(TraceInfo::new);
        self.gc_data
            .gc_host_alloc_types
            .insert(ty.index(), (ty, trace_info));
    }

    /// Performs a lazy allocation of the `GcStore` within this store, returning
    /// the previous allocation if it's already present.
    ///
    /// This method will, if necessary, allocate a new `GcStore` -- linear
    /// memory and all. This is a blocking operation due to
    /// `ResourceLimiterAsync` which means that this should only be executed
    /// in a fiber context at this time.
    #[inline]
    pub(crate) async fn ensure_gc_store(
        &mut self,
        limiter: Option<&mut StoreResourceLimiter<'_>>,
    ) -> Result<&mut GcStore> {
        if self.gc_store.is_some() {
            return Ok(self.gc_store.as_mut().unwrap());
        }
        self.allocate_gc_store(limiter).await
    }

    #[inline(never)]
    async fn allocate_gc_store(
        &mut self,
        limiter: Option<&mut StoreResourceLimiter<'_>>,
    ) -> Result<&mut GcStore> {
        log::trace!("allocating GC heap for store {:?}", self.id());

        assert!(self.gc_store.is_none());
        assert_eq!(
            self.vm_store_context.gc_heap.get_mut().base.as_non_null(),
            NonNull::dangling(),
        );
        assert_eq!(self.vm_store_context.gc_heap.get_mut().current_length(), 0);

        let engine = self.engine();
        let mem_ty = engine.tunables().gc_heap_memory_type();
        ensure!(
            engine.features().gc_types(),
            "cannot allocate a GC store when GC is disabled at configuration time"
        );
        let gc_runtime = engine
            .gc_runtime()
            .context("no GC runtime: GC disabled at compile time or configuration time")?;

        // First, allocate the memory that will be our GC heap's storage.
        let mut request = InstanceAllocationRequest {
            id: InstanceId::reserved_value(),
            runtime_info: engine.empty_module_runtime_info(),
            imports: vm::Imports::default(),
            store: self,
            limiter,
        };

        let (mem_alloc_index, mem) = engine
            .allocator()
            .allocate_memory(
                &mut request,
                &mem_ty,
                None,
                wasmtime_environ::MemoryKind::GcHeap,
            )
            .await?;

        // Then, allocate the actual GC heap, passing in that memory
        // storage.
        let (index, mut heap) =
            match engine
                .allocator()
                .allocate_gc_heap(engine, &**gc_runtime, mem_alloc_index)
            {
                Ok(pair) => pair,
                Err(e) => unsafe {
                    engine
                        .allocator()
                        .deallocate_memory(None, mem_alloc_index, mem);
                    return Err(e);
                },
            };
        heap.attach(mem);

        let gc_store = GcStore::new(index, heap, engine.tunables().gc_zeal_alloc_counter);
        *self.vm_store_context.gc_heap.get_mut() = gc_store.vmmemory_definition();
        Ok(self.gc_store.insert(gc_store))
    }

    /// Tests whether there is a pending exception.
    pub fn has_pending_exception(&self) -> bool {
        self.gc_data.pending_exception.is_some()
    }

    #[inline]
    pub(crate) fn gc_roots(&self) -> &RootSet {
        &self.gc_data.gc_roots
    }

    #[inline]
    pub(crate) fn gc_roots_mut(&mut self) -> &mut RootSet {
        &mut self.gc_data.gc_roots
    }

    #[inline]
    pub(crate) fn enter_gc_lifo_scope(&self) -> usize {
        self.gc_data.gc_roots.enter_lifo_scope()
    }

    #[inline]
    pub(crate) fn exit_gc_lifo_scope(&mut self, scope: usize) {
        self.gc_data
            .gc_roots
            .exit_lifo_scope(self.gc_store.as_mut(), scope);
    }

    /// Is type `sub` a subtype of `sup`?
    ///
    /// Equivalent to `self.engine().signatures().is_subtype(sub, sup)` but
    /// caches results store-locally to avoid contention on the engine's type
    /// registry lock. See the documentation of the `subtype_check_cache` field
    /// for details.
    pub(crate) fn is_subtype_cached(
        &mut self,
        sub: wasmtime_environ::VMSharedTypeIndex,
        sup: wasmtime_environ::VMSharedTypeIndex,
    ) -> bool {
        const MAX_SIZE: usize = 1 << 16; // 64k entries

        let key = (u64::from(sub.as_u32()) << 32) | u64::from(sup.as_u32());
        let engine_answer = || self.engine.signatures().is_subtype(sub, sup);
        if self.gc_data.subtype_check_cache.len() < MAX_SIZE {
            *self
                .gc_data
                .subtype_check_cache
                .entry(key)
                .or_insert_with(engine_answer)
        } else {
            self.gc_data
                .subtype_check_cache
                .get(&key)
                .copied()
                .unwrap_or_else(engine_answer)
        }
    }
}

/// RAII type to temporarily disable the GC zeal allocation counter.
struct WithoutGcZealAllocCounter<'a> {
    store: &'a mut StoreOpaque,
    counter: Option<NonZeroU32>,
}

impl Deref for WithoutGcZealAllocCounter<'_> {
    type Target = StoreOpaque;

    fn deref(&self) -> &Self::Target {
        &self.store
    }
}

impl DerefMut for WithoutGcZealAllocCounter<'_> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.store
    }
}

impl Drop for WithoutGcZealAllocCounter<'_> {
    fn drop(&mut self) {
        self.store.replace_gc_zeal_alloc_counter(self.counter);
    }
}

impl<'a> WithoutGcZealAllocCounter<'a> {
    pub fn new(store: &'a mut StoreOpaque) -> Self {
        let counter = store.replace_gc_zeal_alloc_counter(None);
        WithoutGcZealAllocCounter { store, counter }
    }
}

/// Given that we've hit a `GcHeapOutOfMemory` error, should we try freeing up
/// space by collecting first or by growing the GC heap first?
///
/// * `bytes_needed`: the number of bytes the mutator wants to allocate
///
/// * `gc_heap_capacity`: The current size of the GC heap.
///
/// * `last_gc_heap_usage`: The precise GC heap usage after the last collection.
#[track_caller]
fn should_collect_first(
    bytes_needed: u64,
    gc_heap_capacity: usize,
    last_gc_heap_usage: usize,
) -> bool {
    debug_assert!(last_gc_heap_usage <= gc_heap_capacity);

    // If we haven't allocated the GC heap yet, there's nothing to collect.
    //
    // Make sure to grow in this scenario even when the GC zeal infrastructure
    // passes `bytes_needed = 0`. This way our retry-after-gc logic doesn't
    // auto-fail on its second attempt, which would be bad because it doesn't
    // necessarily retry more than once.
    if gc_heap_capacity == 0 {
        return false;
    }

    // The GC zeal infrastructure will use `bytes_needed = 0` to trigger extra
    // collections.
    if bytes_needed == 0 {
        return true;
    }

    let Ok(bytes_needed) = usize::try_from(bytes_needed) else {
        // No point wasting time on collection if we will never be able to
        // satisfy the allocation.
        return false;
    };

    if bytes_needed > isize::MAX.cast_unsigned() {
        // Similarly, no allocation can be larger than `isize::MAX` in Rust (or
        // LLVM), so don't bother wasting time on collection if we will never be
        // able to satisfy the allocation.
        return false;
    }

    let Some(predicted_usage) = last_gc_heap_usage.checked_add(bytes_needed) else {
        // If we can't represent our predicted usage as a `usize`, we won't be
        // able to grow the GC heap to that size, so try collecting first to
        // free up space.
        return true;
    };

    // Common case: to balance collection frequency (and its time overhead) with
    // GC heap growth (and its space overhead), only prefer growing first if the
    // predicted GC heap utilization is greater than half the GC heap's
    // capacity.
    predicted_usage < gc_heap_capacity / 2
}

#[cfg(test)]
mod tests {
    use super::should_collect_first;
    use crate::{AsContextMut, Config, Engine, ExternRef, Result, Store};

    #[test]
    fn test_should_collect_first() {
        // No GC heap yet special case.
        for bytes_needed in 0..256 {
            assert_eq!(should_collect_first(bytes_needed, 0, 0), false);
        }

        // GC zeal special case.
        for cap in 1..256 {
            for usage in 0..=cap {
                assert_eq!(should_collect_first(0, cap, usage), true);
            }
        }

        let max_alloc_usize = isize::MAX.cast_unsigned();
        let max_alloc_u64 = u64::try_from(max_alloc_usize).unwrap();

        // Allocation size larger than `isize::MAX` --> will never succeed, do
        // not bother collecting.
        assert_eq!(
            should_collect_first(max_alloc_u64 + 1, max_alloc_usize, 0),
            false,
        );

        // Predicted usage overflow --> growth will likely fail, collect first.
        assert_eq!(should_collect_first(1, usize::MAX, usize::MAX), true);

        // Common case: predicted usage is low --> we likely have more than
        // enough space already, so collect first.
        assert_eq!(should_collect_first(16, 1024, 64), true);

        // Common case: predicted usage is high --> plausible we may not have
        // enough space, and we want to amortize the cost of collections, so
        // grow first.
        assert_eq!(should_collect_first(16, 1024, 512), false);
    }

    #[test]
    fn gc_heap_initial_size() -> Result<()> {
        let mut config = Config::new();
        config.gc_heap_initial_size(1 << 20);
        let engine = Engine::new(&config)?;
        let mut store = Store::new(&engine, ());
        ExternRef::new(&mut store, 1)?;

        let gc_store = store.as_context_mut().0.unwrap_gc_store();
        assert_eq!(gc_store.gc_heap_capacity(), 1 << 20);
        Ok(())
    }
}