Merge #2824

2824: Use a pool of stack for running Wasm code r=syrusakbary a=Amanieu Previous we would allocate a new stack on every entry into Wasm code. This PR caches the stacks in a pool which allows them to be reused. Co-authored-by: Amanieu d'Antras <[email protected]>
wasmerio · Mar 18, 2022 · 02d27cd · 02d27cd
2 parents 704d935 + a1f6a22
commit 02d27cd
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 2 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/lib/vm/Cargo.toml b/lib/vm/Cargo.toml
@@ -26,6 +26,7 @@ loupe = { version = "0.1", features = ["enable-indexmap"] }
 enum-iterator = "0.7.0"
 corosensei = { version = "0.1.2" }
 scopeguard = "1.1.0"
+lazy_static = "1.4.0"
 
 [target.'cfg(target_vendor = "apple")'.dependencies]
 mach = "0.3.2"

diff --git a/lib/vm/src/trap/traphandlers.rs b/lib/vm/src/trap/traphandlers.rs
@@ -7,6 +7,7 @@
 use super::trapcode::TrapCode;
 use crate::vmcontext::{VMFunctionBody, VMFunctionEnvironment, VMTrampoline};
 use backtrace::Backtrace;
+use corosensei::stack::DefaultStack;
 use corosensei::trap::{CoroutineTrapHandler, TrapHandlerRegs};
 use corosensei::{CoroutineResult, ScopedCoroutine, Yielder};
 use scopeguard::defer;
@@ -19,7 +20,7 @@ use std::mem;
 use std::mem::MaybeUninit;
 use std::ptr::{self, NonNull};
 use std::sync::atomic::{compiler_fence, AtomicPtr, Ordering};
-use std::sync::Once;
+use std::sync::{Mutex, Once};
 
 cfg_if::cfg_if! {
     if #[cfg(unix)] {
@@ -835,8 +836,18 @@ fn on_wasm_stack<F: FnOnce() -> T, T>(
     trap_handler: &(dyn TrapHandler + 'static),
     f: F,
 ) -> Result<T, UnwindReason> {
+    // Allocating a new stack is pretty expensive since it involves several
+    // system calls. We therefore keep a cache of pre-allocated stacks which
+    // allows them to be reused multiple times.
+    // FIXME(Amanieu): We should refactor this to avoid the lock.
+    lazy_static::lazy_static! {
+        static ref STACK_POOL: Mutex<Vec<DefaultStack>> = Mutex::new(vec![]);
+    }
+    let stack = STACK_POOL.lock().unwrap().pop().unwrap_or_default();
+    let mut stack = scopeguard::guard(stack, |stack| STACK_POOL.lock().unwrap().push(stack));
+
     // Create a coroutine with a new stack to run the function on.
-    let mut coro = ScopedCoroutine::new(move |yielder, ()| {
+    let mut coro = ScopedCoroutine::with_stack(&mut *stack, move |yielder, ()| {
         // Save the yielder to TLS so that it can be used later.
         YIELDER.with(|cell| cell.set(Some(yielder.into())));