Made Stack Size parametrable (for #3760)

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<PathBuf>,
 
+    /// The stack size (default is 1048576)
+    #[clap(long = "stack-size")]
+    pub(crate) stack_size: Option<usize>,
+
     /// Application arguments
     #[clap(value_name = "ARGS")]
     pub(crate) args: Vec<String>,
@@ -181,6 +185,9 @@ impl RunWithPathBuf {
     }
 
     fn inner_module_run(&self, store: &mut Store, instance: Instance) -> Result<i32> {
+        if self.stack_size.is_some() {
+            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

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<usize>,
     /// The function or command to invoke.
     #[clap(short, long, aliases = &["command", "invoke"])]
     entrypoint: Option<String>,
@@ -101,6 +104,9 @@ impl RunUnstable {
         module: &Module,
         store: &mut Store,
     ) -> Result<(), Error> {
+        if self.stack_size.is_some() {
+            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 +124,9 @@ impl RunUnstable {
         mut cache: ModuleCache,
         store: &mut Store,
     ) -> Result<(), Error> {
+        if self.stack_size.is_some() {
+            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())?,

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;

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,6 +29,13 @@ 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 DEFAULT_STACK_SIZE: AtomicUsize = AtomicUsize::new(1024 * 1024);
+
+/// Default stack size is 1MB.
+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! {
     if #[cfg(unix)] {
         /// Function which may handle custom signals while processing traps.
@@ -624,7 +631,12 @@ 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.load(Ordering::Relaxed),
+        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<F: FnOnce() -> T, T>(
+    stack_size: usize,
     trap_handler: Option<*const TrapHandlerFn<'static>>,
     f: F,
 ) -> Result<T, UnwindReason> {
@@ -851,7 +864,11 @@ fn on_wasm_stack<F: FnOnce() -> T, T>(
     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 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.