1 //! This example demonstrates:
2 //! - how to enable memory protection keys (MPK) in a Wasmtime embedding (see
3 //! [`build_engine`])
4 //! - the expected memory compression from using MPK: it will probe the system
5 //! by creating larger and larger memory pools until system memory is
6 //! exhausted (see [`probe_engine_size`]). Then, it prints a comparison of the
7 //! memory used in both the MPK enabled and MPK disabled configurations.
8 //!
9 //! You can execute this example with:
10 //!
11 //! ```console
12 //! $ cargo run --example mpk
13 //! ```
14 //!
15 //! Append `-- --help` for details about the configuring the memory size of the
16 //! pool. Also, to inspect interesting configuration values used for
17 //! constructing the pool, turn on logging:
18 //!
19 //! ```console
20 //! $ RUST_LOG=debug cargo run --example mpk -- --memory-size 512MiB
21 //! ```
22 //!
23 //! Note that MPK support is limited to x86 Linux systems. OS limits on the
24 //! number of virtual memory areas (VMAs) can significantly restrict the total
25 //! number MPK-striped memory slots; each MPK-protected slot ends up using a new
26 //! VMA entry. On Linux, one can raise this limit:
27 //!
28 //! ```console
29 //! $ sysctl vm.max_map_count
30 //! 65530
31 //! $ sysctl vm.max_map_count=$LARGER_LIMIT
32 //! ```
33
34 use bytesize::ByteSize;
35 use clap::Parser;
36 use log::{info, warn};
37 use std::str::FromStr;
38 use wasmtime::format_err;
39 use wasmtime::*;
40
main() -> Result<()>41 fn main() -> Result<()> {
42 env_logger::init();
43 let args = Args::parse();
44 info!("{args:?}");
45
46 let without_mpk = probe_engine_size(&args, Enabled::No)?;
47 println!("without MPK:\t{}", without_mpk.to_string());
48
49 if PoolingAllocationConfig::are_memory_protection_keys_available() {
50 let with_mpk = probe_engine_size(&args, Enabled::Yes)?;
51 println!("with MPK:\t{}", with_mpk.to_string());
52 println!(
53 "\t\t{}x more slots per reserved memory",
54 with_mpk.compare(&without_mpk)
55 );
56 } else {
57 println!("with MPK:\tunavailable\t\tunavailable");
58 }
59
60 Ok(())
61 }
62
63 #[derive(Debug, Parser)]
64 #[command(author, version, about, long_about = None)]
65 struct Args {
66 /// The maximum number of bytes for each WebAssembly linear memory in the
67 /// pool.
68 #[arg(long, default_value = "128MiB", value_parser = parse_byte_size)]
69 memory_size: u64,
70
71 /// The maximum number of bytes a memory is considered static; see
72 /// `Config::memory_reservation` for more details and the default
73 /// value if unset.
74 #[arg(long, value_parser = parse_byte_size)]
75 memory_reservation: Option<u64>,
76
77 /// The size in bytes of the guard region to expect between static memory
78 /// slots; see [`Config::memory_guard_size`] for more details and the
79 /// default value if unset.
80 #[arg(long, value_parser = parse_byte_size)]
81 memory_guard_size: Option<u64>,
82 }
83
84 /// Parse a human-readable byte size--e.g., "512 MiB"--into the correct number
85 /// of bytes.
parse_byte_size(value: &str) -> Result<u64>86 fn parse_byte_size(value: &str) -> Result<u64> {
87 let size = ByteSize::from_str(value).map_err(|e| format_err!(e))?;
88 Ok(size.as_u64())
89 }
90
91 /// Find the engine with the largest number of memories we can create on this
92 /// machine.
probe_engine_size(args: &Args, mpk: Enabled) -> Result<Pool>93 fn probe_engine_size(args: &Args, mpk: Enabled) -> Result<Pool> {
94 let mut search = ExponentialSearch::new();
95 let mut mapped_bytes = 0;
96 while !search.done() {
97 match build_engine(&args, search.next(), mpk) {
98 Ok(rb) => {
99 // TODO: assert!(rb >= mapped_bytes);
100 mapped_bytes = rb;
101 search.record(true)
102 }
103 Err(e) => {
104 warn!("failed engine allocation, continuing search: {e:?}");
105 search.record(false)
106 }
107 }
108 }
109 Ok(Pool {
110 num_memories: search.next(),
111 mapped_bytes,
112 })
113 }
114
115 #[derive(Debug)]
116 struct Pool {
117 num_memories: u32,
118 mapped_bytes: usize,
119 }
120 impl Pool {
121 /// Print a human-readable, tab-separated description of this structure.
to_string(&self) -> String122 fn to_string(&self) -> String {
123 let human_size = ByteSize::b(self.mapped_bytes as u64).display().si();
124 format!(
125 "{} memory slots\t{} reserved",
126 self.num_memories, human_size
127 )
128 }
129 /// Return the number of times more memory slots in `self` than `other`
130 /// after normalizing by the mapped bytes sizes. Rounds to three decimal
131 /// places arbitrarily; no significance intended.
compare(&self, other: &Pool) -> f64132 fn compare(&self, other: &Pool) -> f64 {
133 let size_ratio = other.mapped_bytes as f64 / self.mapped_bytes as f64;
134 let slots_ratio = self.num_memories as f64 / other.num_memories as f64;
135 let times_more_efficient = slots_ratio * size_ratio;
136 (times_more_efficient * 1000.0).round() / 1000.0
137 }
138 }
139
140 /// Exponentially increase the `next` value until the attempts fail, then
141 /// perform a binary search to find the maximum attempted value that still
142 /// succeeds.
143 #[derive(Debug)]
144 struct ExponentialSearch {
145 /// Determines if we are in the growth phase.
146 growing: bool,
147 /// The last successful value tried; this is the algorithm's lower bound.
148 last: u32,
149 /// The next value to try; this is the algorithm's upper bound.
150 next: u32,
151 }
152 impl ExponentialSearch {
new() -> Self153 fn new() -> Self {
154 Self {
155 growing: true,
156 last: 0,
157 next: 1,
158 }
159 }
next(&self) -> u32160 fn next(&self) -> u32 {
161 self.next
162 }
record(&mut self, success: bool)163 fn record(&mut self, success: bool) {
164 if !success {
165 self.growing = false
166 }
167 let diff = if self.growing {
168 (self.next - self.last) * 2
169 } else {
170 (self.next - self.last + 1) / 2
171 };
172 if success {
173 self.last = self.next;
174 self.next = self.next + diff;
175 } else {
176 self.next = self.next - diff;
177 }
178 }
done(&self) -> bool179 fn done(&self) -> bool {
180 self.last == self.next
181 }
182 }
183
184 /// Build a pool-allocated engine with `num_memories` slots.
build_engine(args: &Args, num_memories: u32, enable_mpk: Enabled) -> Result<usize>185 fn build_engine(args: &Args, num_memories: u32, enable_mpk: Enabled) -> Result<usize> {
186 // Configure the memory pool.
187 let mut pool = PoolingAllocationConfig::default();
188 let max_memory_size =
189 usize::try_from(args.memory_size).expect("memory size should fit in `usize`");
190 pool.max_memory_size(max_memory_size)
191 .total_memories(num_memories)
192 .memory_protection_keys(enable_mpk);
193
194 // Configure the engine itself.
195 let mut config = Config::new();
196 if let Some(memory_reservation) = args.memory_reservation {
197 config.memory_reservation(memory_reservation);
198 }
199 if let Some(memory_guard_size) = args.memory_guard_size {
200 config.memory_guard_size(memory_guard_size);
201 }
202 config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool));
203
204 // Measure memory use before and after the engine is built.
205 let mapped_bytes_before = num_bytes_mapped()?;
206 let engine = Engine::new(&config)?;
207 let mapped_bytes_after = num_bytes_mapped()?;
208
209 // Ensure we actually use the engine somehow.
210 engine.increment_epoch();
211
212 let mapped_bytes = mapped_bytes_after - mapped_bytes_before;
213 info!("{num_memories}-slot pool ({enable_mpk:?}): {mapped_bytes} bytes mapped");
214 Ok(mapped_bytes)
215 }
216
217 /// Add up the sizes of all the mapped virtual memory regions for the current
218 /// Linux process.
219 ///
220 /// This manually parses `/proc/self/maps` to avoid a rather-large `proc-maps`
221 /// dependency. We do expect this example to be Linux-specific anyways. For
222 /// reference, lines of that file look like:
223 ///
224 /// ```text
225 /// 5652d4418000-5652d441a000 r--p 00000000 00:23 84629427 /usr/bin/...
226 /// ```
227 ///
228 /// We parse the start and end addresses: <start>-<end> [ignore the rest].
229 #[cfg(target_os = "linux")]
num_bytes_mapped() -> Result<usize>230 fn num_bytes_mapped() -> Result<usize> {
231 use std::fs::File;
232 use std::io::{BufRead, BufReader};
233
234 let file = File::open("/proc/self/maps")?;
235 let reader = BufReader::new(file);
236 let mut total = 0;
237 for line in reader.lines() {
238 let line = line?;
239 let range = line
240 .split_whitespace()
241 .next()
242 .ok_or(format_err!("parse failure: expected whitespace"))?;
243 let mut addresses = range.split("-");
244 let start = addresses.next().ok_or(format_err!(
245 "parse failure: expected dash-separated address"
246 ))?;
247 let start = usize::from_str_radix(start, 16)?;
248 let end = addresses.next().ok_or(format_err!(
249 "parse failure: expected dash-separated address"
250 ))?;
251 let end = usize::from_str_radix(end, 16)?;
252
253 total += end - start;
254 }
255 Ok(total)
256 }
257
258 #[cfg(not(target_os = "linux"))]
num_bytes_mapped() -> Result<usize>259 fn num_bytes_mapped() -> Result<usize> {
260 wasmtime::bail!("this example can only read virtual memory maps on Linux")
261 }
262