Respect pooling allocation options in wasmtime serve (bytecodealliance#8525)

* Respect pooling allocation options in `wasmtime serve`

This commit updates the processing of pooling allocator options in
`wasmtime serve`. Previously the pooling allocator was enabled by
default but the options to configure it weren't processed due to how
this default-enable was implemented. The option to enable it by default
for `wasmtime serve`, but only `wasmtime serve`, is now processed
differently in a way that handles various other
pooling-allocator-related options.

Closes bytecodealliance#8504

* Fix compile of bench api

* Fix test build

* Ignore newly added test as it's flaky

Author: alexcrichton

crates/bench-api/src/lib.rs

@@ -435,7 +435,7 @@ impl BenchState {
         execution_end: extern "C" fn(*mut u8),
         make_wasi_cx: impl FnMut() -> Result<WasiCtx> + 'static,
     ) -> Result<Self> {
-        let mut config = options.config(Some(&Triple::host().to_string()))?;
+        let mut config = options.config(Some(&Triple::host().to_string()), None)?;
         // NB: always disable the compilation cache.
         config.disable_cache();
         let engine = Engine::new(&config)?;

crates/cli-flags/src/lib.rs

@@ -432,7 +432,11 @@ impl CommonOptions {
         Ok(())
     }
 
-    pub fn config(&mut self, target: Option<&str>) -> Result<Config> {
+    pub fn config(
+        &mut self,
+        target: Option<&str>,
+        pooling_allocator_default: Option<bool>,
+    ) -> Result<Config> {
         self.configure();
         let mut config = Config::new();
 
@@ -563,7 +567,7 @@ impl CommonOptions {
         }
 
         match_feature! {
-            ["pooling-allocator" : self.opts.pooling_allocator]
+            ["pooling-allocator" : self.opts.pooling_allocator.or(pooling_allocator_default)]
             enable => {
                 if enable {
                     let mut cfg = wasmtime::PoolingAllocationConfig::default();

src/commands/compile.rs

@@ -61,7 +61,7 @@ impl CompileCommand {
     pub fn execute(mut self) -> Result<()> {
         self.common.init_logging()?;
 
-        let mut config = self.common.config(self.target.as_deref())?;
+        let mut config = self.common.config(self.target.as_deref(), None)?;
 
         if let Some(path) = self.emit_clif {
             if !path.exists() {

src/commands/explore.rs

@@ -30,7 +30,7 @@ impl ExploreCommand {
     pub fn execute(mut self) -> Result<()> {
         self.common.init_logging()?;
 
-        let config = self.common.config(self.target.as_deref())?;
+        let config = self.common.config(self.target.as_deref(), None)?;
 
         let bytes =
             Cow::Owned(std::fs::read(&self.module).with_context(|| {

src/commands/run.rs

@@ -74,7 +74,7 @@ impl RunCommand {
     pub fn execute(mut self) -> Result<()> {
         self.run.common.init_logging()?;
 
-        let mut config = self.run.common.config(None)?;
+        let mut config = self.run.common.config(None, None)?;
 
         if self.run.common.wasm.timeout.is_some() {
             config.epoch_interruption(true);

src/commands/serve.rs

@@ -247,23 +247,12 @@ impl ServeCommand {
     async fn serve(mut self) -> Result<()> {
         use hyper::server::conn::http1;
 
-        let mut config = self.run.common.config(None)?;
-        match self.run.common.opts.pooling_allocator {
-            // If explicitly enabled on the CLI then the pooling allocator was
-            // already configured in the `config` method above. If the allocator
-            // is explicitly disabled, then we don't want it. In both cases do
-            // nothing.
-            Some(true) | Some(false) => {}
-
-            // Otherwise though if not explicitly specified then always enable
-            // the pooling allocator. The `wasmtime serve` use case is
-            // tailor-made for pooling allocation and there's no downside to
-            // enabling it.
-            None => {
-                let cfg = wasmtime::PoolingAllocationConfig::default();
-                config.allocation_strategy(wasmtime::InstanceAllocationStrategy::Pooling(cfg));
-            }
-        }
+        // If not explicitly specified then always enable the pooling allocator.
+        // The `wasmtime serve` use case is tailor-made for pooling allocation
+        // and there's not much downside to enabling it.
+        let pooling_allocator_default = Some(true);
+
+        let mut config = self.run.common.config(None, pooling_allocator_default)?;
         config.wasm_component_model(true);
         config.async_support(true);
 

src/commands/wast.rs

@@ -23,7 +23,7 @@ impl WastCommand {
     pub fn execute(mut self) -> Result<()> {
         self.common.init_logging()?;
 
-        let config = self.common.config(None)?;
+        let config = self.common.config(None, None)?;
         let store = Store::new(&Engine::new(&config)?, ());
         let mut wast_context = WastContext::new(store);
 

tests/all/cli_tests.rs

@@ -1640,36 +1640,51 @@ mod test_programs {
             // all remaining output is left to be captured by future requests
             // send to the server.
             let mut line = String::new();
-            BufReader::new(child.stderr.as_mut().unwrap()).read_line(&mut line)?;
-            let addr_start = line.find("127.0.0.1").unwrap();
-            let addr = &line[addr_start..];
-            let addr_end = addr.find("/").unwrap();
-            let addr = &addr[..addr_end];
-            Ok(WasmtimeServe {
-                child: Some(child),
-                addr: addr.parse().unwrap(),
-            })
+            let mut reader = BufReader::new(child.stderr.take().unwrap());
+            reader.read_line(&mut line)?;
+
+            match line.find("127.0.0.1").and_then(|addr_start| {
+                let addr = &line[addr_start..];
+                let addr_end = addr.find("/")?;
+                addr[..addr_end].parse().ok()
+            }) {
+                Some(addr) => {
+                    assert!(reader.buffer().is_empty());
+                    child.stderr = Some(reader.into_inner());
+                    Ok(WasmtimeServe {
+                        child: Some(child),
+                        addr,
+                    })
+                }
+                None => {
+                    child.kill()?;
+                    child.wait()?;
+                    reader.read_to_string(&mut line)?;
+                    bail!("failed to start child: {line}")
+                }
+            }
         }
 
         /// Completes this server gracefully by printing the output on failure.
-        fn finish(mut self) -> Result<()> {
+        fn finish(mut self) -> Result<String> {
             let mut child = self.child.take().unwrap();
 
             // If the child process has already exited then collect the output
             // and test if it succeeded. Otherwise it's still running so kill it
             // and then reap it. Assume that if it's still running then the test
             // has otherwise passed so no need to print the output.
-            if child.try_wait()?.is_some() {
-                let output = child.wait_with_output()?;
-                if output.status.success() {
-                    return Ok(());
-                }
-                bail!("child failed {output:?}");
+            let known_failure = if child.try_wait()?.is_some() {
+                false
             } else {
                 child.kill()?;
-                child.wait_with_output()?;
+                true
+            };
+            let output = child.wait_with_output()?;
+            if !known_failure && !output.status.success() {
+                bail!("child failed {output:?}");
             }
-            Ok(())
+
+            Ok(String::from_utf8_lossy(&output.stderr).into_owned())
         }
 
         /// Send a request to this server and wait for the response.
@@ -1775,6 +1790,34 @@ mod test_programs {
         server.finish()?;
         Ok(())
     }
+
+    #[tokio::test]
+    #[ignore] // TODO: printing stderr in the child and killing the child at the
+              // end of this test race so the stderr may be present or not. Need
+              // to implement a more graceful shutdown routine for `wasmtime
+              // serve`.
+    async fn cli_serve_respect_pooling_options() -> Result<()> {
+        let server = WasmtimeServe::new(CLI_SERVE_ECHO_ENV_COMPONENT, |cmd| {
+            cmd.arg("-Opooling-total-memories=0").arg("-Scli");
+        })?;
+
+        let result = server
+            .send_request(
+                hyper::Request::builder()
+                    .uri("http://localhost/")
+                    .header("env", "FOO")
+                    .body(String::new())
+                    .context("failed to make request")?,
+            )
+            .await;
+        assert!(result.is_err());
+        let stderr = server.finish()?;
+        assert!(
+            stderr.contains("maximum concurrent memory limit of 0 reached"),
+            "bad stderr: {stderr}",
+        );
+        Ok(())
+    }
 }
 
 #[test]

tests/disas.rs

@@ -190,7 +190,7 @@ impl Test {
         // Use wasmtime::Config with its `emit_clif` option to get Wasmtime's
         // code generator to jettison CLIF out the back.
         let tempdir = TempDir::new().context("failed to make a tempdir")?;
-        let mut config = self.opts.config(Some(&self.config.target))?;
+        let mut config = self.opts.config(Some(&self.config.target), None)?;
         match self.config.test {
             TestKind::Clif | TestKind::Optimize => {
                 config.emit_clif(tempdir.path());