Skip to content

Commit

Permalink
Respect pooling allocation options in wasmtime serve
Browse files Browse the repository at this point in the history
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
  • Loading branch information
alexcrichton committed May 2, 2024
1 parent 2c40953 commit 4666af7
Show file tree
Hide file tree
Showing 7 changed files with 72 additions and 41 deletions.
8 changes: 6 additions & 2 deletions crates/cli-flags/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -421,7 +421,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();

Expand Down Expand Up @@ -546,7 +550,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();
Expand Down
2 changes: 1 addition & 1 deletion src/commands/compile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
2 changes: 1 addition & 1 deletion src/commands/explore.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(|| {
Expand Down
2 changes: 1 addition & 1 deletion src/commands/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
23 changes: 6 additions & 17 deletions src/commands/serve.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
2 changes: 1 addition & 1 deletion src/commands/wast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
74 changes: 56 additions & 18 deletions tests/all/cli_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1640,36 +1640,50 @@ 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) => {
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.
Expand Down Expand Up @@ -1775,6 +1789,30 @@ mod test_programs {
server.finish()?;
Ok(())
}

#[tokio::test]
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]
Expand Down

0 comments on commit 4666af7

Please sign in to comment.