use page_size=1 for instrumented components

.github/workflows/main.yml

@@ -22,6 +22,8 @@ jobs:
           toolchain: 1.93.1   # Pin to 1.93 to use WASI 0.2.6, before viceroy catches up to 0.2.9
       - uses: bytecodealliance/actions/wasmtime/setup@v1
       - uses: bytecodealliance/actions/wasm-tools/setup@v1
+        with:
+          version: 1.244.0
       - uses: bytecodealliance/actions/wit-bindgen/setup@v1
         with:
           version: ${{ env.WIT_BINDGEN_VERSION }}

Makefile

@@ -4,8 +4,8 @@ all: build-components build-cli
 build-cli:
 	cargo build --all-features --release
 build-components:
-	cargo build -p debug --target wasm32-wasip2 --release
-	cargo build -p recorder --target wasm32-wasip2 --release
+	RUSTFLAGS="-C link-arg=--page-size=1" cargo build -p debug --target wasm32-wasip2 --release
+	RUSTFLAGS="-C link-arg=--page-size=1" cargo build -p recorder --target wasm32-wasip2 --release
 	cp target/wasm32-wasip2/release/debug.wasm assets/debug.wasm
 	cp target/wasm32-wasip2/release/recorder.wasm assets/recorder.wasm
 
@@ -25,7 +25,7 @@ test-record:
 	$(MAKE) run-record WASM=tests/rust.wasm
 	# test the same trace with a different wasm replay
 	target/release/proxy-component instrument -m replay tests/rust.debug.wasm
-	wasmtime --invoke 'start()' composed.wasm < trace.out
+	wasmtime -W custom-page-sizes=y --invoke 'start()' composed.wasm < trace.out
 	# build-only test
 	target/release/proxy-component instrument -m record tests/calculator.wasm
 	target/release/proxy-component instrument -m replay tests/calculator.wasm
@@ -40,13 +40,13 @@ test-dialog:
 
 run-fuzz:
 	target/release/proxy-component instrument -m fuzz $(WASM)
-	wasmtime --invoke 'start()' composed.wasm
+	wasmtime -W custom-page-sizes=y --invoke 'start()' composed.wasm
 
 run-record:
 	target/release/proxy-component instrument -m record $(WASM)
 	$(MAKE) run-viceroy URL=localhost:7676
 	target/release/proxy-component instrument -m replay $(WASM)
-	wasmtime --invoke 'start()' composed.wasm < trace.out
+	wasmtime -W custom-page-sizes=y --invoke 'start()' composed.wasm < trace.out
 	# test host replay
 	target/release/proxy-component instrument -m replay --use-host-recorder $(WASM)
 	target/release/proxy-component run composed.wasm --invoke 'start()' --trace trace.out

assets/proj_cargo.toml

@@ -11,3 +11,4 @@ path = "lib.rs"
 wit-bindgen.workspace = true
 wasm-wave.workspace = true
 arbitrary.workspace = true
+talc.workspace = true

assets/workspace_cargo.toml

@@ -7,6 +7,7 @@ wit-bindgen = { version = "0.55.0", default-features = false, features = ["bitfl
 wasm-wave = { git = "https://github.com/chenyan2002/wasm-tools.git", branch = "extend-wave", version = "0.239.0", default-features = false }
 #wasm-wave = { path = "/Users/chenyan/src/bytecodealliance/wasm-tools/crates/wasm-wave", default-features = false }
 arbitrary = "1.4.2"
+talc = "5.0.3"
 
 [profile.release]
 codegen-units = 1

components/debug/Cargo.toml

@@ -8,4 +8,5 @@ crate-type = ["cdylib"]
 
 [dependencies]
 wit-bindgen.workspace = true
-getrandom = "0.3"
+getrandom = "0.3"
+talc = "5.0.3"

components/debug/src/lib.rs

@@ -1,3 +1,12 @@
+#[cfg(all(not(target_feature = "atomics"), target_family = "wasm"))]
+#[global_allocator]
+static TALC: talc::wasm::WasmArenaTalc = {
+    use core::mem::MaybeUninit;
+    static mut MEMORY: [MaybeUninit<u8>; 0x80000] = [MaybeUninit::uninit(); 0x80000];
+    // SAFETY: the memory for MEMORY is never modified externally. It's the allocator's.
+    unsafe { talc::wasm::new_wasm_arena_allocator(&raw mut MEMORY) }
+};
+
 mod bindings {
     wit_bindgen::generate!({
         path: "../../assets/util.wit",

components/recorder/Cargo.toml

@@ -10,3 +10,4 @@ crate-type = ["cdylib"]
 trace = { path = "../../crates/trace" }
 serde_json.workspace = true
 wit-bindgen.workspace = true
+talc = "5.0.3"

components/recorder/src/lib.rs

@@ -1,3 +1,12 @@
+#[cfg(all(not(target_feature = "atomics"), target_family = "wasm"))]
+#[global_allocator]
+static TALC: talc::wasm::WasmArenaTalc = {
+    use core::mem::MaybeUninit;
+    static mut MEMORY: [MaybeUninit<u8>; 0x80000] = [MaybeUninit::uninit(); 0x80000];
+    // SAFETY: the memory for MEMORY is never modified externally. It's the allocator's.
+    unsafe { talc::wasm::new_wasm_arena_allocator(&raw mut MEMORY) }
+};
+
 mod bindings {
     wit_bindgen::generate!({
         path: "../../assets/recorder.wit",

src/codegen/mod.rs

@@ -22,6 +22,8 @@ pub struct GenerateArgs {
     /// The path to the output file.
     #[arg(short, long, default_value("lib.rs"))]
     pub output_file: PathBuf,
+    #[arg(long)]
+    pub use_custom_allocator: bool,
 }
 #[derive(clap::ValueEnum, clap::Parser, Clone)]
 pub enum GenerateMode {
@@ -84,7 +86,19 @@ impl GenerateArgs {
         let traits = trait_generator.generate();
         drop(trait_generator);
         state.output.extend(traits);
-        let file = state.into_output_file();
+        let mut file = state.into_output_file();
+        if self.use_custom_allocator {
+            file.items.push(parse_quote! {
+                #[cfg(all(not(target_feature = "atomics"), target_family = "wasm"))]
+                #[global_allocator]
+                static TALC: talc::wasm::WasmArenaTalc = {
+                    use core::mem::MaybeUninit;
+                    static mut MEMORY: [MaybeUninit<u8>; 0x80000] = [MaybeUninit::uninit(); 0x80000];
+                    // SAFETY: the memory for MEMORY is never modified externally. It's the allocator's.
+                    unsafe { talc::wasm::new_wasm_arena_allocator(&raw mut MEMORY) }
+                };
+            });
+        }
         let output = prettyplease::unparse(&file);
         std::fs::write(&self.output_file, output)?;
         Ok(())

src/instrument.rs

@@ -17,8 +17,12 @@ pub struct InstrumentArgs {
     /// Whether to use the host recorder implementation or link the recorder component
     #[arg(long)]
     pub use_host_recorder: bool,
+    /// Whether to use a one-byte page size for the wasm component
+    #[arg(long, default_value_t = true)]
+    pub one_byte_page_size: bool,
 }
 
+// TODO: Add one-byte and regular page size wasm binaries
 const DEBUG_WASM: &[u8] = include_bytes!("../assets/debug.wasm");
 const RECORDER_WASM: &[u8] = include_bytes!("../assets/recorder.wasm");
 
@@ -63,10 +67,27 @@ pub fn run(args: InstrumentArgs) -> Result<()> {
     }
 
     // 5. Generate Rust binding for both import and export interface
-    bindgen(&tmp_dir, &wit_dir, &args.mode, "imports", "record_imports")?;
-    bindgen(&tmp_dir, &wit_dir, &args.mode, "exports", "record_exports")?;
+    bindgen(
+        &tmp_dir,
+        &wit_dir,
+        &args.mode,
+        "imports",
+        "record_imports",
+        args.one_byte_page_size,
+    )?;
+    bindgen(
+        &tmp_dir,
+        &wit_dir,
+        &args.mode,
+        "exports",
+        "record_exports",
+        args.one_byte_page_size,
+    )?;
     // 6. cargo build
     let mut cmd = Command::new("cargo");
+    if args.one_byte_page_size {
+        cmd.env("RUSTFLAGS", "-C link-arg=--page-size=1");
+    }
     cmd.arg("build")
         .arg("--target=wasm32-unknown-unknown")
         .current_dir(tmp_dir.as_path());
@@ -128,6 +149,7 @@ fn bindgen(
     mode: &Mode,
     world_name: &str,
     dest_name: &str,
+    use_custom_allocator: bool,
 ) -> Result<()> {
     let out_dir = tmp_dir.join(dest_name);
     let status = Command::new("wit-bindgen")
@@ -154,6 +176,7 @@ fn bindgen(
         bindings: binding_file.clone(),
         output_file: out_dir.join("lib.rs"),
         mode: codegen_mode,
+        use_custom_allocator,
     };
     codegen_opt.generate()?;
     let status = Command::new("mv")

src/run.rs

@@ -74,6 +74,7 @@ pub fn run(args: RunArgs) -> anyhow::Result<()> {
     let mut config = Config::new();
     config
         .consume_fuel(true)
+        .wasm_custom_page_sizes(true)
         //.debug_info(true)
         .wasm_backtrace_details(WasmBacktraceDetails::Enable);
     let engine = Engine::new(&config)?;