From 63488a4111262fd0955f05c2c29922b58bf1f7e5 Mon Sep 17 00:00:00 2001 From: ptitSeb Date: Wed, 19 Apr 2023 14:41:18 +0200 Subject: [PATCH 1/2] Made Stack Size parametrable (for #3760) --- lib/cli/src/commands/run.rs | 9 +++++++++ lib/cli/src/commands/run_unstable.rs | 13 +++++++++++++ lib/vm/src/trap/mod.rs | 4 ++-- lib/vm/src/trap/traphandlers.rs | 21 +++++++++++++++++++-- 4 files changed, 43 insertions(+), 4 deletions(-) diff --git a/lib/cli/src/commands/run.rs b/lib/cli/src/commands/run.rs index a784c4dea75..e2509393bf4 100644 --- a/lib/cli/src/commands/run.rs +++ b/lib/cli/src/commands/run.rs @@ -85,6 +85,10 @@ pub struct RunWithoutFile { #[clap(name = "COREDUMP PATH", long = "coredump-on-trap", parse(from_os_str))] coredump_on_trap: Option, + /// The stack size (default is 1048576) + #[clap(long = "stack-size")] + pub(crate) stack_size: Option, + /// Application arguments #[clap(value_name = "ARGS")] pub(crate) args: Vec, @@ -181,6 +185,11 @@ impl RunWithPathBuf { } fn inner_module_run(&self, store: &mut Store, instance: Instance) -> Result { + if self.stack_size.is_some() { + unsafe { + wasmer_vm::set_stack_size(self.stack_size.unwrap()); + } + } // If this module exports an _initialize function, run that first. if let Ok(initialize) = instance.exports.get_function("_initialize") { initialize diff --git a/lib/cli/src/commands/run_unstable.rs b/lib/cli/src/commands/run_unstable.rs index 1bc55f7d625..24379e45073 100644 --- a/lib/cli/src/commands/run_unstable.rs +++ b/lib/cli/src/commands/run_unstable.rs @@ -48,6 +48,9 @@ pub struct RunUnstable { wasi: crate::commands::run::Wasi, #[clap(flatten)] wcgi: WcgiOptions, + /// The stack size (default is 1048576) + #[clap(long = "stack-size")] + stack_size: Option, /// The function or command to invoke. #[clap(short, long, aliases = &["command", "invoke"])] entrypoint: Option, @@ -101,6 +104,11 @@ impl RunUnstable { module: &Module, store: &mut Store, ) -> Result<(), Error> { + if self.stack_size.is_some() { + unsafe { + wasmer_vm::set_stack_size(self.stack_size.unwrap()); + } + } if wasmer_emscripten::is_emscripten_module(module) { self.execute_emscripten_module() } else if wasmer_wasix::is_wasi_module(module) || wasmer_wasix::is_wasix_module(module) { @@ -118,6 +126,11 @@ impl RunUnstable { mut cache: ModuleCache, store: &mut Store, ) -> Result<(), Error> { + if self.stack_size.is_some() { + unsafe { + wasmer_vm::set_stack_size(self.stack_size.unwrap()); + } + } let id = match self.entrypoint.as_deref() { Some(cmd) => cmd, None => infer_webc_entrypoint(container.manifest())?, diff --git a/lib/vm/src/trap/mod.rs b/lib/vm/src/trap/mod.rs index 6a467c2d438..a9c31aeae96 100644 --- a/lib/vm/src/trap/mod.rs +++ b/lib/vm/src/trap/mod.rs @@ -10,8 +10,8 @@ mod traphandlers; pub use trap::Trap; pub use traphandlers::{ - catch_traps, on_host_stack, raise_lib_trap, raise_user_trap, wasmer_call_trampoline, - TrapHandlerFn, + catch_traps, on_host_stack, raise_lib_trap, raise_user_trap, set_stack_size, + wasmer_call_trampoline, TrapHandlerFn, }; pub use traphandlers::{init_traps, resume_panic}; pub use wasmer_types::TrapCode; diff --git a/lib/vm/src/trap/traphandlers.rs b/lib/vm/src/trap/traphandlers.rs index 89b9a141da5..d918a89c942 100644 --- a/lib/vm/src/trap/traphandlers.rs +++ b/lib/vm/src/trap/traphandlers.rs @@ -29,6 +29,18 @@ use wasmer_types::TrapCode; // On Arm64, the udf alows for a 16bits values, so we'll use the same 0xC? to store the trapinfo static MAGIC: u8 = 0xc0; +static mut DEFAULT_STACK_SIZE: usize = 1024 * 1024; + +/// Default stack size is 1MB. +/// # Safety +/// +/// This function should only be called before starting any wasm process +pub unsafe fn set_stack_size(size: usize) { + unsafe { + DEFAULT_STACK_SIZE = size.max(8 * 1024).min(100 * 1024 * 1024); + } +} + cfg_if::cfg_if! { if #[cfg(unix)] { /// Function which may handle custom signals while processing traps. @@ -624,7 +636,7 @@ where // Ensure that per-thread initialization is done. lazy_per_thread_init()?; - on_wasm_stack(trap_handler, closure).map_err(UnwindReason::into_trap) + on_wasm_stack(DEFAULT_STACK_SIZE, trap_handler, closure).map_err(UnwindReason::into_trap) } // We need two separate thread-local variables here: @@ -841,6 +853,7 @@ unsafe fn unwind_with(reason: UnwindReason) -> ! { /// bounded. Stack overflows and other traps can be caught and execution /// returned to the root of the stack. fn on_wasm_stack T, T>( + stack_size: usize, trap_handler: Option<*const TrapHandlerFn<'static>>, f: F, ) -> Result { @@ -851,7 +864,11 @@ fn on_wasm_stack T, T>( lazy_static::lazy_static! { static ref STACK_POOL: Mutex> = Mutex::new(vec![]); } - let stack = STACK_POOL.lock().unwrap().pop().unwrap_or_default(); + let stack = STACK_POOL + .lock() + .unwrap() + .pop() + .unwrap_or_else(|| DefaultStack::new(stack_size).unwrap()); 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. From 672f64b82fb106fa4556021de4931c3134b87592 Mon Sep 17 00:00:00 2001 From: ptitSeb Date: Wed, 19 Apr 2023 15:20:23 +0200 Subject: [PATCH 2/2] Remove unsafe and use Atomic instead --- lib/cli/src/commands/run.rs | 4 +--- lib/cli/src/commands/run_unstable.rs | 8 ++------ lib/vm/src/trap/traphandlers.rs | 20 ++++++++++---------- 3 files changed, 13 insertions(+), 19 deletions(-) diff --git a/lib/cli/src/commands/run.rs b/lib/cli/src/commands/run.rs index e2509393bf4..e84d068ca9a 100644 --- a/lib/cli/src/commands/run.rs +++ b/lib/cli/src/commands/run.rs @@ -186,9 +186,7 @@ impl RunWithPathBuf { fn inner_module_run(&self, store: &mut Store, instance: Instance) -> Result { if self.stack_size.is_some() { - unsafe { - wasmer_vm::set_stack_size(self.stack_size.unwrap()); - } + wasmer_vm::set_stack_size(self.stack_size.unwrap()); } // If this module exports an _initialize function, run that first. if let Ok(initialize) = instance.exports.get_function("_initialize") { diff --git a/lib/cli/src/commands/run_unstable.rs b/lib/cli/src/commands/run_unstable.rs index 24379e45073..9b3c9924062 100644 --- a/lib/cli/src/commands/run_unstable.rs +++ b/lib/cli/src/commands/run_unstable.rs @@ -105,9 +105,7 @@ impl RunUnstable { store: &mut Store, ) -> Result<(), Error> { if self.stack_size.is_some() { - unsafe { - wasmer_vm::set_stack_size(self.stack_size.unwrap()); - } + wasmer_vm::set_stack_size(self.stack_size.unwrap()); } if wasmer_emscripten::is_emscripten_module(module) { self.execute_emscripten_module() @@ -127,9 +125,7 @@ impl RunUnstable { store: &mut Store, ) -> Result<(), Error> { if self.stack_size.is_some() { - unsafe { - wasmer_vm::set_stack_size(self.stack_size.unwrap()); - } + wasmer_vm::set_stack_size(self.stack_size.unwrap()); } let id = match self.entrypoint.as_deref() { Some(cmd) => cmd, diff --git a/lib/vm/src/trap/traphandlers.rs b/lib/vm/src/trap/traphandlers.rs index d918a89c942..a5c75cf83a0 100644 --- a/lib/vm/src/trap/traphandlers.rs +++ b/lib/vm/src/trap/traphandlers.rs @@ -20,7 +20,7 @@ use std::mem; #[cfg(unix)] use std::mem::MaybeUninit; use std::ptr::{self, NonNull}; -use std::sync::atomic::{compiler_fence, AtomicPtr, Ordering}; +use std::sync::atomic::{compiler_fence, AtomicPtr, AtomicUsize, Ordering}; use std::sync::{Mutex, Once}; use wasmer_types::TrapCode; @@ -29,16 +29,11 @@ use wasmer_types::TrapCode; // On Arm64, the udf alows for a 16bits values, so we'll use the same 0xC? to store the trapinfo static MAGIC: u8 = 0xc0; -static mut DEFAULT_STACK_SIZE: usize = 1024 * 1024; +static DEFAULT_STACK_SIZE: AtomicUsize = AtomicUsize::new(1024 * 1024); /// Default stack size is 1MB. -/// # Safety -/// -/// This function should only be called before starting any wasm process -pub unsafe fn set_stack_size(size: usize) { - unsafe { - DEFAULT_STACK_SIZE = size.max(8 * 1024).min(100 * 1024 * 1024); - } +pub fn set_stack_size(size: usize) { + DEFAULT_STACK_SIZE.store(size.max(8 * 1024).min(100 * 1024 * 1024), Ordering::Relaxed); } cfg_if::cfg_if! { @@ -636,7 +631,12 @@ where // Ensure that per-thread initialization is done. lazy_per_thread_init()?; - on_wasm_stack(DEFAULT_STACK_SIZE, trap_handler, closure).map_err(UnwindReason::into_trap) + on_wasm_stack( + DEFAULT_STACK_SIZE.load(Ordering::Relaxed), + trap_handler, + closure, + ) + .map_err(UnwindReason::into_trap) } // We need two separate thread-local variables here: