fix(benches): browser bench fixes (#821)

* fix(benches): make browser benches work again

* Update crates/benches/binary/README.md

Co-authored-by: Hendrik Eeckhaut <[email protected]>

* Update crates/benches/browser/wasm/Cargo.toml

Co-authored-by: Hendrik Eeckhaut <[email protected]>

* add --release flag

---------

Co-authored-by: Hendrik Eeckhaut <[email protected]>

Author: themighty1

crates/benches/binary/README.md

@@ -38,7 +38,7 @@ With a Chrome browser installed on your system, make sure you're in the `crates/
 directory, build the wasm module, build the binaries, and then run the script:
 ```sh
 cd browser/wasm
-rustup run nightly wasm-pack build --release --target web
+wasm-pack build --release --target web
 cd ../../binary
 cargo build --release --features browser-bench
 sudo ./bench.sh

crates/benches/binary/bench.toml

@@ -38,8 +38,8 @@ upload-delay = 25
 download = 250
 download-delay = 25
 upload-size = 1024
-# Setting download-size higher than 45000 will cause a `Maximum call stack size exceeded`
-# error in the browser.
-download-size = [1024, 4096, 16384, 45000]
+# It was observed that setting download-size > 30K causes browser errors that need to
+# be investigated.
+download-size = [1024, 4096, 16384]
 defer-decryption = true
 memory-profile = true

crates/benches/binary/benches.Dockerfile

@@ -5,14 +5,15 @@ COPY . .
 ARG BENCH_TYPE=native
 
 RUN \
+  rustup update; \
   if [ "$BENCH_TYPE" = "browser" ]; then \
     # ring's build script needs clang.
     apt update && apt install -y clang; \
     rustup install nightly; \
     rustup component add rust-src --toolchain nightly; \
     cargo install wasm-pack; \
     cd crates/benches/browser/wasm; \
-    rustup run nightly wasm-pack build --release --target web; \
+    wasm-pack build --release --target web; \
     cd ../../binary; \
     cargo build --release --features browser-bench; \
   else \

crates/benches/browser/native/src/lib.rs

@@ -4,14 +4,7 @@
 //! components. The native component is responsible for starting the browser,
 //! loading the wasm component and driving it.
 
-use std::{env, net::IpAddr};
-
-use serio::{stream::IoStreamExt, SinkExt as _};
-use tlsn_benches_browser_core::{
-    msg::{Config, Runtime},
-    FramedIo,
-};
-use tlsn_benches_library::{AsyncIo, ProverKind, ProverTrait};
+use std::{env, net::IpAddr, time::Duration};
 
 use anyhow::{anyhow, Context, Result};
 use async_trait::async_trait;
@@ -24,10 +17,17 @@ use chromiumoxide::{
 };
 use futures::{Future, FutureExt, StreamExt};
 use rust_embed::RustEmbed;
+use serio::{stream::IoStreamExt, SinkExt as _};
 use tokio::{io, io::AsyncWriteExt, net::TcpListener, task::JoinHandle};
 use tracing::{debug, error, info};
 use warp::Filter;
 
+use tlsn_benches_browser_core::{
+    msg::{Config, Runtime},
+    FramedIo,
+};
+use tlsn_benches_library::{AsyncIo, ProverKind, ProverTrait};
+
 /// The IP on which the wasm component is served.
 pub static DEFAULT_WASM_IP: &str = "127.0.0.1";
 /// The IP of the websocket relay.
@@ -100,6 +100,12 @@ impl ProverTrait for BrowserProver {
 
         relays.push(spawn_websocket_relay(ws_ip, ws_port).await?);
 
+        // Create a framed connection to the wasm component.
+        let (wasm_left, wasm_right) = tokio::io::duplex(1 << 16);
+
+        relays.push(spawn_port_relay(wasm_to_native_port, Box::new(wasm_right)).await?);
+        let mut wasm_io = FramedIo::new(Box::new(wasm_left));
+
         let http_server = spawn_http_server(wasm_ip, wasm_port)?;
 
         // Relay data from the wasm component to the server.
@@ -108,12 +114,6 @@ impl ProverTrait for BrowserProver {
         // Relay data from the wasm component to the verifier.
         relays.push(spawn_port_relay(wasm_to_verifier_port, verifier_io).await?);
 
-        // Create a framed connection to the wasm component.
-        let (wasm_left, wasm_right) = tokio::io::duplex(1 << 16);
-
-        relays.push(spawn_port_relay(wasm_to_native_port, Box::new(wasm_right)).await?);
-        let mut wasm_io = FramedIo::new(Box::new(wasm_left));
-
         info!("spawning browser");
 
         // Note that the browser must be spawned only when the WebSocket relay is
@@ -129,6 +129,10 @@ impl ProverTrait for BrowserProver {
         )
         .await?;
 
+        // Without this sleep, it was observed that `wasm_io.send(Config)`
+        // msg does not reach the browser component.
+        tokio::time::sleep(Duration::from_secs(2)).await;
+
         info!("sending config to the browser component");
 
         wasm_io
@@ -267,14 +271,15 @@ async fn spawn_browser(
     tokio::spawn(register_listeners(&page).await?);
 
     page.wait_for_navigation().await?;
+
     // Note that `format!` needs double {{ }} in order to escape them.
     let _ = page
         .evaluate_function(&format!(
             r#"
                 async function() {{
-                    await window.worker.init();
+                    await window.benchWorker.init();
                     // Do not `await` run() or else it will block the browser. 
-                    window.worker.run("{}", {}, {}, {}, {});
+                    window.benchWorker.run("{}", {}, {}, {}, {});
                 }}
             "#,
             ws_ip, ws_port, wasm_to_server_port, wasm_to_verifier_port, wasm_to_native_port

crates/benches/browser/wasm/.cargo/config.toml

@@ -1,11 +1,16 @@
 [build]
 target = "wasm32-unknown-unknown"
 
+[unstable]
+build-std = ["panic_abort", "std"]
+
 [target.wasm32-unknown-unknown]
 rustflags = [
     "-C",
     "target-feature=+atomics,+bulk-memory,+mutable-globals",
+    "-C",
+    # 4GB
+    "link-arg=--max-memory=4294967296",
+    "--cfg",
+    'getrandom_backend="wasm_js"',
 ]
-
-[unstable]
-build-std = ["panic_abort", "std"]

crates/benches/browser/wasm/Cargo.toml

@@ -18,9 +18,11 @@ tlsn-wasm = { path = "../../../wasm" }
 serio = { workspace = true }
 
 anyhow = { workspace = true }
+rayon = { workspace = true }
 tracing = { workspace = true }
-wasm-bindgen = { version = "0.2.87" }
-wasm-bindgen-futures = { version = "0.4.37" }
+wasm-bindgen = { version = "0.2" }
+wasm-bindgen-futures = { version = "0.4" }
+web-spawn = { workspace = true, features = ["no-bundler"] }
 web-time = { workspace = true }
 # Use the patched ws_stream_wasm to fix the issue https://github.com/najamelan/ws_stream_wasm/issues/12#issuecomment-1711902958
 ws_stream_wasm = { version = "0.7.4", git = "https://github.com/tlsnotary/ws_stream_wasm", rev = "2ed12aad9f0236e5321f577672f309920b2aef51", features = [

crates/benches/browser/wasm/pkg/index.js

@@ -1,7 +1,5 @@
 import * as Comlink from "./comlink.mjs";
 
-async function init() {
-    const worker = Comlink.wrap(new Worker("worker.js", { type: "module" }));
-    window.worker = worker;
-}
-init();
+const benchWorker = Comlink.wrap(new Worker("worker.js", { type: "module" }));
+
+window.benchWorker = benchWorker;

crates/benches/browser/wasm/pkg/worker.js

@@ -1,18 +1,14 @@
 import * as Comlink from "./comlink.mjs";
 
-import init, { wasm_main, initialize } from './tlsn_benches_browser_wasm.js';
+import init_wasm, * as wasm from './tlsn_benches_browser_wasm.js';
 
-class Worker {
+class BenchWorker {
   async init() {
     try {
-      await init();
-      // Tracing may interfere with the benchmark results. We should enable it only for debugging.
-      // init_logging({
-      //   level: 'Debug',
-      //   crate_filters: undefined,
-      //   span_events: undefined,
-      // });
-      await initialize({ thread_count: navigator.hardwareConcurrency });
+      await init_wasm();
+      // Using Error level since excessive logging may interfere with the
+      // benchmark results. 
+      await wasm.initialize_bench({ level: "Error" }, navigator.hardwareConcurrency);
     } catch (e) {
       console.error(e);
       throw e;
@@ -27,7 +23,7 @@ class Worker {
     wasm_to_native_port
   ) {
     try {
-      await wasm_main(
+      await wasm.wasm_main(
         ws_ip,
         ws_port,
         wasm_to_server_port,
@@ -40,6 +36,6 @@ class Worker {
   }
 }
 
-const worker = new Worker();
+const worker = new BenchWorker();
 
 Comlink.expose(worker);

crates/benches/browser/wasm/rust-toolchain.toml

@@ -1,2 +1,4 @@
 [toolchain]
-channel = "nightly"
+channel = "nightly"
+components = ["rust-src"]
+targets = ["wasm32-unknown-unknown"]

crates/benches/browser/wasm/src/lib.rs

@@ -5,18 +5,19 @@
 //! Conceptually the browser prover consists of the native and the wasm
 //! components.
 
+use anyhow::Result;
 use serio::{stream::IoStreamExt, SinkExt as _};
+use tracing::info;
+use wasm_bindgen::prelude::*;
+use web_time::Instant;
+use ws_stream_wasm::WsMeta;
+
 use tlsn_benches_browser_core::{
     msg::{Config, Runtime},
     FramedIo,
 };
 use tlsn_benches_library::run_prover;
-
-use anyhow::Result;
-use tracing::info;
-use wasm_bindgen::prelude::*;
-use web_time::Instant;
-use ws_stream_wasm::WsMeta;
+use tlsn_wasm::LoggingConfig;
 
 #[wasm_bindgen]
 pub async fn wasm_main(
@@ -85,6 +86,9 @@ pub async fn main(
     let cfg: Config = native_io.expect_next().await?;
 
     let start_time = Instant::now();
+
+    info!("running the prover");
+
     run_prover(
         cfg.upload_size,
         cfg.download_size,
@@ -100,3 +104,12 @@ pub async fn main(
 
     Ok(())
 }
+
+/// Initializes the module.
+#[wasm_bindgen]
+pub async fn initialize_bench(
+    logging_config: Option<LoggingConfig>,
+    thread_count: usize,
+) -> Result<(), JsValue> {
+    tlsn_wasm::initialize(logging_config, thread_count).await
+}