From 611f79d261a4f0ee11c10d3aa590df0af2f8b1ee Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Tue, 24 Sep 2024 15:47:06 -0500
Subject: [PATCH] Add `-Wasync-stack-size`, automatically configure it too
 (#9302) (#9305)

* Add `-Wasync-stack-size`, automatically configure it too

Recently the `wasmtime` CLI switched from sync to async to handle
interrupting host APIs in WASI. Previously the CLI provided no means
to configure the async stack size, however, which now means that when
increasing the max wasm stack the CLI prints an error that cannot be
resolved about the wasm stack being larger than the async stack. This
commit fixes this issue by both adding a configuration option for the
async stack size and additionally automatically increasing the async
stack size when only the wasm stack size is provided.

Closes #9298

* Fix feature-gated build
---
 Cargo.toml                  | 17 +++++++++++++++--
 crates/cli-flags/src/lib.rs | 21 +++++++++++++++++++++
 tests/all/cli_tests.rs      | 15 +++++++++++++++
 3 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index e7f58b73bf54..dc4e0fe5e54c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -435,12 +435,25 @@ gc = ["wasmtime-cli-flags/gc"]
 
 # CLI subcommands for the `wasmtime` executable. See `wasmtime $cmd --help`
 # for more information on each subcommand.
-serve = ["wasi-http", "component-model", "dep:http-body-util", "dep:http"]
+serve = [
+  "wasi-http",
+  "component-model",
+  "dep:http-body-util",
+  "dep:http",
+  "wasmtime-cli-flags/async",
+]
 explore = ["dep:wasmtime-explorer", "dep:tempfile"]
 wast = ["dep:wasmtime-wast"]
 config = ["cache"]
 compile = ["cranelift"]
-run = ["dep:wasmtime-wasi", "wasmtime/runtime", "dep:listenfd", "dep:wasi-common", "dep:tokio"]
+run = [
+  "dep:wasmtime-wasi",
+  "wasmtime/runtime",
+  "dep:listenfd",
+  "dep:wasi-common",
+  "dep:tokio",
+  "wasmtime-cli-flags/async",
+]
 
 [[test]]
 name = "host_segfault"
diff --git a/crates/cli-flags/src/lib.rs b/crates/cli-flags/src/lib.rs
index 8a7bb6dc9b6b..adeb216ec0da 100644
--- a/crates/cli-flags/src/lib.rs
+++ b/crates/cli-flags/src/lib.rs
@@ -199,6 +199,12 @@ wasmtime_option_group! {
         /// Maximum stack size, in bytes, that wasm is allowed to consume before a
         /// stack overflow is reported.
         pub max_wasm_stack: Option<usize>,
+        /// Stack size, in bytes, that will be allocated for async stacks.
+        ///
+        /// Note that this must be larger than `max-wasm-stack` and the
+        /// difference between the two is how much stack the host has to execute
+        /// on.
+        pub async_stack_size: Option<usize>,
         /// Allow unknown exports when running commands.
         pub unknown_exports_allow: Option<bool>,
         /// Allow the main module to import unknown functions, using an
@@ -652,8 +658,23 @@ impl CommonOptions {
             anyhow::bail!("memory protection keys require the pooling allocator");
         }
 
+        match_feature! {
+            ["async" : self.wasm.async_stack_size]
+            size => config.async_stack_size(size),
+            _ => err,
+        }
+
         if let Some(max) = self.wasm.max_wasm_stack {
             config.max_wasm_stack(max);
+
+            // If `-Wasync-stack-size` isn't passed then automatically adjust it
+            // to the wasm stack size provided here too. That prevents the need
+            // to pass both when one can generally be inferred from the other.
+            #[cfg(feature = "async")]
+            if self.wasm.async_stack_size.is_none() {
+                const DEFAULT_HOST_STACK: usize = 512 << 10;
+                config.async_stack_size(max + DEFAULT_HOST_STACK);
+            }
         }
 
         if let Some(enable) = self.wasm.relaxed_simd_deterministic {
diff --git a/tests/all/cli_tests.rs b/tests/all/cli_tests.rs
index dfb31842a91b..f93ea38e81a7 100644
--- a/tests/all/cli_tests.rs
+++ b/tests/all/cli_tests.rs
@@ -1093,6 +1093,21 @@ fn mpk_without_pooling() -> Result<()> {
     Ok(())
 }
 
+// Very basic use case: compile binary wasm file and run specific function with arguments.
+#[test]
+fn increase_stack_size() -> Result<()> {
+    run_wasmtime(&[
+        "run",
+        "--invoke",
+        "simple",
+        &format!("-Wmax-wasm-stack={}", 5 << 20),
+        "-Ccache=n",
+        "tests/all/cli_tests/simple.wat",
+        "4",
+    ])?;
+    Ok(())
+}
+
 mod test_programs {
     use super::{get_wasmtime_command, run_wasmtime};
     use anyhow::{bail, Context, Result};