Allow host functions to call Wasm functions (#590)

* add StashArena data structure to wasmi_arena

* rename wasmi_arena::Index to ArenaIndex

* import core::ops::Index[Mut] since wasmi_arena::Index got renamed

* properly forward to Stash::is_empty

* put shared engine resources behind RwLock

This will later allow multiple EngineExecutors to share a single read-only EngineResources.

* remove unnecessary clone call

* remove redundant clone in stash tests

* add proper Default impls to Stash[Arena]

* add EngineStacks struct

* integrate new EngineStacks type

* add cached_stacks option to Config

* add Config::cached_stacks option

* fix various compile warnings and errors

* add test for host calling wasm back

* remove [Arena]Stash data structure again

* apply clippy suggestions

* apply rustfmt suggestions

* apply rustfmt suggestions v2

* apply rustfmt suggestions v3

* apply rustfmt suggestions v4

crates/wasmi/Cargo.toml

@@ -19,6 +19,7 @@ wasmi_arena = { version = "0.1", path = "../arena", default-features = false }
 spin = { version = "0.9", default-features = false, features = [
     "mutex",
     "spin_mutex",
+    "rwlock",
 ] }
 
 [dev-dependencies]

crates/wasmi/src/engine/config.rs

@@ -1,13 +1,18 @@
 use super::stack::StackLimits;
 use wasmparser::WasmFeatures;
 
+/// The default amount of stacks kept in the cache at most.
+const DEFAULT_CACHED_STACKS: usize = 2;
+
 /// Configuration for an [`Engine`].
 ///
 /// [`Engine`]: [`crate::Engine`]
 #[derive(Debug, Copy, Clone)]
 pub struct Config {
     /// The limits set on the value stack and call stack.
     stack_limits: StackLimits,
+    /// The amount of Wasm stacks to keep in cache at most.
+    cached_stacks: usize,
     /// Is `true` if the `mutable-global` Wasm proposal is enabled.
     mutable_global: bool,
     /// Is `true` if the `sign-extension` Wasm proposal is enabled.
@@ -22,6 +27,7 @@ impl Default for Config {
     fn default() -> Self {
         Self {
             stack_limits: StackLimits::default(),
+            cached_stacks: DEFAULT_CACHED_STACKS,
             mutable_global: true,
             sign_extension: true,
             saturating_float_to_int: true,
@@ -42,6 +48,21 @@ impl Config {
         self.stack_limits
     }
 
+    /// Sets the maximum amount of cached stacks for reuse for the [`Config`].
+    ///
+    /// # Note
+    ///
+    /// Defaults to 2.
+    pub fn set_cached_stacks(&mut self, amount: usize) -> &mut Self {
+        self.cached_stacks = amount;
+        self
+    }
+
+    /// Returns the maximum amount of cached stacks for reuse of the [`Config`].
+    pub(super) fn cached_stacks(&self) -> usize {
+        self.cached_stacks
+    }
+
     /// Enable or disable the [`mutable-global`] Wasm proposal for the [`Config`].
     ///
     /// # Note

crates/wasmi/src/engine/mod.rs

@@ -14,7 +14,6 @@ mod traits;
 #[cfg(test)]
 mod tests;
 
-pub(crate) use self::func_args::{FuncParams, FuncResults};
 pub use self::{
     bytecode::DropKeep,
     code_map::FuncBody,
@@ -37,15 +36,18 @@ use self::{
     func_types::FuncTypeRegistry,
     stack::{FuncFrame, Stack, ValueStack},
 };
+pub(crate) use self::{
+    func_args::{FuncParams, FuncResults},
+    func_types::DedupFuncType,
+};
 use super::{func::FuncEntityInternal, AsContextMut, Func};
 use crate::{
     core::{Trap, TrapCode},
     FuncType,
 };
-use alloc::sync::Arc;
+use alloc::{sync::Arc, vec::Vec};
 use core::sync::atomic::{AtomicU32, Ordering};
-pub use func_types::DedupFuncType;
-use spin::mutex::Mutex;
+use spin::{Mutex, RwLock};
 use wasmi_arena::{ArenaIndex, GuardedEntity};
 
 /// The outcome of a `wasmi` function execution.
@@ -100,7 +102,7 @@ type Guarded<Idx> = GuardedEntity<EngineIdx, Idx>;
 ///   Most of its API has a `&self` receiver, so can be shared easily.
 #[derive(Debug, Clone)]
 pub struct Engine {
-    inner: Arc<Mutex<EngineInner>>,
+    inner: Arc<EngineInner>,
 }
 
 impl Default for Engine {
@@ -117,18 +119,18 @@ impl Engine {
     /// Users should ues [`Engine::default`] to construct a default [`Engine`].
     pub fn new(config: &Config) -> Self {
         Self {
-            inner: Arc::new(Mutex::new(EngineInner::new(config))),
+            inner: Arc::new(EngineInner::new(config)),
         }
     }
 
     /// Returns a shared reference to the [`Config`] of the [`Engine`].
     pub fn config(&self) -> Config {
-        *self.inner.lock().config()
+        self.inner.config()
     }
 
     /// Allocates a new function type to the engine.
     pub(super) fn alloc_func_type(&self, func_type: FuncType) -> DedupFuncType {
-        self.inner.lock().func_types.alloc_func_type(func_type)
+        self.inner.alloc_func_type(func_type)
     }
 
     /// Resolves a deduplicated function type into a [`FuncType`] entity.
@@ -141,8 +143,7 @@ impl Engine {
     where
         F: FnOnce(&FuncType) -> R,
     {
-        // Note: The clone operation on FuncType is intentionally cheap.
-        f(self.inner.lock().func_types.resolve_func_type(func_type))
+        self.inner.resolve_func_type(func_type, f)
     }
 
     /// Allocates the instructions of a Wasm function body to the [`Engine`].
@@ -159,7 +160,6 @@ impl Engine {
         I::IntoIter: ExactSizeIterator,
     {
         self.inner
-            .lock()
             .alloc_func_body(len_locals, max_stack_height, insts)
     }
 
@@ -176,11 +176,7 @@ impl Engine {
     /// If the [`FuncBody`] is invalid for the [`Engine`].
     #[cfg(test)]
     pub(crate) fn resolve_inst(&self, func_body: FuncBody, index: usize) -> Option<Instruction> {
-        self.inner
-            .lock()
-            .code_map
-            .get_instr(func_body, index)
-            .copied()
+        self.inner.resolve_inst(func_body, index)
     }
 
     /// Executes the given [`Func`] using the given arguments `params` and stores the result into `results`.
@@ -211,17 +207,133 @@ impl Engine {
         Params: CallParams,
         Results: CallResults,
     {
-        self.inner.lock().execute_func(ctx, func, params, results)
+        self.inner.execute_func(ctx, func, params, results)
     }
 }
 
 /// The internal state of the `wasmi` engine.
 #[derive(Debug)]
 pub struct EngineInner {
+    /// Engine resources shared across multiple engine executors.
+    res: RwLock<EngineResources>,
+    /// Reusable engine stacks for Wasm execution.
+    ///
+    /// Concurrently executing Wasm executions each require their own stack to
+    /// operate on. Therefore a Wasm engine is required to provide stacks and
+    /// ideally recycles old ones since creation of a new stack is rather expensive.
+    stacks: Mutex<EngineStacks>,
+}
+
+/// The engine's stacks for reuse.
+///
+/// Rquired for efficient concurrent Wasm executions.
+#[derive(Debug)]
+pub struct EngineStacks {
+    /// Stacks to be (re)used.
+    stacks: Vec<Stack>,
+    /// Stack limits for newly constructed engine stacks.
+    limits: StackLimits,
+    /// How many stacks should be kept for reuse at most.
+    keep: usize,
+}
+
+impl EngineStacks {
+    /// Creates new [`EngineStacks`] with the given [`StackLimits`].
+    pub fn new(config: &Config) -> Self {
+        Self {
+            stacks: Vec::new(),
+            limits: config.stack_limits(),
+            keep: config.cached_stacks(),
+        }
+    }
+
+    /// Reuse or create a new [`Stack`] if none was available.
+    pub fn reuse_or_new(&mut self) -> Stack {
+        match self.stacks.pop() {
+            Some(stack) => stack,
+            None => Stack::new(self.limits),
+        }
+    }
+
+    /// Disose and recycle the `stack`.
+    pub fn recycle(&mut self, stack: Stack) {
+        if self.stacks.len() < self.keep {
+            self.stacks.push(stack);
+        }
+    }
+}
+
+impl EngineInner {
+    /// Creates a new [`EngineInner`] with the given [`Config`].
+    fn new(config: &Config) -> Self {
+        Self {
+            res: RwLock::new(EngineResources::new(config)),
+            stacks: Mutex::new(EngineStacks::new(config)),
+        }
+    }
+
+    fn config(&self) -> Config {
+        self.res.read().config
+    }
+
+    fn alloc_func_type(&self, func_type: FuncType) -> DedupFuncType {
+        self.res.write().func_types.alloc_func_type(func_type)
+    }
+
+    fn alloc_func_body<I>(&self, len_locals: usize, max_stack_height: usize, insts: I) -> FuncBody
+    where
+        I: IntoIterator<Item = Instruction>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        self.res
+            .write()
+            .code_map
+            .alloc(len_locals, max_stack_height, insts)
+    }
+
+    fn resolve_func_type<F, R>(&self, func_type: DedupFuncType, f: F) -> R
+    where
+        F: FnOnce(&FuncType) -> R,
+    {
+        f(self.res.read().func_types.resolve_func_type(func_type))
+    }
+
+    #[cfg(test)]
+    fn resolve_inst(&self, func_body: FuncBody, index: usize) -> Option<Instruction> {
+        self.res
+            .read()
+            .code_map
+            .get_instr(func_body, index)
+            .copied()
+    }
+
+    fn execute_func<Params, Results>(
+        &self,
+        ctx: impl AsContextMut,
+        func: Func,
+        params: Params,
+        results: Results,
+    ) -> Result<<Results as CallResults>::Results, Trap>
+    where
+        Params: CallParams,
+        Results: CallResults,
+    {
+        let res = self.res.read();
+        let mut stack = self.stacks.lock().reuse_or_new();
+        let results =
+            EngineExecutor::new(&mut stack).execute_func(ctx, &res, func, params, results);
+        self.stacks.lock().recycle(stack);
+        results
+    }
+}
+
+/// Engine resources that are immutable during function execution.
+///
+/// Can be shared by multiple engine executors.
+#[derive(Debug)]
+pub struct EngineResources {
     /// The configuration with which the [`Engine`] has been created.
     config: Config,
-    /// The value and call stacks.
-    stack: Stack,
     /// Stores all Wasm function bodies that the interpreter is aware of.
     code_map: CodeMap,
     /// Deduplicated function types.
@@ -233,37 +345,29 @@ pub struct EngineInner {
     func_types: FuncTypeRegistry,
 }
 
-impl EngineInner {
-    /// Creates a new [`EngineInner`] with the given [`Config`].
-    pub fn new(config: &Config) -> Self {
+impl EngineResources {
+    /// Creates a new [`EngineResources`] with the given [`Config`].
+    fn new(config: &Config) -> Self {
         let engine_idx = EngineIdx::new();
         Self {
             config: *config,
-            stack: Stack::new(config.stack_limits()),
             code_map: CodeMap::default(),
             func_types: FuncTypeRegistry::new(engine_idx),
         }
     }
+}
 
-    /// Returns a shared reference to the [`Config`] of the [`Engine`].
-    pub fn config(&self) -> &Config {
-        &self.config
-    }
+/// The internal state of the `wasmi` engine.
+#[derive(Debug)]
+pub struct EngineExecutor<'stack> {
+    /// The value and call stacks.
+    stack: &'stack mut Stack,
+}
 
-    /// Allocates the instructions of a Wasm function body to the [`Engine`].
-    ///
-    /// Returns a [`FuncBody`] reference to the allocated function body.
-    pub fn alloc_func_body<I>(
-        &mut self,
-        len_locals: usize,
-        max_stack_height: usize,
-        insts: I,
-    ) -> FuncBody
-    where
-        I: IntoIterator<Item = Instruction>,
-        I::IntoIter: ExactSizeIterator,
-    {
-        self.code_map.alloc(len_locals, max_stack_height, insts)
+impl<'stack> EngineExecutor<'stack> {
+    /// Creates a new [`EngineExecutor`] with the given [`StackLimits`].
+    fn new(stack: &'stack mut Stack) -> Self {
+        Self { stack }
     }
 
     /// Executes the given [`Func`] using the given arguments `args` and stores the result into `results`.
@@ -273,9 +377,10 @@ impl EngineInner {
     /// - If the given arguments `args` do not match the expected parameters of `func`.
     /// - If the given `results` do not match the the length of the expected results of `func`.
     /// - When encountering a Wasm trap during the execution of `func`.
-    pub fn execute_func<Params, Results>(
+    fn execute_func<Params, Results>(
         &mut self,
         mut ctx: impl AsContextMut,
+        res: &EngineResources,
         func: Func,
         params: Params,
         results: Results,
@@ -287,14 +392,14 @@ impl EngineInner {
         self.initialize_args(params);
         match func.as_internal(ctx.as_context()) {
             FuncEntityInternal::Wasm(wasm_func) => {
-                let mut frame = self.stack.call_wasm_root(wasm_func, &self.code_map)?;
+                let mut frame = self.stack.call_wasm_root(wasm_func, &res.code_map)?;
                 let mut cache = InstanceCache::from(frame.instance());
-                self.execute_wasm_func(ctx.as_context_mut(), &mut frame, &mut cache)?;
+                self.execute_wasm_func(ctx.as_context_mut(), res, &mut frame, &mut cache)?;
             }
             FuncEntityInternal::Host(host_func) => {
                 let host_func = host_func.clone();
                 self.stack
-                    .call_host_root(ctx.as_context_mut(), host_func, &self.func_types)?;
+                    .call_host_root(ctx.as_context_mut(), host_func, &res.func_types)?;
             }
         };
         let results = self.write_results_back(results);
@@ -335,6 +440,7 @@ impl EngineInner {
     fn execute_wasm_func(
         &mut self,
         mut ctx: impl AsContextMut,
+        res: &EngineResources,
         frame: &mut FuncFrame,
         cache: &mut InstanceCache,
     ) -> Result<(), Trap> {
@@ -350,7 +456,7 @@ impl EngineInner {
                 CallOutcome::NestedCall(called_func) => {
                     match called_func.as_internal(ctx.as_context()) {
                         FuncEntityInternal::Wasm(wasm_func) => {
-                            *frame = self.stack.call_wasm(frame, wasm_func, &self.code_map)?;
+                            *frame = self.stack.call_wasm(frame, wasm_func, &res.code_map)?;
                         }
                         FuncEntityInternal::Host(host_func) => {
                             cache.reset_default_memory_bytes();
@@ -359,7 +465,7 @@ impl EngineInner {
                                 ctx.as_context_mut(),
                                 frame,
                                 host_func,
-                                &self.func_types,
+                                &res.func_types,
                             )?;
                         }
                     }