wasmtime_bench_api/

lib.rs

//! A C API for benchmarking Wasmtime's WebAssembly compilation, instantiation,
//! and execution.
//!
//! The API expects calls that match the following state machine:
//!
//! ```text
//!               |
//!               |
//!               V
//! .---> wasm_bench_create
//! |        |        |
//! |        |        |
//! |        |        V
//! |        |   wasm_bench_compile
//! |        |     |            |
//! |        |     |            |     .----.
//! |        |     |            |     |    |
//! |        |     |            V     V    |
//! |        |     |     wasm_bench_instantiate <------.
//! |        |     |            |        |             |
//! |        |     |            |        |             |
//! |        |     |            |        |             |
//! |        |     |     .------'        '-----> wasm_bench_execute
//! |        |     |     |                             |
//! |        |     |     |                             |
//! |        V     V     V                             |
//! '------ wasm_bench_free <--------------------------'
//!               |
//!               |
//!               V
//! ```
//!
//! All API calls must happen on the same thread.
//!
//! Functions which return pointers use null as an error value. Function which
//! return `int` use `0` as OK and non-zero as an error value.
//!
//! # Example
//!
//! ```
//! use std::ptr;
//! use wasmtime_bench_api::*;
//!
//! let working_dir = std::env::current_dir().unwrap().display().to_string();
//! let stdout_path = "./stdout.log";
//! let stderr_path = "./stderr.log";
//!
//! // Functions to start/end timers for compilation.
//! //
//! // The `compilation_timer` pointer configured in the `WasmBenchConfig` is
//! // passed through.
//! extern "C" fn compilation_start(timer: *mut u8) {
//!     // Start your compilation timer here.
//! }
//! extern "C" fn compilation_end(timer: *mut u8) {
//!     // End your compilation timer here.
//! }
//!
//! // Similar for instantiation.
//! extern "C" fn instantiation_start(timer: *mut u8) {
//!     // Start your instantiation timer here.
//! }
//! extern "C" fn instantiation_end(timer: *mut u8) {
//!     // End your instantiation timer here.
//! }
//!
//! // Similar for execution.
//! extern "C" fn execution_start(timer: *mut u8) {
//!     // Start your execution timer here.
//! }
//! extern "C" fn execution_end(timer: *mut u8) {
//!     // End your execution timer here.
//! }
//!
//! let config = WasmBenchConfig {
//!     version: WASM_BENCH_CONFIG_VERSION,
//!     size: std::mem::size_of::<WasmBenchConfig>(),
//!     working_dir_ptr: working_dir.as_ptr(),
//!     working_dir_len: working_dir.len(),
//!     stdout_path_ptr: stdout_path.as_ptr(),
//!     stdout_path_len: stdout_path.len(),
//!     stderr_path_ptr: stderr_path.as_ptr(),
//!     stderr_path_len: stderr_path.len(),
//!     stdin_path_ptr: ptr::null(),
//!     stdin_path_len: 0,
//!     compilation_timer: ptr::null_mut(),
//!     compilation_start,
//!     compilation_end,
//!     instantiation_timer: ptr::null_mut(),
//!     instantiation_start,
//!     instantiation_end,
//!     execution_timer: ptr::null_mut(),
//!     execution_start,
//!     execution_end,
//!     execution_flags_ptr: ptr::null(),
//!     execution_flags_len: 0,
//!     cache: false,
//! };
//!
//! let mut bench_api = ptr::null_mut();
//! unsafe {
//!     let code = wasm_bench_create(&config, &mut bench_api);
//!     assert_eq!(code, OK);
//!     assert!(!bench_api.is_null());
//! };
//!
//! let wasm = wat::parse_bytes(br#"
//!     (module
//!         (func $bench_start (import "bench" "start"))
//!         (func $bench_end (import "bench" "end"))
//!         (func $start (export "_start")
//!             call $bench_start
//!             i32.const 1
//!             i32.const 2
//!             i32.add
//!             drop
//!             call $bench_end
//!         )
//!     )
//! "#).unwrap();
//!
//! // This will call the `compilation_{start,end}` timing functions on success.
//! let code = unsafe { wasm_bench_compile(bench_api, wasm.as_ptr(), wasm.len()) };
//! assert_eq!(code, OK);
//!
//! // This will call the `instantiation_{start,end}` timing functions on success.
//! let code = unsafe { wasm_bench_instantiate(bench_api) };
//! assert_eq!(code, OK);
//!
//! // This will call the `execution_{start,end}` timing functions on success.
//! let code = unsafe { wasm_bench_execute(bench_api) };
//! assert_eq!(code, OK);
//!
//! unsafe {
//!     wasm_bench_free(bench_api);
//! }
//! ```

mod unsafe_send_sync;

use crate::unsafe_send_sync::UnsafeSendSync;
use clap::Parser;
use std::os::raw::{c_int, c_void};
use std::slice;
use std::{env, path::PathBuf};
use wasmtime::{
    CodeBuilder, CodeHint, Engine, Linker, Module, Result, Store, error::Context as _, format_err,
};
use wasmtime_cli_flags::CommonOptions;
use wasmtime_wasi::cli::{InputFile, OutputFile};
use wasmtime_wasi::{FsPerms, I32Exit, WasiCtx, WasiCtxView, WasiView, p1::WasiP1Ctx};

pub type ExitCode = c_int;
pub const OK: ExitCode = 0;
pub const ERR: ExitCode = -1;

/// The current version of `WasmBenchConfig`.
///
/// Bumped once per release that appends a field, and never otherwise, so a
/// field added in version `N` is present exactly when the caller's `version` is
/// at least `N`.
pub const WASM_BENCH_CONFIG_VERSION: u32 = 1;

// Randomize the location of heap objects to avoid accidental locality being an
// uncontrolled variable that obscures performance evaluation in our
// experiments.
#[cfg(feature = "shuffling-allocator")]
#[global_allocator]
static ALLOC: shuffling_allocator::ShufflingAllocator<std::alloc::System> =
    shuffling_allocator::wrap!(&std::alloc::System);

/// Generate `WasmBenchConfig` and a `Config` accessor for each of its fields.
///
/// Tag each field with the `#[since_version = N]` that introduced it. Version 0
/// fields are mandatory, so their accessors return the field directly; later
/// fields may be missing from an older caller, so theirs return `Option`.
macro_rules! define_config {
    (
        $(#[doc = $struct_doc:expr])*
        pub struct $name:ident {
            $(
                $(#[doc = $field_doc:expr])*
                #[since_version = $version:tt]
                pub $field:ident : $ty:ty,
            )*
        }
    ) => {
        $(#[doc = $struct_doc])*
        #[repr(C)]
        pub struct $name {
            $(
                $(#[doc = $field_doc])*
                pub $field : $ty,
            )*
        }

        impl Config<'_> {
            $(
                define_config!(@accessor $name, $field, $ty, $version);
            )*
        }
    };

    // Version 0 fields are mandatory. `Config::new` rejects any caller too
    // small to hold them, so they are always in bounds.
    (@accessor $name:ident, $field:ident, $ty:ty, 0) => {
        fn $field(&self) -> $ty {
            unsafe { self.read(std::ptr::addr_of!((*self.ptr).$field)) }
        }
    };

    // Fields added later may lie outside an older caller's allocation, so they
    // are only read once `has_field` confirms this caller provided them.
    (@accessor $name:ident, $field:ident, $ty:ty, $version:tt) => {
        fn $field(&self) -> Option<$ty> {
            let offset = std::mem::offset_of!($name, $field);
            if self.has_field($version, offset, size_of::<$ty>()) {
                Some(unsafe { self.read(std::ptr::addr_of!((*self.ptr).$field)) })
            } else {
                None
            }
        }
    };
}

define_config! {
    /// Configuration options for the benchmark.
    ///
    /// This structure is always passed across the bench API by pointer, and new
    /// fields are only ever appended to the end, so that a caller built against
    /// an older version of this structure can still drive a newer bench API
    /// library.
    ///
    /// We only ever read one through [`Config`], never as a reference.
    pub struct WasmBenchConfig {
        /// The version of this structure that the caller was built against.
        ///
        /// Each field's documentation names the version that added it; a field
        /// is only present when `version` is at least that.
        #[since_version = 0]
        pub version: u32,

        /// The size of this structure, in bytes, as the caller knows it.
        ///
        /// Fields beyond these first `size` bytes were not allocated by the
        /// caller and must not be read.
        ///
        /// This acts as a secondary backstop to `version`.
        #[since_version = 0]
        pub size: usize,

        /// The working directory where benchmarks should be executed.
        #[since_version = 0]
        pub working_dir_ptr: *const u8,
        #[since_version = 0]
        pub working_dir_len: usize,

        /// The file path that should be created and used as `stdout`.
        #[since_version = 0]
        pub stdout_path_ptr: *const u8,
        #[since_version = 0]
        pub stdout_path_len: usize,

        /// The file path that should be created and used as `stderr`.
        #[since_version = 0]
        pub stderr_path_ptr: *const u8,
        #[since_version = 0]
        pub stderr_path_len: usize,

        /// The (optional) file path that should be opened and used as `stdin`.
        /// If not provided, then the WASI context will not have a `stdin`
        /// initialized.
        #[since_version = 0]
        pub stdin_path_ptr: *const u8,
        #[since_version = 0]
        pub stdin_path_len: usize,

        /// The functions to start and stop performance timers/counters during
        /// Wasm compilation.
        #[since_version = 0]
        pub compilation_timer: *mut u8,
        #[since_version = 0]
        pub compilation_start: extern "C" fn(*mut u8),
        #[since_version = 0]
        pub compilation_end: extern "C" fn(*mut u8),

        /// The functions to start and stop performance timers/counters during
        /// Wasm instantiation.
        #[since_version = 0]
        pub instantiation_timer: *mut u8,
        #[since_version = 0]
        pub instantiation_start: extern "C" fn(*mut u8),
        #[since_version = 0]
        pub instantiation_end: extern "C" fn(*mut u8),

        /// The functions to start and stop performance timers/counters during
        /// Wasm execution.
        #[since_version = 0]
        pub execution_timer: *mut u8,
        #[since_version = 0]
        pub execution_start: extern "C" fn(*mut u8),
        #[since_version = 0]
        pub execution_end: extern "C" fn(*mut u8),

        /// The (optional) flags to use when running Wasmtime. These correspond
        /// to the flags used when running Wasmtime from the command line.
        #[since_version = 0]
        pub execution_flags_ptr: *const u8,
        #[since_version = 0]
        pub execution_flags_len: usize,

        /// Whether to enable Wasmtime's compilation cache.
        ///
        /// Benchmark runs that measure compilation want this off, so that they
        /// measure compilation rather than a cache lookup. Benchmark runs that
        /// measure only instantiation or execution want this on, to avoid
        /// spending time recompiling the Wasm.
        #[since_version = 1]
        pub cache: bool,
    }
}

/// A caller-provided [`WasmBenchConfig`], read through its pointer.
///
/// Holding the pointer rather than a reference is what keeps this sound. A
/// `&WasmBenchConfig` would assert that the whole struct is initialized, which
/// is not true of its padding bytes, nor of the fields belonging to versions
/// newer than the caller's.
pub struct Config<'a> {
    ptr: *const WasmBenchConfig,
    _lifetime: std::marker::PhantomData<&'a WasmBenchConfig>,
}

impl<'a> Config<'a> {
    /// The fields every caller must provide, regardless of how old it is.
    const MIN_SIZE: usize =
        std::mem::offset_of!(WasmBenchConfig, execution_flags_len) + size_of::<usize>();

    /// Wrap a caller-provided config pointer, after checking it is large enough
    /// to hold every mandatory field.
    ///
    /// # Safety
    ///
    /// `ptr` must be non-null, aligned, and point to a `version` and `size`
    /// that accurately describe the fields that follow, all of which must
    /// remain valid for `'a`.
    unsafe fn new(ptr: *const WasmBenchConfig) -> Result<Config<'a>> {
        assert!(!ptr.is_null());
        let config = Config {
            ptr,
            _lifetime: std::marker::PhantomData,
        };

        // Reads `size`, which every caller provides, so this is in bounds even
        // for the oldest one.
        let size = config.size();
        if size < Self::MIN_SIZE {
            return Err(format_err!(
                "given `WasmBenchConfig` is {} bytes, but at least {} are required",
                size,
                Self::MIN_SIZE,
            ));
        }

        Ok(config)
    }

    /// Read one field out of the caller's config.
    ///
    /// # Safety
    ///
    /// `field` must be a pointer to a field this caller actually provided,
    /// derived from `self.ptr`.
    unsafe fn read<T>(&self, field: *const T) -> T {
        assert!(!field.is_null());
        assert!(field.is_aligned());
        unsafe { field.read() }
    }

    /// Did the caller provide the field that version `version` added, which
    /// occupies `field_size` bytes at `offset`?
    ///
    /// Both checks must pass. `version` alone is the field's real name, but
    /// trusting it would let a caller that bumped its version without growing
    /// the struct send us past its allocation; `size` alone cannot tell a field
    /// apart from the padding that may follow the one before it.
    fn has_field(&self, version: u32, offset: usize, field_size: usize) -> bool {
        self.version() >= version && self.size() >= offset + field_size
    }

    /// Whether to enable the compilation cache.
    ///
    /// Conservatively defaults to disabled for callers that predate the `cache`
    /// field, so we don't break compilation measurements.
    fn cache_enabled(&self) -> bool {
        self.cache().unwrap_or(false)
    }

    fn working_dir(&self) -> Result<PathBuf> {
        let working_dir =
            unsafe { std::slice::from_raw_parts(self.working_dir_ptr(), self.working_dir_len()) };
        let working_dir = std::str::from_utf8(working_dir)
            .context("given working directory is not valid UTF-8")?;
        Ok(working_dir.into())
    }

    fn stdout_path(&self) -> Result<PathBuf> {
        let stdout_path =
            unsafe { std::slice::from_raw_parts(self.stdout_path_ptr(), self.stdout_path_len()) };
        let stdout_path =
            std::str::from_utf8(stdout_path).context("given stdout path is not valid UTF-8")?;
        Ok(stdout_path.into())
    }

    fn stderr_path(&self) -> Result<PathBuf> {
        let stderr_path =
            unsafe { std::slice::from_raw_parts(self.stderr_path_ptr(), self.stderr_path_len()) };
        let stderr_path =
            std::str::from_utf8(stderr_path).context("given stderr path is not valid UTF-8")?;
        Ok(stderr_path.into())
    }

    fn stdin_path(&self) -> Result<Option<PathBuf>> {
        let stdin_path_ptr = self.stdin_path_ptr();
        if stdin_path_ptr.is_null() {
            return Ok(None);
        }

        let stdin_path =
            unsafe { std::slice::from_raw_parts(stdin_path_ptr, self.stdin_path_len()) };
        let stdin_path =
            std::str::from_utf8(stdin_path).context("given stdin path is not valid UTF-8")?;
        Ok(Some(stdin_path.into()))
    }

    fn execution_flags(&self) -> Result<CommonOptions> {
        let execution_flags_ptr = self.execution_flags_ptr();
        let flags = if execution_flags_ptr.is_null() {
            ""
        } else {
            let execution_flags = unsafe {
                std::slice::from_raw_parts(execution_flags_ptr, self.execution_flags_len())
            };
            std::str::from_utf8(execution_flags)
                .context("given execution flags string is not valid UTF-8")?
        };
        let options = CommonOptions::try_parse_from(
            ["wasmtime"]
                .into_iter()
                .chain(flags.split(' ').filter(|s| !s.is_empty())),
        )
        .context("failed to parse options")?;
        Ok(options)
    }
}

/// Exposes a C-compatible way of creating the engine from the bytes of a single
/// Wasm module.
///
/// On success, the `out_bench_ptr` is initialized to a pointer to a structure
/// that contains the engine's initialized state, and `0` is returned. On
/// failure, a non-zero status code is returned and `out_bench_ptr` is left
/// untouched.
///
/// The `config` pointer only needs to remain valid for the duration of this
/// call, as does everything it points to.
#[unsafe(no_mangle)]
pub extern "C" fn wasm_bench_create(
    config: *const WasmBenchConfig,
    out_bench_ptr: *mut *mut c_void,
) -> ExitCode {
    let result = (|| -> Result<_> {
        let config = unsafe { Config::new(config)? };

        // Clone paths for the core Wasm WASI closure.
        let working_dir = config.working_dir()?;
        let stdout_path = config.stdout_path()?;
        let stderr_path = config.stderr_path()?;
        let stdin_path = config.stdin_path()?;
        let options = config.execution_flags()?;

        let state = Box::new(BenchState::new(
            options,
            config.cache_enabled(),
            config.compilation_timer(),
            config.compilation_start(),
            config.compilation_end(),
            config.instantiation_timer(),
            config.instantiation_start(),
            config.instantiation_end(),
            config.execution_timer(),
            config.execution_start(),
            config.execution_end(),
            move || {
                let mut cx = WasiCtx::builder();

                let stdout = std::fs::File::create(&stdout_path)
                    .with_context(|| format!("failed to create {}", stdout_path.display()))?;
                cx.stdout(OutputFile::new(stdout));

                let stderr = std::fs::File::create(&stderr_path)
                    .with_context(|| format!("failed to create {}", stderr_path.display()))?;
                cx.stderr(OutputFile::new(stderr));

                if let Some(stdin_path) = &stdin_path {
                    let stdin = std::fs::File::open(stdin_path)
                        .with_context(|| format!("failed to open {}", stdin_path.display()))?;
                    cx.stdin(InputFile::new(stdin));
                }

                // Allow access to the working directory so that the benchmark can read
                // its input workload(s).
                cx.preopened_dir(working_dir.clone(), ".", FsPerms::ReadOnly)?;

                // Pass this env var along so that the benchmark program can use smaller
                // input workload(s) if it has them and that has been requested.
                if let Ok(val) = env::var("WASM_BENCH_USE_SMALL_WORKLOAD") {
                    cx.env("WASM_BENCH_USE_SMALL_WORKLOAD", &val);
                }

                Ok(cx.build_p1())
            },
        )?);
        Ok(Box::into_raw(state) as _)
    })();

    if let Ok(bench_ptr) = result {
        unsafe {
            assert!(!out_bench_ptr.is_null());
            *out_bench_ptr = bench_ptr;
        }
    }

    to_exit_code(result.map(|_| ()))
}

/// Free the engine state allocated by this library.
#[unsafe(no_mangle)]
pub extern "C" fn wasm_bench_free(state: *mut c_void) {
    assert!(!state.is_null());
    unsafe {
        drop(Box::from_raw(state as *mut BenchState));
    }
}

/// Compile the Wasm benchmark module.
#[unsafe(no_mangle)]
pub extern "C" fn wasm_bench_compile(
    state: *mut c_void,
    wasm_bytes: *const u8,
    wasm_bytes_length: usize,
) -> ExitCode {
    let state = unsafe { (state as *mut BenchState).as_mut().unwrap() };
    let wasm_bytes = unsafe { slice::from_raw_parts(wasm_bytes, wasm_bytes_length) };
    let result = state.compile(wasm_bytes).context("failed to compile");
    to_exit_code(result)
}

/// Instantiate the Wasm benchmark module.
#[unsafe(no_mangle)]
pub extern "C" fn wasm_bench_instantiate(state: *mut c_void) -> ExitCode {
    let state = unsafe { (state as *mut BenchState).as_mut().unwrap() };
    let result = state.instantiate().context("failed to instantiate");
    to_exit_code(result)
}

/// Execute the Wasm benchmark module.
#[unsafe(no_mangle)]
pub extern "C" fn wasm_bench_execute(state: *mut c_void) -> ExitCode {
    let state = unsafe { (state as *mut BenchState).as_mut().unwrap() };
    let result = state.execute().context("failed to execute");
    to_exit_code(result)
}

/// Helper function for converting a Rust result to a C error code.
///
/// This will print an error indicating some information regarding the failure.
fn to_exit_code<T>(result: impl Into<Result<T>>) -> ExitCode {
    match result.into() {
        Ok(_) => OK,
        Err(error) => {
            eprintln!("{error:?}");
            ERR
        }
    }
}

/// This structure contains the actual Rust implementation of the state required
/// to manage the Wasmtime engine between calls.
struct BenchState {
    linker: Linker<HostState>,
    component_linker: wasmtime::component::Linker<HostState>,
    compilation_timer: *mut u8,
    compilation_start: extern "C" fn(*mut u8),
    compilation_end: extern "C" fn(*mut u8),
    instantiation_timer: *mut u8,
    instantiation_start: extern "C" fn(*mut u8),
    instantiation_end: extern "C" fn(*mut u8),
    make_wasi_cx: Box<dyn FnMut() -> Result<WasiP1Ctx>>,
    module: Option<Module>,
    component: Option<wasmtime::component::Component>,
    store_and_instance: Option<(Store<HostState>, Instance)>,
    epoch_interruption: bool,
    fuel: Option<u64>,
}

enum Instance {
    Core(wasmtime::Instance),
    Component(wasmtime::component::Instance),
}

struct HostState {
    wasi: WasiP1Ctx,
    #[cfg(feature = "wasi-nn")]
    wasi_nn: wasmtime_wasi_nn::witx::WasiNnCtx,
}

impl WasiView for HostState {
    fn ctx(&mut self) -> WasiCtxView<'_> {
        self.wasi.ctx()
    }
}

impl BenchState {
    fn new(
        mut options: CommonOptions,
        cache: bool,
        compilation_timer: *mut u8,
        compilation_start: extern "C" fn(*mut u8),
        compilation_end: extern "C" fn(*mut u8),
        instantiation_timer: *mut u8,
        instantiation_start: extern "C" fn(*mut u8),
        instantiation_end: extern "C" fn(*mut u8),
        execution_timer: *mut u8,
        execution_start: extern "C" fn(*mut u8),
        execution_end: extern "C" fn(*mut u8),
        make_wasi_cx: impl FnMut() -> Result<WasiP1Ctx> + 'static,
    ) -> Result<Self> {
        let mut config = options.config(None)?;

        // NB: set this either way, since `CommonOptions` enables the cache by
        // default and we must not break compilation benchmarks.
        config.cache(if cache {
            Some(wasmtime::Cache::from_file(None)?)
        } else {
            None
        });

        let engine = Engine::new(&config)?;
        let mut linker = Linker::<HostState>::new(&engine);
        let mut component_linker = wasmtime::component::Linker::<HostState>::new(&engine);

        // Define the benchmarking start/end functions.
        let execution_timer = unsafe {
            // Safe because this bench API's contract requires that its methods
            // are only ever called from a single thread.
            UnsafeSendSync::new(execution_timer)
        };
        linker.func_wrap("bench", "start", move || {
            execution_start(*execution_timer.get());
            Ok(())
        })?;
        linker.func_wrap("bench", "end", move || {
            execution_end(*execution_timer.get());
            Ok(())
        })?;

        // Define the same benchmarking functions on the component linker.
        let mut bench_instance = component_linker.instance("bench")?;
        bench_instance.func_wrap(
            "start",
            move |_: wasmtime::StoreContextMut<'_, _>, (): ()| {
                execution_start(*execution_timer.get());
                Ok(())
            },
        )?;
        bench_instance.func_wrap("end", move |_: wasmtime::StoreContextMut<'_, _>, (): ()| {
            execution_end(*execution_timer.get());
            Ok(())
        })?;

        let epoch_interruption = options.wasm.epoch_interruption.unwrap_or(false);
        let fuel = options.wasm.fuel;

        if options.wasi.common != Some(false) {
            wasmtime_wasi::p1::add_to_linker_sync(&mut linker, |cx| &mut cx.wasi)?;
            wasmtime_wasi::p2::add_to_linker_sync(&mut component_linker)?;
        }

        #[cfg(feature = "wasi-nn")]
        if options.wasi.nn == Some(true) {
            wasmtime_wasi_nn::witx::add_to_linker(&mut linker, |cx| &mut cx.wasi_nn)?;
        }

        Ok(Self {
            linker,
            component_linker,
            compilation_timer,
            compilation_start,
            compilation_end,
            instantiation_timer,
            instantiation_start,
            instantiation_end,
            make_wasi_cx: Box::new(make_wasi_cx) as _,
            module: None,
            component: None,
            store_and_instance: None,
            epoch_interruption,
            fuel,
        })
    }

    fn compile(&mut self, bytes: &[u8]) -> Result<()> {
        self.module = None;
        self.component = None;

        let mut builder = CodeBuilder::new(self.linker.engine());
        builder.wasm_binary(bytes, None)?;

        match builder.hint() {
            Some(CodeHint::Component) => {
                (self.compilation_start)(self.compilation_timer);
                let component = builder.compile_component()?;
                (self.compilation_end)(self.compilation_timer);
                self.component = Some(component);
            }
            Some(CodeHint::Module) | None => {
                (self.compilation_start)(self.compilation_timer);
                let module = builder.compile_module()?;
                (self.compilation_end)(self.compilation_timer);
                self.module = Some(module);
            }
        }

        Ok(())
    }

    fn instantiate(&mut self) -> Result<()> {
        self.store_and_instance = None;

        let host = HostState {
            wasi: (self.make_wasi_cx)().context("failed to create a WASI context")?,
            #[cfg(feature = "wasi-nn")]
            wasi_nn: {
                let (backends, registry) = wasmtime_wasi_nn::preload(&[])?;
                wasmtime_wasi_nn::witx::WasiNnCtx::new(backends, registry)
            },
        };

        // NB: Start measuring instantiation time *after* we've created the
        // WASI context, since that needs to do file I/O to setup
        // stdin/stdout/stderr.
        (self.instantiation_start)(self.instantiation_timer);

        let mut store = Store::new(self.linker.engine(), host);
        if self.epoch_interruption {
            store.set_epoch_deadline(1);
        }
        if let Some(fuel) = self.fuel {
            store.set_fuel(fuel).unwrap();
        }

        let instance = match &self.component {
            Some(component) => {
                Instance::Component(self.component_linker.instantiate(&mut store, component)?)
            }
            None => {
                let module = self
                    .module
                    .as_ref()
                    .expect("compile the module before instantiating it");
                Instance::Core(self.linker.instantiate(&mut store, &module)?)
            }
        };

        (self.instantiation_end)(self.instantiation_timer);

        self.store_and_instance = Some((store, instance));
        Ok(())
    }

    fn execute(&mut self) -> Result<()> {
        match self.store_and_instance.take() {
            Some((mut store, Instance::Component(instance))) => {
                let command =
                    wasmtime_wasi::p2::bindings::sync::Command::new(&mut store, &instance)?;
                match command.wasi_cli_run().call_run(&mut store)? {
                    Ok(()) => Ok(()),
                    Err(()) => Err(format_err!("calling `run` failed")),
                }
            }
            Some((mut store, Instance::Core(instance))) => {
                let start_func = instance.get_typed_func::<(), ()>(&mut store, "_start")?;
                match start_func.call(&mut store, ()) {
                    Ok(_) => Ok(()),
                    Err(trap) => {
                        // Since _start will likely return by using the system `exit` call, we must
                        // check the trap code to see if it actually represents a successful exit.
                        if let Some(exit) = trap.downcast_ref::<I32Exit>() {
                            if exit.0 == 0 {
                                return Ok(());
                            }
                        }

                        Err(trap)
                    }
                }
            }
            None => {
                panic!("instantiate before execution")
            }
        }
    }
}