From 57efd6087933f3e642bc89afd1f838693dfb5b90 Mon Sep 17 00:00:00 2001 From: Andrew Brown Date: Fri, 1 Dec 2023 16:46:12 -0800 Subject: [PATCH] mpk: add an example testing the memory limits (#7609) * mpk: allow checking for MPK without a config instance It is inconvenient to have to construct a `PoolingAllocationConfig` in order to check if memory protection keys are available. This removes the unused `&self` restriction. * mpk: improve logging of calculated slab layout When double-checking the slab layout calculations it is quite convenient to see the total slab size. This helps in correlating with mapped regions. * mpk: add an example testing the memory limits This adds an example that can be run with `cargo run --example mpk`. Not only does the example demonstrate how to build a pool-allocated engine that uses MPK, it performs an exponential search to find the maximum number of slots the system can support, with and without MPK. * review: document Linux requirement * review: `env_logger::init` * review: replace `proc-maps` with manual parsing * vet: audit `bytesize` * fix: provide `main` for non-Linux systems * fix: move `cfg` to avoid unused code --- Cargo.lock | 43 +-- Cargo.toml | 1 + .../instance/allocator/pooling/memory_pool.rs | 5 +- crates/wasmtime/src/config.rs | 2 +- examples/mpk.rs | 264 ++++++++++++++++++ supply-chain/audits.toml | 5 + supply-chain/imports.lock | 7 + 7 files changed, 307 insertions(+), 20 deletions(-) create mode 100644 examples/mpk.rs diff --git a/Cargo.lock b/Cargo.lock index 8630e5ebe2d4..ae5e9bbb3e21 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -121,7 +121,7 @@ checksum = "a564d521dd56509c4c47480d00b80ee55f7e385ae48db5744c67ad50c92d2ebf" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", ] [[package]] @@ -240,6 +240,12 @@ version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "89b2fd2a0dcf38d7971e2194b6b6eebab45ae01067456a7fd93d5547a61b70be" +[[package]] +name = "bytesize" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3e368af43e418a04d52505cf3dbc23dda4e3407ae2fa99fd0e4f308ce546acc" + [[package]] name = "camino" version = "1.1.4" @@ -459,7 +465,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", ] [[package]] @@ -501,7 +507,7 @@ version = "0.0.0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", ] [[package]] @@ -947,7 +953,7 @@ checksum = "53e0efad4403bfc52dc201159c4b842a246a14b98c64b55dfd0f2d89729dfeb8" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", ] [[package]] @@ -2332,7 +2338,7 @@ checksum = "4eca7ac642d82aa35b60049a6eccb4be6be75e599bd2e9adb5f875a737654af2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", ] [[package]] @@ -2509,9 +2515,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.29" +version = "2.0.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c324c494eba9d92503e6f1ef2e6df781e78f6a7705a0202d9801b198807d518a" +checksum = "239814284fd6f1a4ffe4ca893952cdd93c224b6a1571c9a9eadd670295c0c9e2" dependencies = [ "proc-macro2", "quote", @@ -2624,7 +2630,7 @@ checksum = "463fe12d7993d3b327787537ce8dd4dfa058de32fc2b195ef3cde03dc4771e8f" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", ] [[package]] @@ -2688,7 +2694,7 @@ checksum = "630bdcf245f78637c13ec01ffae6187cca34625e8c63150d424b59e55af2675e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", ] [[package]] @@ -2745,7 +2751,7 @@ checksum = "5f4f31f56159e98206da9efd823404b79b6ef3143b4a7ab76e67b1751b25a4ab" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", ] [[package]] @@ -3032,7 +3038,7 @@ dependencies = [ "once_cell", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", "wasm-bindgen-shared", ] @@ -3054,7 +3060,7 @@ checksum = "54681b18a46765f095758388f2d0cf16eb8d4169b639ab575a8f5693af210c7b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", "wasm-bindgen-backend", "wasm-bindgen-shared", ] @@ -3313,6 +3319,7 @@ dependencies = [ "async-trait", "bstr", "bytes", + "bytesize", "clap", "component-macro-test", "component-test-util", @@ -3380,7 +3387,7 @@ dependencies = [ "component-macro-test-helpers", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", "tracing", "wasmtime", "wasmtime-component-util", @@ -3641,7 +3648,7 @@ version = "16.0.0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", ] [[package]] @@ -3855,7 +3862,7 @@ dependencies = [ "proc-macro2", "quote", "shellexpand", - "syn 2.0.29", + "syn 2.0.32", "witx", ] @@ -3865,7 +3872,7 @@ version = "16.0.0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", "wiggle", "wiggle-generate", ] @@ -3954,7 +3961,7 @@ dependencies = [ "glob", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", ] [[package]] @@ -4097,7 +4104,7 @@ dependencies = [ "anyhow", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.32", "wit-bindgen-core", "wit-bindgen-rust", "wit-component 0.18.2", diff --git a/Cargo.toml b/Cargo.toml index b18a39960c2e..071809bdfd63 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -89,6 +89,7 @@ serde = { workspace = true } serde_json = { workspace = true } walkdir = { workspace = true } test-programs-artifacts = { workspace = true } +bytesize = "1.3.0" [target.'cfg(windows)'.dev-dependencies] windows-sys = { workspace = true, features = ["Win32_System_Memory"] } diff --git a/crates/runtime/src/instance/allocator/pooling/memory_pool.rs b/crates/runtime/src/instance/allocator/pooling/memory_pool.rs index 25c000101dc1..3bf096927974 100644 --- a/crates/runtime/src/instance/allocator/pooling/memory_pool.rs +++ b/crates/runtime/src/instance/allocator/pooling/memory_pool.rs @@ -189,7 +189,10 @@ impl MemoryPool { // region to start--`PROT_NONE`. let constraints = SlabConstraints::new(&config.limits, tunables, pkeys.len())?; let layout = calculate(&constraints)?; - log::debug!("creating memory pool: {constraints:?} -> {layout:?}"); + log::debug!( + "creating memory pool: {constraints:?} -> {layout:?} (total: {})", + layout.total_slab_bytes()? + ); let mut mapping = Mmap::accessible_reserved(0, layout.total_slab_bytes()?) .context("failed to create memory pool mapping")?; diff --git a/crates/wasmtime/src/config.rs b/crates/wasmtime/src/config.rs index cc42af4e3151..eb149561965f 100644 --- a/crates/wasmtime/src/config.rs +++ b/crates/wasmtime/src/config.rs @@ -2409,7 +2409,7 @@ impl PoolingAllocationConfig { /// same method that [`MpkEnabled::Auto`] does. See /// [`PoolingAllocationConfig::memory_protection_keys`] for more /// information. - pub fn are_memory_protection_keys_available(&self) -> bool { + pub fn are_memory_protection_keys_available() -> bool { mpk::is_supported() } } diff --git a/examples/mpk.rs b/examples/mpk.rs new file mode 100644 index 000000000000..05300f96c326 --- /dev/null +++ b/examples/mpk.rs @@ -0,0 +1,264 @@ +//! This example demonstrates: +//! - how to enable memory protection keys (MPK) in a Wasmtime embedding (see +//! [`build_engine`]) +//! - the expected memory compression from using MPK: it will probe the system +//! by creating larger and larger memory pools until system memory is +//! exhausted (see [`probe_engine_size`]). Then, it prints a comparison of the +//! memory used in both the MPK enabled and MPK disabled configurations. +//! +//! You can execute this example with: +//! +//! ```console +//! $ cargo run --example mpk +//! ``` +//! +//! Append `-- --help` for details about the configuring the memory size of the +//! pool. Also, to inspect interesting configuration values used for +//! constructing the pool, turn on logging: +//! +//! ```console +//! $ RUST_LOG=debug cargo run --example mpk -- --memory-size 512MiB +//! ``` +//! +//! Note that MPK support is limited to x86 Linux systems. OS limits on the +//! number of virtual memory areas (VMAs) can significantly restrict the total +//! number MPK-striped memory slots; each MPK-protected slot ends up using a new +//! VMA entry. On Linux, one can raise this limit: +//! +//! ```console +//! $ sysctl vm.max_map_count +//! 65530 +//! $ sysctl vm.max_map_count=$LARGER_LIMIT +//! ``` + +use anyhow::{anyhow, Result}; +use bytesize::ByteSize; +use clap::Parser; +use log::{info, warn}; +use std::str::FromStr; +use wasmtime::*; + +fn main() -> Result<()> { + env_logger::init(); + let args = Args::parse(); + info!("{:?}", args); + + let without_mpk = probe_engine_size(&args, MpkEnabled::Disable)?; + println!("without MPK:\t{}", without_mpk.to_string()); + + if PoolingAllocationConfig::are_memory_protection_keys_available() { + let with_mpk = probe_engine_size(&args, MpkEnabled::Enable)?; + println!("with MPK:\t{}", with_mpk.to_string()); + println!( + "\t\t{}x more slots per reserved memory", + with_mpk.compare(&without_mpk) + ); + } else { + println!("with MPK:\tunavailable\t\tunavailable"); + } + + Ok(()) +} + +#[derive(Debug, Parser)] +#[command(author, version, about, long_about = None)] +struct Args { + /// The maximum number of bytes for each WebAssembly linear memory in the + /// pool. + #[arg(long, default_value = "128MiB", value_parser = parse_byte_size)] + memory_size: u64, + + /// The maximum number of bytes a memory is considered static; see + /// `Config::static_memory_maximum_size` for more details and the default + /// value if unset. + #[arg(long, value_parser = parse_byte_size)] + static_memory_maximum_size: Option, + + /// The size in bytes of the guard region to expect between static memory + /// slots; see [`Config::static_memory_guard_size`] for more details and the + /// default value if unset. + #[arg(long, value_parser = parse_byte_size)] + static_memory_guard_size: Option, +} + +/// Parse a human-readable byte size--e.g., "512 MiB"--into the correct number +/// of bytes. +fn parse_byte_size(value: &str) -> Result { + let size = ByteSize::from_str(value).map_err(|e| anyhow!(e))?; + Ok(size.as_u64()) +} + +/// Find the engine with the largest number of memories we can create on this +/// machine. +fn probe_engine_size(args: &Args, mpk: MpkEnabled) -> Result { + let mut search = ExponentialSearch::new(); + let mut mapped_bytes = 0; + while !search.done() { + match build_engine(&args, search.next(), mpk) { + Ok(rb) => { + // TODO: assert!(rb >= mapped_bytes); + mapped_bytes = rb; + search.record(true) + } + Err(e) => { + warn!("failed engine allocation, continuing search: {:?}", e); + search.record(false) + } + } + } + Ok(Pool { + num_memories: search.next(), + mapped_bytes, + }) +} + +#[derive(Debug)] +#[allow(dead_code)] +struct Pool { + num_memories: u32, + mapped_bytes: usize, +} +impl Pool { + /// Print a human-readable, tab-separated description of this structure. + fn to_string(&self) -> String { + let human_size = ByteSize::b(self.mapped_bytes as u64).to_string_as(true); + format!( + "{} memory slots\t{} reserved", + self.num_memories, human_size + ) + } + /// Return the number of times more memory slots in `self` than `other` + /// after normalizing by the mapped bytes sizes. Rounds to three decimal + /// places arbitrarily; no significance intended. + fn compare(&self, other: &Pool) -> f64 { + let size_ratio = other.mapped_bytes as f64 / self.mapped_bytes as f64; + let slots_ratio = self.num_memories as f64 / other.num_memories as f64; + let times_more_efficient = slots_ratio * size_ratio; + (times_more_efficient * 1000.0).round() / 1000.0 + } +} + +/// Exponentially increase the `next` value until the attempts fail, then +/// perform a binary search to find the maximum attempted value that still +/// succeeds. +#[derive(Debug)] +struct ExponentialSearch { + /// Determines if we are in the growth phase. + growing: bool, + /// The last successful value tried; this is the algorithm's lower bound. + last: u32, + /// The next value to try; this is the algorithm's upper bound. + next: u32, +} +impl ExponentialSearch { + fn new() -> Self { + Self { + growing: true, + last: 0, + next: 1, + } + } + fn next(&self) -> u32 { + self.next + } + fn record(&mut self, success: bool) { + if !success { + self.growing = false + } + let diff = if self.growing { + (self.next - self.last) * 2 + } else { + (self.next - self.last + 1) / 2 + }; + if success { + self.last = self.next; + self.next = self.next + diff; + } else { + self.next = self.next - diff; + } + } + fn done(&self) -> bool { + self.last == self.next + } +} + +/// Build a pool-allocated engine with `num_memories` slots. +fn build_engine(args: &Args, num_memories: u32, enable_mpk: MpkEnabled) -> Result { + // Configure the memory pool. + let mut pool = PoolingAllocationConfig::default(); + let memory_pages = args.memory_size / u64::from(wasmtime_environ::WASM_PAGE_SIZE); + pool.memory_pages(memory_pages); + pool.total_memories(num_memories) + .memory_protection_keys(enable_mpk); + + // Configure the engine itself. + let mut config = Config::new(); + if let Some(static_memory_maximum_size) = args.static_memory_maximum_size { + config.static_memory_maximum_size(static_memory_maximum_size); + } + if let Some(static_memory_guard_size) = args.static_memory_guard_size { + config.static_memory_guard_size(static_memory_guard_size); + } + config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); + + // Measure memory use before and after the engine is built. + let mapped_bytes_before = num_bytes_mapped()?; + let engine = Engine::new(&config)?; + let mapped_bytes_after = num_bytes_mapped()?; + + // Ensure we actually use the engine somehow. + engine.increment_epoch(); + + let mapped_bytes = mapped_bytes_after - mapped_bytes_before; + info!( + "{}-slot pool ({:?}): {} bytes mapped", + num_memories, enable_mpk, mapped_bytes + ); + Ok(mapped_bytes) +} + +/// Add up the sizes of all the mapped virtual memory regions for the current +/// Linux process. +/// +/// This manually parses `/proc/self/maps` to avoid a rather-large `proc-maps` +/// dependency. We do expect this example to be Linux-specific anyways. For +/// reference, lines of that file look like: +/// +/// ```text +/// 5652d4418000-5652d441a000 r--p 00000000 00:23 84629427 /usr/bin/... +/// ``` +/// +/// We parse the start and end addresses: - [ignore the rest]. +#[cfg(target_os = "linux")] +fn num_bytes_mapped() -> Result { + use std::fs::File; + use std::io::{BufRead, BufReader}; + + let file = File::open("/proc/self/maps")?; + let reader = BufReader::new(file); + let mut total = 0; + for line in reader.lines() { + let line = line?; + let range = line + .split_whitespace() + .next() + .ok_or(anyhow!("parse failure: expected whitespace"))?; + let mut addresses = range.split("-"); + let start = addresses + .next() + .ok_or(anyhow!("parse failure: expected dash-separated address"))?; + let start = usize::from_str_radix(start, 16)?; + let end = addresses + .next() + .ok_or(anyhow!("parse failure: expected dash-separated address"))?; + let end = usize::from_str_radix(end, 16)?; + + total += end - start; + } + Ok(total) +} + +#[cfg(not(target_os = "linux"))] +fn num_bytes_mapped() -> Result { + anyhow::bail!("this example can only read virtual memory maps on Linux") +} diff --git a/supply-chain/audits.toml b/supply-chain/audits.toml index b4d557e07bd0..064868b341ba 100644 --- a/supply-chain/audits.toml +++ b/supply-chain/audits.toml @@ -833,6 +833,11 @@ criteria = "safe-to-deploy" version = "3.11.1" notes = "I am the author of this crate." +[[audits.bytesize]] +who = "Andrew Brown " +criteria = "safe-to-deploy" +version = "1.3.0" + [[audits.camino]] who = "Pat Hickey " criteria = "safe-to-deploy" diff --git a/supply-chain/imports.lock b/supply-chain/imports.lock index 632f38005ea2..fe9e858e9038 100644 --- a/supply-chain/imports.lock +++ b/supply-chain/imports.lock @@ -1013,6 +1013,13 @@ user-id = 3618 user-login = "dtolnay" user-name = "David Tolnay" +[[publisher.syn]] +version = "2.0.32" +when = "2023-09-10" +user-id = 3618 +user-login = "dtolnay" +user-name = "David Tolnay" + [[publisher.system-interface]] version = "0.26.0" when = "2023-06-30"