diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 3ea1112..ebf66bf 100644 --- a/.github/workflows/main.yml +++ b/.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 }} diff --git a/Cargo.lock b/Cargo.lock index 3784148..df93e9a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,12 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" +[[package]] +name = "allocator-api2" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c880a97d28a3681c0267bd29cff89621202715b065127cd445fa0f0fe0aa2880" + [[package]] name = "ambient-authority" version = "0.0.2" @@ -156,7 +162,7 @@ version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" dependencies = [ - "allocator-api2", + "allocator-api2 0.2.21", ] [[package]] @@ -560,6 +566,7 @@ name = "debug" version = "0.1.0" dependencies = [ "getrandom 0.3.4", + "talc", "wit-bindgen 0.52.0", ] @@ -1221,6 +1228,15 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + [[package]] name = "log" version = "0.4.29" @@ -1552,6 +1568,7 @@ name = "recorder" version = "0.1.0" dependencies = [ "serde_json", + "talc", "trace", "wit-bindgen 0.52.0", ] @@ -1573,7 +1590,7 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08effbc1fa53aaebff69521a5c05640523fab037b34a4a2c109506bc938246fa" dependencies = [ - "allocator-api2", + "allocator-api2 0.2.21", "bumpalo", "hashbrown 0.15.5", "log", @@ -1647,6 +1664,12 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "semver" version = "1.0.27" @@ -1830,6 +1853,16 @@ dependencies = [ "winx", ] +[[package]] +name = "talc" +version = "5.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3081e3f3ad542dc4cd0ae720f752598f8880fd364dc2d031132e3509f673aa23" +dependencies = [ + "allocator-api2 0.4.0", + "lock_api", +] + [[package]] name = "target-lexicon" version = "0.13.5" diff --git a/Makefile b/Makefile index 66d4b1d..9457206 100644 --- a/Makefile +++ b/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 diff --git a/assets/debug.wasm b/assets/debug.wasm index d22efa2..33a5ad9 100644 Binary files a/assets/debug.wasm and b/assets/debug.wasm differ diff --git a/assets/proj_cargo.toml b/assets/proj_cargo.toml index 58aa0dd..6a52f39 100644 --- a/assets/proj_cargo.toml +++ b/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 diff --git a/assets/recorder.wasm b/assets/recorder.wasm index 3a7c079..dee6ab1 100644 Binary files a/assets/recorder.wasm and b/assets/recorder.wasm differ diff --git a/assets/workspace_cargo.toml b/assets/workspace_cargo.toml index 1eedad2..191c989 100644 --- a/assets/workspace_cargo.toml +++ b/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 diff --git a/components/debug/Cargo.toml b/components/debug/Cargo.toml index 8666812..6609ae2 100644 --- a/components/debug/Cargo.toml +++ b/components/debug/Cargo.toml @@ -8,4 +8,5 @@ crate-type = ["cdylib"] [dependencies] wit-bindgen.workspace = true -getrandom = "0.3" \ No newline at end of file +getrandom = "0.3" +talc = "5.0.3" diff --git a/components/debug/src/lib.rs b/components/debug/src/lib.rs index 2c8467c..b9e9bbe 100644 --- a/components/debug/src/lib.rs +++ b/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; 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", diff --git a/components/recorder/Cargo.toml b/components/recorder/Cargo.toml index e64a57a..ddd92bc 100644 --- a/components/recorder/Cargo.toml +++ b/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" diff --git a/components/recorder/src/lib.rs b/components/recorder/src/lib.rs index 8566a28..4b3939d 100644 --- a/components/recorder/src/lib.rs +++ b/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; 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", diff --git a/src/codegen/mod.rs b/src/codegen/mod.rs index 9de3718..b445bfa 100644 --- a/src/codegen/mod.rs +++ b/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; 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(()) diff --git a/src/instrument.rs b/src/instrument.rs index 06fd281..ce97122 100644 --- a/src/instrument.rs +++ b/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") diff --git a/src/run.rs b/src/run.rs index b71dc7c..2db054f 100644 --- a/src/run.rs +++ b/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)?;