From 62ca987dde6d3028d64f0ac2384c78548a835b09 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Wed, 15 Apr 2026 02:42:53 -0300 Subject: [PATCH 01/17] feat: add integration benchmarks with mock server Allocation-counting benchmarks that run against Bartender (mock server) to track real protocol-level allocation regressions in CI. Scenarios: connect-to-ready, send message, send+receive round-trip, reconnect. Outputs customSmallerIsBetter JSON for github-action-benchmark. Optional `dhat-heap` feature for full call-tree profiling via DHAT viewer. Also switches e2e tests from SqliteStore to InMemoryBackend and disables the sqlite-storage default feature on both test crates for faster builds. --- .github/workflows/bench-integration.yml | 119 ++++++ .gitignore | 1 + Cargo.lock | 17 +- Cargo.toml | 7 + tests/bench-integration/Cargo.toml | 35 ++ tests/bench-integration/src/counting_alloc.rs | 49 +++ tests/bench-integration/src/main.rs | 353 ++++++++++++++++++ tests/e2e/Cargo.toml | 11 +- tests/e2e/src/lib.rs | 17 +- 9 files changed, 592 insertions(+), 17 deletions(-) create mode 100644 .github/workflows/bench-integration.yml create mode 100644 tests/bench-integration/Cargo.toml create mode 100644 tests/bench-integration/src/counting_alloc.rs create mode 100644 tests/bench-integration/src/main.rs diff --git a/.github/workflows/bench-integration.yml b/.github/workflows/bench-integration.yml new file mode 100644 index 000000000..d64b8df34 --- /dev/null +++ b/.github/workflows/bench-integration.yml @@ -0,0 +1,119 @@ +name: Integration Benchmark + +on: + pull_request: + branches: [main] + push: + branches: [main] + +permissions: + contents: write + pull-requests: write + deployments: write + +concurrency: + group: bench-integration-${{ github.head_ref || github.run_id }} + cancel-in-progress: true + +env: + CARGO_TERM_COLOR: always + PROTOC_VERSION: "3.25.3" + SCCACHE_GHA_ENABLED: "true" + RUSTC_WRAPPER: "sccache" + +jobs: + bench-integration: + name: Integration Benchmark + runs-on: ubuntu-latest + services: + mock-server: + image: ghcr.io/whiskeysockets-devtools/bartender:latest + credentials: + username: ${{ github.actor }} + password: ${{ secrets.BARTENDER_GHCR_TOKEN }} + ports: + - 8080:8080 + env: + CHATSTATE_TTL_SECS: "3" + options: --log-driver none + steps: + - uses: actions/checkout@v6 + + - uses: dtolnay/rust-toolchain@master + with: + toolchain: nightly-2026-04-05 + + - name: Install protoc + uses: taiki-e/install-action@v2 + with: + tool: protoc@${{ env.PROTOC_VERSION }} + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Cache Rust registry + uses: Swatinem/rust-cache@v2 + with: + prefix-key: ${{ runner.os }}-cargo-bench-int + cache-targets: "false" + + - name: Wait for mock server + run: | + for i in $(seq 1 30); do + if curl -sk https://localhost:8080/ > /dev/null 2>&1; then + echo "Mock server is ready" + exit 0 + fi + sleep 1 + done + echo "Mock server failed to become ready" + exit 1 + + - name: Run integration benchmarks + env: + MOCK_SERVER_URL: "wss://127.0.0.1:8080/ws/chat" + RUST_LOG: warn + run: | + cargo run -p bench-integration --release \ + > bench_integration_results.json \ + 2> bench_integration.log + cat bench_integration.log >&2 + echo "Results:" + cat bench_integration_results.json + + - name: Upload benchmark artifacts + uses: actions/upload-artifact@v4 + with: + name: bench-integration-results + path: | + bench_integration_results.json + bench_integration.log + retention-days: 30 + + # Push to main: store baseline + - name: Store baseline (push to main) + if: github.event_name == 'push' + uses: benchmark-action/github-action-benchmark@v1 + with: + name: "whatsapp-rust integration benchmarks" + tool: "customSmallerIsBetter" + output-file-path: bench_integration_results.json + github-token: ${{ secrets.GITHUB_TOKEN }} + auto-push: true + benchmark-data-dir-path: dev/bench-integration + max-items-in-chart: 100 + summary-always: true + + # PR: compare against baseline + - name: Compare against baseline (pull request) + if: github.event_name == 'pull_request' + uses: benchmark-action/github-action-benchmark@v1 + with: + name: "whatsapp-rust integration benchmarks" + tool: "customSmallerIsBetter" + output-file-path: bench_integration_results.json + github-token: ${{ secrets.GITHUB_TOKEN }} + auto-push: false + save-data-file: false + benchmark-data-dir-path: dev/bench-integration + summary-always: true diff --git a/.gitignore b/.gitignore index 539e2d082..a157f7df1 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ docs .claude __pycache__ .codex +dhat-heap.json diff --git a/Cargo.lock b/Cargo.lock index 498b50e9c..e442964a5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -144,6 +144,22 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "bench-integration" +version = "0.0.0" +dependencies = [ + "anyhow", + "dhat", + "e2e-tests", + "env_logger", + "log", + "serde", + "serde_json", + "tokio", + "wacore", + "whatsapp-rust", +] + [[package]] name = "bincode" version = "1.3.3" @@ -608,7 +624,6 @@ dependencies = [ "wacore", "wacore-binary", "whatsapp-rust", - "whatsapp-rust-sqlite-storage", "whatsapp-rust-tokio-transport", "whatsapp-rust-ureq-http-client", ] diff --git a/Cargo.toml b/Cargo.toml index edb16d2c7..59507d659 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ members = [ ".", "http_clients/ureq-client", "storages/sqlite-storage", + "tests/bench-integration", "tests/e2e", "transports/tokio-transport", "wacore", @@ -161,3 +162,9 @@ inherits = "release" lto = "thin" debug = 1 strip = false + +# Profiling: optimized with debug symbols for DHAT / heaptrack / perf +[profile.profiling] +inherits = "release" +debug = 1 +strip = false diff --git a/tests/bench-integration/Cargo.toml b/tests/bench-integration/Cargo.toml new file mode 100644 index 000000000..30dae0baa --- /dev/null +++ b/tests/bench-integration/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "bench-integration" +version = "0.0.0" +edition = "2024" +publish = false + +[package.metadata.cargo-shear] +ignored = ["dhat"] + +[[bin]] +name = "bench-integration" +path = "src/main.rs" + +[features] +dhat-heap = ["dep:dhat"] + +[dependencies] +anyhow = { workspace = true } +dhat = { version = "0.3", optional = true } +e2e-tests = { path = "../e2e" } +env_logger = { workspace = true } +log = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true, features = ["std"] } +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "sync", "time"] } +wacore = { path = "../../wacore" } +whatsapp-rust = { path = "../..", default-features = false, features = [ + "danger-skip-tls-verify", + "debug-diagnostics", + "moka-cache", + "simd", + "tokio-runtime", + "tokio-native", + "signal", +] } diff --git a/tests/bench-integration/src/counting_alloc.rs b/tests/bench-integration/src/counting_alloc.rs new file mode 100644 index 000000000..130049772 --- /dev/null +++ b/tests/bench-integration/src/counting_alloc.rs @@ -0,0 +1,49 @@ +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::atomic::{AtomicU64, Ordering}; + +static ALLOC_COUNT: AtomicU64 = AtomicU64::new(0); +static ALLOC_BYTES: AtomicU64 = AtomicU64::new(0); + +pub struct CountingAlloc; + +unsafe impl GlobalAlloc for CountingAlloc { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + ALLOC_COUNT.fetch_add(1, Ordering::Relaxed); + ALLOC_BYTES.fetch_add(layout.size() as u64, Ordering::Relaxed); + unsafe { System.alloc(layout) } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct AllocSnapshot { + alloc_count: u64, + alloc_bytes: u64, +} + +impl AllocSnapshot { + pub fn now() -> Self { + Self { + alloc_count: ALLOC_COUNT.load(Ordering::Relaxed), + alloc_bytes: ALLOC_BYTES.load(Ordering::Relaxed), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct AllocDelta { + pub alloc_count: u64, + pub alloc_bytes: u64, +} + +impl AllocDelta { + pub fn between(before: AllocSnapshot, after: AllocSnapshot) -> Self { + Self { + alloc_count: after.alloc_count.saturating_sub(before.alloc_count), + alloc_bytes: after.alloc_bytes.saturating_sub(before.alloc_bytes), + } + } +} diff --git a/tests/bench-integration/src/main.rs b/tests/bench-integration/src/main.rs new file mode 100644 index 000000000..82c6c30f2 --- /dev/null +++ b/tests/bench-integration/src/main.rs @@ -0,0 +1,353 @@ +#[cfg(not(feature = "dhat-heap"))] +mod counting_alloc; + +#[cfg(feature = "dhat-heap")] +#[global_allocator] +static GLOBAL: dhat::Alloc = dhat::Alloc; + +#[cfg(not(feature = "dhat-heap"))] +#[global_allocator] +static GLOBAL: counting_alloc::CountingAlloc = counting_alloc::CountingAlloc; + +use std::time::Instant; + +use e2e_tests::{TestClient, text_msg}; +use serde::Serialize; + +// --------------------------------------------------------------------------- +// Measurement helpers +// --------------------------------------------------------------------------- + +/// Snapshot-based measurement: runs `f`, returns wall time + alloc delta. +#[cfg(not(feature = "dhat-heap"))] +async fn measure( + mut f: F, +) -> anyhow::Result<(T, std::time::Duration, counting_alloc::AllocDelta)> +where + F: AsyncFnMut() -> anyhow::Result, +{ + let before = counting_alloc::AllocSnapshot::now(); + let t0 = Instant::now(); + let result = f().await?; + let elapsed = t0.elapsed(); + let delta = counting_alloc::AllocDelta::between(before, counting_alloc::AllocSnapshot::now()); + Ok((result, elapsed, delta)) +} + +/// DHAT mode: just runs `f` and returns wall time (DHAT captures everything). +#[cfg(feature = "dhat-heap")] +async fn measure(mut f: F) -> anyhow::Result<(T, std::time::Duration)> +where + F: AsyncFnMut() -> anyhow::Result, +{ + let t0 = Instant::now(); + let result = f().await?; + let elapsed = t0.elapsed(); + Ok((result, elapsed)) +} + +// --------------------------------------------------------------------------- +// Results collection +// --------------------------------------------------------------------------- + +#[derive(Serialize)] +struct BenchEntry { + name: String, + unit: String, + value: u64, +} + +struct BenchResults(Vec); + +impl BenchResults { + fn new() -> Self { + Self(Vec::new()) + } + + fn record(&mut self, scenario: &str, metric: &str, unit: &str, value: u64) { + self.0.push(BenchEntry { + name: format!("integration::{scenario}::{metric}"), + unit: unit.to_string(), + value, + }); + } + + fn record_wall(&mut self, scenario: &str, elapsed: std::time::Duration) { + self.record( + scenario, + "wall_ms", + "milliseconds", + elapsed.as_millis() as u64, + ); + } + + #[cfg(not(feature = "dhat-heap"))] + fn record_measured( + &mut self, + scenario: &str, + elapsed: std::time::Duration, + delta: &counting_alloc::AllocDelta, + ) { + self.record(scenario, "alloc_count", "allocations", delta.alloc_count); + self.record(scenario, "alloc_bytes", "bytes", delta.alloc_bytes); + self.record_wall(scenario, elapsed); + } + + #[cfg(not(feature = "dhat-heap"))] + fn record_measured_amortized( + &mut self, + scenario: &str, + n: u64, + elapsed: std::time::Duration, + delta: &counting_alloc::AllocDelta, + ) { + self.record( + scenario, + "alloc_count", + "allocations", + delta.alloc_count / n, + ); + self.record(scenario, "alloc_bytes", "bytes", delta.alloc_bytes / n); + self.record( + scenario, + "wall_ms", + "milliseconds", + elapsed.as_millis() as u64 / n, + ); + } +} + +// --------------------------------------------------------------------------- +// Scenarios +// --------------------------------------------------------------------------- + +/// Measure allocations from client creation through Connected (ready). +async fn bench_connect_to_ready(results: &mut BenchResults) -> anyhow::Result<()> { + let m = measure(async || TestClient::connect("bench_connect").await).await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (client, elapsed, delta) = m; + results.record_measured("connect_to_ready", elapsed, &delta); + client.disconnect().await; + } + #[cfg(feature = "dhat-heap")] + { + let (client, elapsed) = m; + results.record_wall("connect_to_ready", elapsed); + client.disconnect().await; + } + Ok(()) +} + +/// Measure allocations for sending a single DM (sender side only). +/// +/// Both clients are connected before measurement starts. +/// We track the `send_message` call which covers: protobuf encoding, +/// Signal encrypt, node marshal, and WebSocket write. +async fn bench_send_message(results: &mut BenchResults) -> anyhow::Result<()> { + let client_a = TestClient::connect("bench_send_a").await?; + let mut client_b = TestClient::connect("bench_send_b").await?; + let jid_b = client_b.jid().await; + + // Warm up: establish Signal session with a throwaway message + client_a + .client + .send_message(jid_b.clone(), text_msg("warmup-send")) + .await?; + client_b.wait_for_text("warmup-send", 30).await?; + + // -- Single send -- + let m = measure(async || { + client_a + .client + .send_message(jid_b.clone(), text_msg("bench-send-single")) + .await + }) + .await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (_, elapsed, delta) = m; + results.record_measured("send_message", elapsed, &delta); + } + #[cfg(feature = "dhat-heap")] + { + let (_, elapsed) = m; + results.record_wall("send_message", elapsed); + } + + client_b.wait_for_text("bench-send-single", 30).await?; + + // -- Amortized: send N messages -- + const N: u64 = 20; + let m = measure(async || { + for i in 0..N { + client_a + .client + .send_message(jid_b.clone(), text_msg(&format!("bench-send-{i}"))) + .await?; + } + Ok(()) + }) + .await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (_, elapsed, delta) = m; + results.record_measured_amortized("send_message_x20_amortized", N, elapsed, &delta); + } + #[cfg(feature = "dhat-heap")] + { + let (_, elapsed) = m; + results.record( + "send_message_x20_amortized", + "wall_ms", + "milliseconds", + elapsed.as_millis() as u64 / N, + ); + } + + for i in 0..N { + client_b + .wait_for_text(&format!("bench-send-{i}"), 30) + .await?; + } + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +/// Measure allocations for a full send+receive round-trip. +/// +/// Both clients are connected and session is warm before measurement. +async fn bench_receive_message(results: &mut BenchResults) -> anyhow::Result<()> { + let client_a = TestClient::connect("bench_recv_a").await?; + let mut client_b = TestClient::connect("bench_recv_b").await?; + let jid_b = client_b.jid().await; + + // Warm up Signal session + client_a + .client + .send_message(jid_b.clone(), text_msg("warmup-recv")) + .await?; + client_b.wait_for_text("warmup-recv", 30).await?; + + // -- Single send+receive -- + let m = measure(async || { + client_a + .client + .send_message(jid_b.clone(), text_msg("bench-recv-single")) + .await?; + client_b.wait_for_text("bench-recv-single", 30).await?; + Ok(()) + }) + .await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (_, elapsed, delta) = m; + results.record_measured("send_and_receive_message", elapsed, &delta); + } + #[cfg(feature = "dhat-heap")] + { + let (_, elapsed) = m; + results.record_wall("send_and_receive_message", elapsed); + } + + // -- Amortized N round-trips -- + const N: u64 = 20; + let m = measure(async || { + for i in 0..N { + let text = format!("bench-recv-{i}"); + client_a + .client + .send_message(jid_b.clone(), text_msg(&text)) + .await?; + client_b.wait_for_text(&text, 30).await?; + } + Ok(()) + }) + .await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (_, elapsed, delta) = m; + results.record_measured_amortized("send_and_receive_x20_amortized", N, elapsed, &delta); + } + #[cfg(feature = "dhat-heap")] + { + let (_, elapsed) = m; + results.record( + "send_and_receive_x20_amortized", + "wall_ms", + "milliseconds", + elapsed.as_millis() as u64 / N, + ); + } + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +/// Measure allocations for a reconnect cycle (disconnect -> reconnect -> ready). +async fn bench_reconnect(results: &mut BenchResults) -> anyhow::Result<()> { + let mut client = TestClient::connect("bench_reconn").await?; + + let m = measure(async || { + client.reconnect_and_wait().await?; + Ok(()) + }) + .await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (_, elapsed, delta) = m; + results.record_measured("reconnect", elapsed, &delta); + } + #[cfg(feature = "dhat-heap")] + { + let (_, elapsed) = m; + results.record_wall("reconnect", elapsed); + } + + client.disconnect().await; + Ok(()) +} + +// --------------------------------------------------------------------------- +// Main +// --------------------------------------------------------------------------- + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("warn")) + .try_init() + .ok(); + + #[cfg(feature = "dhat-heap")] + let _profiler = dhat::Profiler::new_heap(); + + let mut results = BenchResults::new(); + + eprintln!("--- bench: connect_to_ready ---"); + bench_connect_to_ready(&mut results).await?; + + eprintln!("--- bench: send_message ---"); + bench_send_message(&mut results).await?; + + eprintln!("--- bench: receive_message ---"); + bench_receive_message(&mut results).await?; + + eprintln!("--- bench: reconnect ---"); + bench_reconnect(&mut results).await?; + + // Output customSmallerIsBetter JSON to stdout + let json = serde_json::to_string_pretty(&results.0)?; + println!("{json}"); + + eprintln!("--- done: {} metrics collected ---", results.0.len()); + Ok(()) +} diff --git a/tests/e2e/Cargo.toml b/tests/e2e/Cargo.toml index f97e54366..b5d71c233 100644 --- a/tests/e2e/Cargo.toml +++ b/tests/e2e/Cargo.toml @@ -19,8 +19,15 @@ tokio = { workspace = true, features = ["macros", "rt-multi-thread", "sync", "ti uuid = { workspace = true, features = ["v4"] } wacore = { path = "../../wacore" } wacore-binary = { path = "../../wacore/binary" } -whatsapp-rust = { path = "../..", features = ["danger-skip-tls-verify", "debug-diagnostics"] } -whatsapp-rust-sqlite-storage = { path = "../../storages/sqlite-storage" } +whatsapp-rust = { path = "../..", default-features = false, features = [ + "danger-skip-tls-verify", + "debug-diagnostics", + "moka-cache", + "simd", + "tokio-runtime", + "tokio-native", + "signal", +] } whatsapp-rust-tokio-transport = { path = "../../transports/tokio-transport", features = [ "danger-skip-tls-verify", ] } diff --git a/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index c246c47f8..baf323450 100644 --- a/tests/e2e/src/lib.rs +++ b/tests/e2e/src/lib.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use wacore::store::InMemoryBackend; use wacore::store::traits::TcTokenEntry; use wacore::types::events::{ChannelEventHandler, Event}; use wacore_binary::node::Node; @@ -7,20 +8,9 @@ use whatsapp_rust::Jid; use whatsapp_rust::bot::Bot; use whatsapp_rust::store::traits::Backend; use whatsapp_rust::waproto::whatsapp as wa; -use whatsapp_rust_sqlite_storage::SqliteStore; use whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory; use whatsapp_rust_ureq_http_client::UreqHttpClient; -/// Creates a SqliteStore with a unique in-memory database for test isolation. -pub async fn create_test_store(prefix: &str) -> anyhow::Result { - let db = format!( - "file:{}_{}?mode=memory&cache=shared", - prefix, - uuid::Uuid::new_v4() - ); - Ok(SqliteStore::new(&db).await?) -} - /// Returns the mock server WebSocket URL from env, or the default. pub fn mock_server_url() -> String { std::env::var("MOCK_SERVER_URL").unwrap_or_else(|_| "wss://127.0.0.1:8080/ws/chat".to_string()) @@ -64,9 +54,8 @@ impl TestClient { Self::connect_inner(prefix, Some(push_name.to_string())).await } - async fn connect_inner(prefix: &str, push_name: Option) -> anyhow::Result { - let store = create_test_store(prefix).await?; - let backend = Arc::new(store) as Arc; + async fn connect_inner(_prefix: &str, push_name: Option) -> anyhow::Result { + let backend = Arc::new(InMemoryBackend::new()) as Arc; let transport_factory = TokioWebSocketTransportFactory::new().with_url(mock_server_url()); let (event_handler, event_rx) = ChannelEventHandler::new(); From 8efa7a8fca3f9b3b6a5aa50007cd45fde2df07fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Wed, 15 Apr 2026 02:56:04 -0300 Subject: [PATCH 02/17] perf: reduce transport event channel capacity from 10K to 1K DHAT profiling showed `concurrent_queue::Bounded::new` as the #1 allocation site at 6.2 MB per connection. The bounded ring buffer pre-allocates all 10,000 slots upfront. 1,024 slots is still generous headroom -- WA Web processes messages inline and a backlog that large means the client can't keep up anyway. Measured impact (bench-integration with mock server): - connect_to_ready: 5.19 MB -> 4.24 MB (-18.4%) - reconnect: 955 KB -> 522 KB (-45.3%) - DHAT total: 37.9 MB -> 32.1 MB (-15.3%) - send/receive paths: unchanged (as expected) --- transports/tokio-transport/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transports/tokio-transport/src/lib.rs b/transports/tokio-transport/src/lib.rs index cd8ebd462..fbc7215eb 100644 --- a/transports/tokio-transport/src/lib.rs +++ b/transports/tokio-transport/src/lib.rs @@ -15,7 +15,7 @@ use wacore::net::{Transport, TransportEvent, TransportFactory, WHATSAPP_WEB_WS_U pub use tokio_websockets::Connector; -const EVENT_CHANNEL_CAPACITY: usize = 10_000; +const EVENT_CHANNEL_CAPACITY: usize = 1_024; static CRYPTO_PROVIDER_INIT: Once = Once::new(); From f46417eec86f85c5baa0130060897cd3cb6cf46e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Wed, 15 Apr 2026 03:12:56 -0300 Subject: [PATCH 03/17] perf: eliminate intermediate allocations in prekey upload and send paths P2: build_upload_prekeys_request now accepts &[u8] slices instead of Vec, removing 812+ intermediate .to_vec() calls in build_iq(). The NodeBuilder still allocates internally, but the caller no longer creates throwaway Vecs that are immediately consumed. P3: encrypt_group_message, prepare_peer_stanza, and create_sender_key_distribution_message_for_group now accept &ProtocolAddress instead of computing it from &Jid internally. Callers that already have the address (or call multiple functions with the same JID) avoid redundant 64-byte String allocations. In prepare_group_stanza, the address is now built once and shared between SKDM creation and group encryption. --- src/features/signal.rs | 4 ++-- src/send.rs | 2 +- wacore/benches/send_receive_benchmark.rs | 3 ++- wacore/src/iq/prekeys.rs | 8 ++++---- wacore/src/prekeys.rs | 23 ++++++++++------------ wacore/src/send.rs | 25 ++++++++++++------------ 6 files changed, 31 insertions(+), 34 deletions(-) diff --git a/src/features/signal.rs b/src/features/signal.rs index 2ed11995f..9e4b1e706 100644 --- a/src/features/signal.rs +++ b/src/features/signal.rs @@ -151,7 +151,7 @@ impl<'a> Signal<'a> { wacore::send::create_sender_key_distribution_message_for_group( &mut adapter.sender_key_store, group_jid, - &own_jid, + &sender_addr, ) .await?, ) @@ -162,7 +162,7 @@ impl<'a> Signal<'a> { let ciphertext = wacore::send::encrypt_group_message( &mut adapter.sender_key_store, group_jid, - &own_jid, + &sender_addr, plaintext, &mut rng, ) diff --git a/src/send.rs b/src/send.rs index de7800f4c..b886d1124 100644 --- a/src/send.rs +++ b/src/send.rs @@ -883,7 +883,7 @@ impl Client { &mut store_adapter.session_store, &mut store_adapter.identity_store, to, - encryption_jid, + &signal_addr, message, request_id, ) diff --git a/wacore/benches/send_receive_benchmark.rs b/wacore/benches/send_receive_benchmark.rs index 2df58f1ff..3d44002b2 100644 --- a/wacore/benches/send_receive_benchmark.rs +++ b/wacore/benches/send_receive_benchmark.rs @@ -590,11 +590,12 @@ fn setup_group_recv() -> GrpRecvData { #[library_benchmark] #[bench::text(setup = setup_dm_send)] fn bench_dm_send(mut d: DmSendData) { + let signal_addr = d.bob_jid.to_protocol_address(); let node = futures::executor::block_on(prepare_peer_stanza( &mut d.alice.sessions, &mut d.alice.identity, - d.bob_jid.clone(), d.bob_jid, + &signal_addr, &d.msg, "b-001".into(), )) diff --git a/wacore/src/iq/prekeys.rs b/wacore/src/iq/prekeys.rs index 801d2b070..c0094bff6 100644 --- a/wacore/src/iq/prekeys.rs +++ b/wacore/src/iq/prekeys.rs @@ -367,13 +367,13 @@ impl IqSpec for PreKeyUploadSpec { fn build_iq(&self) -> InfoQuery<'static> { let content = PreKeyUtils::build_upload_prekeys_request( self.registration_id, - self.identity_key.public_key_bytes().to_vec(), + self.identity_key.public_key_bytes(), self.signed_pre_key_id, - self.signed_pre_key_public.public_key_bytes().to_vec(), - self.signed_pre_key_signature.clone(), + self.signed_pre_key_public.public_key_bytes(), + &self.signed_pre_key_signature, self.pre_keys .iter() - .map(|(id, pk)| (*id, pk.public_key_bytes().to_vec())), + .map(|(id, pk)| (*id, pk.public_key_bytes())), ); InfoQuery::set( diff --git a/wacore/src/prekeys.rs b/wacore/src/prekeys.rs index a909ed00a..2112c628d 100644 --- a/wacore/src/prekeys.rs +++ b/wacore/src/prekeys.rs @@ -41,30 +41,29 @@ impl PreKeyUtils { NodeBuilder::new("key").children(user_nodes).build() } - pub fn build_upload_prekeys_request( + pub fn build_upload_prekeys_request<'a>( registration_id: u32, - identity_key_bytes: Vec, + identity_key_bytes: &[u8], signed_pre_key_id: u32, - signed_pre_key_public_bytes: Vec, - signed_pre_key_signature: Vec, - pre_keys: impl IntoIterator)>, + signed_pre_key_public_bytes: &[u8], + signed_pre_key_signature: &[u8], + pre_keys: impl IntoIterator, ) -> Vec { let pre_keys = pre_keys.into_iter(); let (lower, upper) = pre_keys.size_hint(); let mut pre_key_nodes = Vec::with_capacity(upper.unwrap_or(lower)); for (pre_key_id, public_bytes) in pre_keys { - let id_bytes = pre_key_id.to_be_bytes()[1..].to_vec(); let node = NodeBuilder::new("key") .children([ - NodeBuilder::new("id").bytes(id_bytes).build(), + NodeBuilder::new("id") + .bytes(pre_key_id.to_be_bytes()[1..].to_vec()) + .build(), NodeBuilder::new("value").bytes(public_bytes).build(), ]) .build(); pre_key_nodes.push(node); } - let registration_id_bytes = registration_id.to_be_bytes().to_vec(); - let signed_pre_key_node = NodeBuilder::new("skey") .children([ NodeBuilder::new("id") @@ -79,13 +78,11 @@ impl PreKeyUtils { ]) .build(); - let type_bytes = vec![5u8]; - vec![ NodeBuilder::new("registration") - .bytes(registration_id_bytes) + .bytes(registration_id.to_be_bytes().to_vec()) .build(), - NodeBuilder::new("type").bytes(type_bytes).build(), + NodeBuilder::new("type").bytes(vec![5u8]).build(), NodeBuilder::new("identity") .bytes(identity_key_bytes) .build(), diff --git a/wacore/src/send.rs b/wacore/src/send.rs index 44ad70363..d9b8cef7f 100644 --- a/wacore/src/send.rs +++ b/wacore/src/send.rs @@ -1,7 +1,7 @@ use crate::client::context::{GroupInfo, SendContextResolver}; use crate::libsignal::protocol::{ - CiphertextMessage, SENDERKEY_MESSAGE_CURRENT_VERSION, SenderKeyMessage, SenderKeyStore, - SignalProtocolError, UsePQRatchet, message_encrypt, process_prekey_bundle, + CiphertextMessage, ProtocolAddress, SENDERKEY_MESSAGE_CURRENT_VERSION, SenderKeyMessage, + SenderKeyStore, SignalProtocolError, UsePQRatchet, message_encrypt, process_prekey_bundle, }; use crate::messages::MessageUtils; use crate::reporting_token::{ @@ -247,7 +247,7 @@ pub fn should_hide_decrypt_fail(msg: &wa::Message) -> bool { pub async fn encrypt_group_message( sender_key_store: &mut S, group_jid: &Jid, - sender_jid: &Jid, + sender_address: &ProtocolAddress, plaintext: &[u8], csprng: &mut R, ) -> Result @@ -255,8 +255,7 @@ where S: SenderKeyStore + ?Sized, R: Rng + CryptoRng, { - let sender_address = sender_jid.to_protocol_address(); - let sender_key_name = make_sender_key_name(group_jid, &sender_address); + let sender_key_name = make_sender_key_name(group_jid, sender_address); log::debug!( "Attempting to load sender key for group {} sender {}", sender_key_name.group_id(), @@ -830,7 +829,7 @@ pub async fn prepare_peer_stanza( session_store: &mut S, identity_store: &mut I, transport_jid: Jid, - encryption_jid: Jid, + signal_address: &ProtocolAddress, message: &wa::Message, request_id: String, ) -> Result @@ -839,10 +838,9 @@ where I: crate::libsignal::protocol::IdentityKeyStore, { let plaintext = MessageUtils::encode_and_pad(message); - let signal_address = encryption_jid.to_protocol_address(); let encrypted_message = - message_encrypt(&plaintext, &signal_address, session_store, identity_store).await?; + message_encrypt(&plaintext, signal_address, session_store, identity_store).await?; let (enc_type, _, serialized_bytes) = extract_ciphertext(encrypted_message) .ok_or_else(|| anyhow!("Unexpected peer encryption message type"))?; @@ -1053,6 +1051,8 @@ pub async fn prepare_group_stanza< let mut phash_for_stanza: Option = None; let mut skdm_encrypted_devices: Vec = Vec::new(); + let sender_address = own_sending_jid.to_protocol_address(); + // Determine if we need to distribute SKDM and to which devices let distribution_list: Option> = if let Some(target_devices) = skdm_target_devices { // Use the specific list of devices that need SKDM @@ -1183,7 +1183,7 @@ pub async fn prepare_group_stanza< let axolotl_skdm_bytes = create_sender_key_distribution_message_for_group( stores.sender_key_store, &to_jid, - &own_sending_jid, + &sender_address, ) .await?; @@ -1247,7 +1247,7 @@ pub async fn prepare_group_stanza< let skmsg = encrypt_group_message( stores.sender_key_store, &to_jid, - &own_sending_jid, + &sender_address, &plaintext, &mut rand::make_rng::(), ) @@ -1334,10 +1334,9 @@ pub async fn prepare_group_stanza< pub async fn create_sender_key_distribution_message_for_group( store: &mut (dyn SenderKeyStore + Send + Sync), group_jid: &Jid, - own_sending_jid: &Jid, + sender_address: &ProtocolAddress, ) -> Result> { - let sender_address = own_sending_jid.to_protocol_address(); - let sender_key_name = make_sender_key_name(group_jid, &sender_address); + let sender_key_name = make_sender_key_name(group_jid, sender_address); let mut rng = rand::make_rng::(); let skdm = crate::libsignal::protocol::create_sender_key_distribution_message( From da9bf0c78ff1976895e347024b9093006bc21d8f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Wed, 15 Apr 2026 10:21:14 -0300 Subject: [PATCH 04/17] perf: pool zlib decompressor to avoid repeated 48KB allocations flate2::Decompress allocates ~48KB of internal zlib state on every Decompress::new(). DHAT showed 24 instances (1.1MB binary protocol + 786KB history sync) created and destroyed during a session. Use a thread_local pool that reuses both the Decompress instance (via reset(true)) and the output Vec between calls. The pooled decompressor is used for both binary protocol node decompression and history sync blob inflation. Measured impact: - DHAT total: 32.1 MB -> 30.6 MB (-1.5 MB, -4.6%) - connect_to_ready: 4.25 MB -> 4.12 MB (-130 KB) --- wacore/binary/src/lib.rs | 1 + wacore/binary/src/util.rs | 19 ++--------- wacore/binary/src/zlib_pool.rs | 61 ++++++++++++++++++++++++++++++++++ wacore/src/history_sync.rs | 18 +++------- 4 files changed, 69 insertions(+), 30 deletions(-) create mode 100644 wacore/binary/src/zlib_pool.rs diff --git a/wacore/binary/src/lib.rs b/wacore/binary/src/lib.rs index bde27e613..ee41589d1 100644 --- a/wacore/binary/src/lib.rs +++ b/wacore/binary/src/lib.rs @@ -11,6 +11,7 @@ pub mod marshal; pub mod node; pub mod token; pub mod util; +pub mod zlib_pool; pub use attrs::{AttrParser, AttrParserRef}; pub use compact_str::CompactString; diff --git a/wacore/binary/src/util.rs b/wacore/binary/src/util.rs index 0be9f7fc2..8ac9f320c 100644 --- a/wacore/binary/src/util.rs +++ b/wacore/binary/src/util.rs @@ -1,27 +1,14 @@ use crate::error::{BinaryError, Result}; +use crate::zlib_pool::decompress_zlib_pooled; use bytes::{Buf, Bytes, BytesMut}; -use flate2::read::ZlibDecoder; use std::borrow::Cow; -use std::io::Read; /// Protocol frames larger than 16 MiB after decompression are rejected. -/// WhatsApp messages are typically small; this guards against malicious -/// or corrupt compressed payloads that would expand into huge allocations. const MAX_DECOMPRESSED_SIZE: u64 = 16 * 1024 * 1024; fn decompress_zlib(compressed: &[u8]) -> Result> { - let estimated = (compressed.len() * 4).clamp(256, 64 * 1024); - let mut out = Vec::with_capacity(estimated); - ZlibDecoder::new(compressed) - .take(MAX_DECOMPRESSED_SIZE + 1) - .read_to_end(&mut out) - .map_err(|e| BinaryError::Zlib(e.to_string()))?; - if out.len() as u64 > MAX_DECOMPRESSED_SIZE { - return Err(BinaryError::Zlib(format!( - "decompressed payload exceeds {MAX_DECOMPRESSED_SIZE} bytes" - ))); - } - Ok(out) + decompress_zlib_pooled(compressed, MAX_DECOMPRESSED_SIZE) + .map_err(|e| BinaryError::Zlib(e.to_string())) } pub fn unpack(data: &[u8]) -> Result> { diff --git a/wacore/binary/src/zlib_pool.rs b/wacore/binary/src/zlib_pool.rs new file mode 100644 index 000000000..3653adf30 --- /dev/null +++ b/wacore/binary/src/zlib_pool.rs @@ -0,0 +1,61 @@ +use flate2::{Decompress, FlushDecompress, Status}; +use std::cell::RefCell; +use std::io; + +thread_local! { + static DECOMPRESSOR: RefCell<(Decompress, Vec)> = RefCell::new(( + Decompress::new(true), + Vec::with_capacity(4096), + )); +} + +/// Decompress zlib data using a pooled decompressor. +/// +/// Reuses the `flate2::Decompress` internal state (~48 KB) and the output +/// buffer across calls on the same thread, avoiding repeated heap allocations. +pub fn decompress_zlib_pooled(compressed: &[u8], max_size: u64) -> io::Result> { + DECOMPRESSOR.with(|cell| { + let (decompressor, scratch) = &mut *cell.borrow_mut(); + decompressor.reset(true); + scratch.clear(); + + let estimated = (compressed.len() * 4).clamp(256, 64 * 1024); + if scratch.capacity() < estimated { + scratch.reserve(estimated - scratch.capacity()); + } + + let mut input_offset = 0; + loop { + let status = decompressor + .decompress_vec( + &compressed[input_offset..], + scratch, + FlushDecompress::Finish, + ) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + + input_offset = decompressor.total_in() as usize; + + if scratch.len() as u64 > max_size { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("decompressed payload exceeds {max_size} bytes"), + )); + } + + match status { + Status::StreamEnd => break, + Status::Ok | Status::BufError => { + // Need more output space + scratch.reserve(scratch.capacity().max(4096)); + } + } + } + + // Return the data by swapping out to avoid cloning. + // The scratch buffer keeps its capacity for the next call. + let mut result = Vec::new(); + std::mem::swap(scratch, &mut result); + Ok(result) + }) +} diff --git a/wacore/src/history_sync.rs b/wacore/src/history_sync.rs index c511352b2..46fec6046 100644 --- a/wacore/src/history_sync.rs +++ b/wacore/src/history_sync.rs @@ -1,7 +1,6 @@ use bytes::Bytes; -use flate2::read::ZlibDecoder; -use std::io::Read; use thiserror::Error; +use wacore_binary::zlib_pool::decompress_zlib_pooled; #[derive(Debug, Error)] pub enum HistorySyncError { @@ -49,23 +48,14 @@ pub fn process_history_sync( compressed_data: Vec, own_user: Option<&str>, retain_blob: bool, - compressed_size_hint: Option, + _compressed_size_hint: Option, ) -> Result { // Hard limit to prevent OOM on malformed blobs. // Typical InitialBootstrap: 5-20 MB decompressed. const MAX_DECOMPRESSED: u64 = 64 * 1024 * 1024; - let estimated = compressed_size_hint - .and_then(|s| usize::try_from(s).ok()) - .map(|s| s * 4) - .unwrap_or_else(|| compressed_data.len() * 4) - .clamp(256, MAX_DECOMPRESSED as usize); - let mut decompressed = Vec::with_capacity(estimated); - { - let decoder = ZlibDecoder::new(compressed_data.as_slice()); - let mut limited = decoder.take(MAX_DECOMPRESSED); - limited.read_to_end(&mut decompressed)?; - } + let decompressed = decompress_zlib_pooled(&compressed_data, MAX_DECOMPRESSED) + .map_err(HistorySyncError::DecompressionError)?; drop(compressed_data); let buf = Bytes::from(decompressed); From a1f68b92ef7fe49d7f7e2d872348c82a9b9e65ed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Wed, 15 Apr 2026 10:33:20 -0300 Subject: [PATCH 05/17] perf: use Bytes for prekey store to eliminate O(n) clones Change SignalStore trait prekey methods from Vec to bytes::Bytes. Bytes::clone() is O(1) (atomic refcount) vs Vec::clone() which copies the entire buffer. The prekey upload path calls store_prekeys_batch twice (pre-upload + mark-uploaded) with 812 keys each. Measured impact: - connect_to_ready allocs: 21,397 -> 19,638 (-1,759 allocs, -8.2%) - connect_to_ready bytes: 4.12 MB -> 4.01 MB (-106 KB) - reconnect bytes: 522 KB -> 511 KB --- Cargo.lock | 1 + src/appstate_sync.rs | 2 +- src/prekeys.rs | 9 +++++---- src/store/signal.rs | 2 +- storages/sqlite-storage/Cargo.toml | 1 + storages/sqlite-storage/src/sqlite_store.rs | 20 ++++++++++++-------- wacore/src/store/in_memory.rs | 11 ++++++----- wacore/src/store/traits.rs | 7 ++++--- 8 files changed, 31 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e442964a5..23e7d5b7d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2420,6 +2420,7 @@ version = "0.5.0" dependencies = [ "async-trait", "bincode 2.0.1", + "bytes", "diesel", "diesel_migrations", "libsqlite3-sys", diff --git a/src/appstate_sync.rs b/src/appstate_sync.rs index 06556e479..0426e94e1 100644 --- a/src/appstate_sync.rs +++ b/src/appstate_sync.rs @@ -59,7 +59,7 @@ mod tests { async fn store_prekey(&self, _: u32, _: &[u8], _: bool) -> StoreResult<()> { Ok(()) } - async fn load_prekey(&self, _: u32) -> StoreResult>> { + async fn load_prekey(&self, _: u32) -> StoreResult> { Ok(None) } async fn remove_prekey(&self, _: u32) -> StoreResult<()> { diff --git a/src/prekeys.rs b/src/prekeys.rs index 110bded44..88f08924f 100644 --- a/src/prekeys.rs +++ b/src/prekeys.rs @@ -144,11 +144,12 @@ impl Client { } // Encode once — reused for both pre-upload store and post-upload mark. - let encoded_batch: Vec<(u32, Vec)> = { + // Uses Bytes so store_prekeys_batch clones are O(1) refcount bumps. + let encoded_batch: Vec<(u32, bytes::Bytes)> = { use prost::Message; keys_to_upload .iter() - .map(|(id, record)| (*id, record.encode_to_vec())) + .map(|(id, record)| (*id, bytes::Bytes::from(record.encode_to_vec()))) .collect() }; @@ -339,7 +340,7 @@ impl Client { // Build a lookup so we preserve the server-requested order. // Dedupe the expected count since the server may send duplicate IDs. - let loaded_map: std::collections::HashMap> = loaded.into_iter().collect(); + let loaded_map: std::collections::HashMap = loaded.into_iter().collect(); let unique_requested: std::collections::HashSet<&u32> = response.prekey_ids.iter().collect(); @@ -358,7 +359,7 @@ impl Client { return Ok(()); }; use prost::Message; - match waproto::whatsapp::PreKeyRecordStructure::decode(record_bytes.as_slice()) { + match waproto::whatsapp::PreKeyRecordStructure::decode(record_bytes.as_ref()) { Ok(record) => { if let Some(pk) = record.public_key { prekey_pubkeys.push(pk); diff --git a/src/store/signal.rs b/src/store/signal.rs index 95dce8a2a..6ac558e85 100644 --- a/src/store/signal.rs +++ b/src/store/signal.rs @@ -276,7 +276,7 @@ impl PreKeyStore for Device { match self.backend.load_prekey(prekey_id).await { Ok(Some(bytes)) => { // Try new format first (protobuf-encoded PreKeyRecordStructure) - if let Ok(record) = PreKeyRecordStructure::decode(bytes.as_slice()) { + if let Ok(record) = PreKeyRecordStructure::decode(bytes.as_ref()) { return Ok(Some(record)); } diff --git a/storages/sqlite-storage/Cargo.toml b/storages/sqlite-storage/Cargo.toml index 59ba57ca6..838c44399 100644 --- a/storages/sqlite-storage/Cargo.toml +++ b/storages/sqlite-storage/Cargo.toml @@ -14,6 +14,7 @@ bundled-sqlite = ["libsqlite3-sys/bundled"] [dependencies] async-trait = { workspace = true } bincode = { version = "2.0.1", features = ["serde"] } +bytes = { workspace = true } diesel = { version = "2.3.7", default-features = false, features = [ "sqlite", "r2d2", diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index d79bfae9e..4748a00ef 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -1,5 +1,6 @@ use crate::schema::*; use async_trait::async_trait; +use bytes::Bytes; use diesel::prelude::*; use diesel::r2d2::{ConnectionManager, Pool}; use diesel::result::{DatabaseErrorKind, Error as DieselError}; @@ -1279,7 +1280,7 @@ impl SignalStore for SqliteStore { )) } - async fn store_prekeys_batch(&self, keys: &[(u32, Vec)], uploaded: bool) -> Result<()> { + async fn store_prekeys_batch(&self, keys: &[(u32, Bytes)], uploaded: bool) -> Result<()> { if keys.is_empty() { return Ok(()); } @@ -1287,7 +1288,7 @@ impl SignalStore for SqliteStore { let pool = self.pool.clone(); let db_semaphore = self.db_semaphore.clone(); let device_id = self.device_id; - let keys = keys.to_vec(); + let keys: Vec<(u32, Vec)> = keys.iter().map(|(id, b)| (*id, b.to_vec())).collect(); const MAX_RETRIES: u32 = 5; @@ -1346,10 +1347,10 @@ impl SignalStore for SqliteStore { )) } - async fn load_prekey(&self, id: u32) -> Result>> { + async fn load_prekey(&self, id: u32) -> Result> { let pool = self.pool.clone(); let device_id = self.device_id; - tokio::task::spawn_blocking(move || -> Result>> { + tokio::task::spawn_blocking(move || -> Result> { let mut conn = pool .get() .map_err(|e| StoreError::Connection(e.to_string()))?; @@ -1360,20 +1361,20 @@ impl SignalStore for SqliteStore { .first(&mut conn) .optional() .map_err(|e| StoreError::Database(e.to_string()))?; - Ok(res) + Ok(res.map(Bytes::from)) }) .await .map_err(|e| StoreError::Database(e.to_string()))? } - async fn load_prekeys_batch(&self, ids: &[u32]) -> Result)>> { + async fn load_prekeys_batch(&self, ids: &[u32]) -> Result> { if ids.is_empty() { return Ok(Vec::new()); } let pool = self.pool.clone(); let device_id = self.device_id; let ids: Vec = ids.iter().map(|&id| id as i32).collect(); - self.with_semaphore(move || -> Result)>> { + self.with_semaphore(move || -> Result> { let mut conn = pool .get() .map_err(|e| StoreError::Connection(e.to_string()))?; @@ -1383,7 +1384,10 @@ impl SignalStore for SqliteStore { .filter(prekeys::device_id.eq(device_id)) .load(&mut conn) .map_err(|e| StoreError::Database(e.to_string()))?; - Ok(rows.into_iter().map(|(id, key)| (id as u32, key)).collect()) + Ok(rows + .into_iter() + .map(|(id, key)| (id as u32, Bytes::from(key))) + .collect()) }) .await } diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index a5f4a2b68..ea942fc19 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -13,6 +13,7 @@ use crate::store::error::Result; use crate::store::traits::*; use async_lock::Mutex; use async_trait::async_trait; +use bytes::Bytes; use wacore_appstate::processor::AppStateMutationMAC; /// Key for the sent-message store: `(chat_jid, message_id)`. @@ -26,7 +27,7 @@ struct SentMessageEntry { /// Key for pre-keys: `id`. struct PreKeyEntry { - record: Vec, + record: Bytes, } /// Key for base-key collision detection: `(address, message_id)`. @@ -168,13 +169,13 @@ impl SignalStore for InMemoryBackend { self.state.lock().await.prekeys.insert( id, PreKeyEntry { - record: record.to_vec(), + record: Bytes::copy_from_slice(record), }, ); Ok(()) } - async fn store_prekeys_batch(&self, keys: &[(u32, Vec)], _uploaded: bool) -> Result<()> { + async fn store_prekeys_batch(&self, keys: &[(u32, Bytes)], _uploaded: bool) -> Result<()> { let mut state = self.state.lock().await; for (id, record) in keys { state.prekeys.insert( @@ -187,7 +188,7 @@ impl SignalStore for InMemoryBackend { Ok(()) } - async fn load_prekey(&self, id: u32) -> Result>> { + async fn load_prekey(&self, id: u32) -> Result> { Ok(self .state .lock() @@ -197,7 +198,7 @@ impl SignalStore for InMemoryBackend { .map(|e| e.record.clone())) } - async fn load_prekeys_batch(&self, ids: &[u32]) -> Result)>> { + async fn load_prekeys_batch(&self, ids: &[u32]) -> Result> { let state = self.state.lock().await; let mut result = Vec::with_capacity(ids.len()); for &id in ids { diff --git a/wacore/src/store/traits.rs b/wacore/src/store/traits.rs index 875e1b6fe..f3717c389 100644 --- a/wacore/src/store/traits.rs +++ b/wacore/src/store/traits.rs @@ -10,6 +10,7 @@ use crate::appstate::hash::HashState; use crate::store::error::Result; use async_trait::async_trait; +use bytes::Bytes; use serde::{Deserialize, Serialize}; use wacore_appstate::processor::AppStateMutationMAC; @@ -116,7 +117,7 @@ pub trait SignalStore: Send + Sync { /// Store multiple pre-keys in a single batch operation. /// Default implementation falls back to individual `store_prekey` calls. - async fn store_prekeys_batch(&self, keys: &[(u32, Vec)], uploaded: bool) -> Result<()> { + async fn store_prekeys_batch(&self, keys: &[(u32, Bytes)], uploaded: bool) -> Result<()> { for (id, record) in keys { self.store_prekey(*id, record, uploaded).await?; } @@ -124,11 +125,11 @@ pub trait SignalStore: Send + Sync { } /// Load a pre-key by ID. - async fn load_prekey(&self, id: u32) -> Result>>; + async fn load_prekey(&self, id: u32) -> Result>; /// Load multiple pre-keys by ID in a single batch operation. /// Returns only the keys that exist. - async fn load_prekeys_batch(&self, ids: &[u32]) -> Result)>> { + async fn load_prekeys_batch(&self, ids: &[u32]) -> Result> { let mut result = Vec::with_capacity(ids.len()); for &id in ids { if let Some(record) = self.load_prekey(id).await? { From d692743fdc70425c067f6fbe596877ecb4564092 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Wed, 15 Apr 2026 11:02:49 -0300 Subject: [PATCH 06/17] perf: CompactString for props cache + Bytes transport + sender buffer reuse Props: Switch AbProp.config_value and AbPropsCache storage from String to CompactString. Most prop values (\"1\", \"true\", \"enabled\") are <=24 bytes and stored inline without heap allocation. Also avoids String clone on cache insertion. Transport: Change Transport::send from Vec to Bytes. The sender task in NoiseSocket now owns a reusable Vec for framing output instead of receiving a fresh one per send. Callers no longer need to pre-allocate an encrypted_buf. Measured impact: - connect_to_ready allocs: 19,638 -> 17,172 (-2,466 allocs, -12.6%) - connect_to_ready bytes: 4.01 MB -> 3.98 MB - DHAT total blocks: 147,279 -> 131,823 (-15,456 blocks, -10.5%) - DHAT total bytes: 29.9 MB -> 29.9 MB (block count reduction, not bytes) --- src/client.rs | 5 +- src/handshake.rs | 4 +- src/socket/noise_socket.rs | 73 ++++++++------------------- src/transport.rs | 2 +- transports/tokio-transport/src/lib.rs | 2 +- wacore/src/iq/props.rs | 19 +++---- wacore/src/net.rs | 2 +- wacore/src/store/ab_props.rs | 7 +-- 8 files changed, 40 insertions(+), 74 deletions(-) diff --git a/src/client.rs b/src/client.rs index 5f399754a..2e055e5d7 100644 --- a/src/client.rs +++ b/src/client.rs @@ -893,10 +893,7 @@ impl Client { /// [`send_node`](Client::send_node) for normal stanza sending. pub async fn send_raw_bytes(&self, plaintext: Vec) -> Result<(), ClientError> { let noise_socket = self.get_noise_socket().await?; - let encrypted_buf = Vec::with_capacity(plaintext.len() + 32); - noise_socket - .encrypt_and_send(plaintext, encrypted_buf) - .await?; + noise_socket.encrypt_and_send(plaintext).await?; self.last_data_sent_ms .store(wacore::time::now_millis().max(0) as u64, Ordering::Relaxed); Ok(()) diff --git a/src/handshake.rs b/src/handshake.rs index 779e2e963..f3f4ee1d2 100644 --- a/src/handshake.rs +++ b/src/handshake.rs @@ -71,7 +71,7 @@ pub async fn do_handshake( // First message includes the WA connection header (with optional edge routing) let framed = wacore::framing::encode_frame(&client_hello_bytes, Some(&header)) .map_err(HandshakeError::Transport)?; - transport.send(framed).await?; + transport.send(bytes::Bytes::from(framed)).await?; // Wait for server response frame let resp_frame = loop { @@ -113,7 +113,7 @@ pub async fn do_handshake( // Subsequent messages don't need the header let framed = wacore::framing::encode_frame(&client_finish_bytes, None) .map_err(HandshakeError::Transport)?; - transport.send(framed).await?; + transport.send(bytes::Bytes::from(framed)).await?; let (write_key, read_key) = handshake_state.finish()?; info!("Handshake complete, switching to encrypted communication"); diff --git a/src/socket/noise_socket.rs b/src/socket/noise_socket.rs index 44292b220..3ed2ad00e 100644 --- a/src/socket/noise_socket.rs +++ b/src/socket/noise_socket.rs @@ -16,7 +16,6 @@ type SendResult = std::result::Result<(), EncryptSendError>; /// A job sent to the dedicated sender task. struct SendJob { plaintext_buf: Vec, - out_buf: Vec, response_tx: oneshot::Sender, } @@ -79,6 +78,7 @@ impl NoiseSocket { send_job_rx: async_channel::Receiver, ) { let mut write_counter: u32 = 0; + let mut out_buf = Vec::with_capacity(4096); while let Ok(job) = send_job_rx.recv().await { let result = Self::process_send_job( @@ -87,15 +87,12 @@ impl NoiseSocket { &write_key, &mut write_counter, job.plaintext_buf, - job.out_buf, + &mut out_buf, ) .await; - // Send result back to caller. Ignore error if receiver was dropped. let _ = job.response_tx.send(result); } - - // Channel closed - NoiseSocket was dropped, task exits naturally } /// Process a single send job: encrypt and send the message. @@ -105,26 +102,21 @@ impl NoiseSocket { write_key: &Arc, write_counter: &mut u32, mut plaintext_buf: Vec, - mut out_buf: Vec, + out_buf: &mut Vec, ) -> SendResult { let counter = *write_counter; - // For small messages, encrypt plaintext_buf in-place then frame into out_buf. - // This avoids the previous triple-copy pattern (plaintext→out→plaintext→out). if plaintext_buf.len() <= INLINE_ENCRYPT_THRESHOLD { if let Err(e) = write_key.encrypt_in_place_with_counter(counter, &mut plaintext_buf) { return Err(EncryptSendError::crypto(anyhow::anyhow!(e.to_string()))); } - // Frame the ciphertext from plaintext_buf into out_buf (single copy) out_buf.clear(); - if let Err(e) = wacore::framing::encode_frame_into(&plaintext_buf, None, &mut out_buf) { + if let Err(e) = wacore::framing::encode_frame_into(&plaintext_buf, None, out_buf) { return Err(EncryptSendError::framing(e)); } } else { - // Offload larger messages to a blocking thread let write_key = write_key.clone(); - let plaintext_arc = Arc::new(plaintext_buf); let plaintext_arc_for_task = plaintext_arc.clone(); @@ -133,9 +125,7 @@ impl NoiseSocket { }) .await; - // Recover ownership so the buffer is dropped at end of scope - plaintext_buf = Arc::try_unwrap(plaintext_arc).unwrap_or_else(|arc| (*arc).clone()); - drop(plaintext_buf); + drop(plaintext_arc); let ciphertext = match encrypt_result { Ok(c) => c, @@ -145,28 +135,27 @@ impl NoiseSocket { }; out_buf.clear(); - if let Err(e) = wacore::framing::encode_frame_into(&ciphertext, None, &mut out_buf) { + if let Err(e) = wacore::framing::encode_frame_into(&ciphertext, None, out_buf) { return Err(EncryptSendError::framing(e)); } } - if let Err(e) = transport.send(out_buf).await { + // Swap out the framed data; the Vec retains capacity after drain + let frame = bytes::Bytes::from(std::mem::take(out_buf)); + if let Err(e) = transport.send(frame).await { return Err(EncryptSendError::transport(e)); } - // Only advance the counter after the encrypted frame was successfully sent. - // If transport.send() fails, we can retry with the same counter value. *write_counter = write_counter.wrapping_add(1); Ok(()) } - pub async fn encrypt_and_send(&self, plaintext_buf: Vec, out_buf: Vec) -> SendResult { + pub async fn encrypt_and_send(&self, plaintext_buf: Vec) -> SendResult { let (response_tx, response_rx) = oneshot::channel(); let job = SendJob { plaintext_buf, - out_buf, response_tx, }; @@ -217,9 +206,8 @@ mod tests { ); let plaintext_buf = Vec::with_capacity(1024); - let encrypted_buf = Vec::with_capacity(1024); - let result = socket.encrypt_and_send(plaintext_buf, encrypted_buf).await; + let result = socket.encrypt_and_send(plaintext_buf).await; assert!(result.is_ok(), "encrypt_and_send should succeed"); } @@ -240,8 +228,9 @@ mod tests { #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl crate::transport::Transport for RecordingTransport { - async fn send(&self, mut data: Vec) -> std::result::Result<(), anyhow::Error> { + async fn send(&self, data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> { if data.len() > 16 { + let mut data = data.to_vec(); // Strip the 3-byte frame header, then decrypt in place data.drain(..3); let counter = self @@ -291,8 +280,7 @@ mod tests { // Use index as the first byte of plaintext to identify this send let mut plaintext = vec![i as u8]; plaintext.extend_from_slice(&[0u8; 99]); - let out_buf = Vec::with_capacity(256); - socket.encrypt_and_send(plaintext, out_buf).await + socket.encrypt_and_send(plaintext).await })); } @@ -324,7 +312,7 @@ mod tests { #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl crate::transport::Transport for SizeRecordingTransport { - async fn send(&self, data: Vec) -> std::result::Result<(), anyhow::Error> { + async fn send(&self, data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> { self.last_size.store(data.len(), Ordering::SeqCst); Ok(()) } @@ -352,13 +340,7 @@ mod tests { for size in test_sizes { let plaintext = vec![0xABu8; size]; - // This is the formula used in client.rs - let buffer_capacity = plaintext.len() + 32; - let encrypted_buf = Vec::with_capacity(buffer_capacity); - - let result = socket - .encrypt_and_send(plaintext.clone(), encrypted_buf) - .await; + let result = socket.encrypt_and_send(plaintext.clone()).await; assert!( result.is_ok(), @@ -376,15 +358,6 @@ mod tests { "Encrypted size for {} byte payload should be {} (got {})", size, expected_max, actual_encrypted_size ); - - // Verify our buffer sizing formula provides enough capacity - assert!( - buffer_capacity >= actual_encrypted_size, - "Buffer capacity {} should be >= encrypted size {} for payload size {}", - buffer_capacity, - actual_encrypted_size, - size - ); } } @@ -399,7 +372,7 @@ mod tests { #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl crate::transport::Transport for NoOpTransport { - async fn send(&self, _data: Vec) -> std::result::Result<(), anyhow::Error> { + async fn send(&self, _data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> { Ok(()) } async fn disconnect(&self) {} @@ -418,16 +391,12 @@ mod tests { ); // Test empty payload - let result = socket - .encrypt_and_send(vec![], Vec::with_capacity(32)) - .await; + let result = socket.encrypt_and_send(vec![]).await; assert!(result.is_ok(), "Empty payload should encrypt successfully"); // Test payload at inline threshold boundary (16KB) let at_threshold = vec![0u8; 16 * 1024]; - let result = socket - .encrypt_and_send(at_threshold, Vec::with_capacity(16 * 1024 + 32)) - .await; + let result = socket.encrypt_and_send(at_threshold).await; assert!( result.is_ok(), "Payload at inline threshold should encrypt successfully" @@ -435,9 +404,7 @@ mod tests { // Test payload just above inline threshold let above_threshold = vec![0u8; 16 * 1024 + 1]; - let result = socket - .encrypt_and_send(above_threshold, Vec::with_capacity(16 * 1024 + 33)) - .await; + let result = socket.encrypt_and_send(above_threshold).await; assert!( result.is_ok(), "Payload above inline threshold should encrypt successfully" diff --git a/src/transport.rs b/src/transport.rs index 69ffde74a..91959d913 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -21,7 +21,7 @@ pub mod mock { #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl Transport for MockTransport { - async fn send(&self, _data: Vec) -> Result<(), anyhow::Error> { + async fn send(&self, _data: bytes::Bytes) -> Result<(), anyhow::Error> { Ok(()) } diff --git a/transports/tokio-transport/src/lib.rs b/transports/tokio-transport/src/lib.rs index fbc7215eb..d2569bd61 100644 --- a/transports/tokio-transport/src/lib.rs +++ b/transports/tokio-transport/src/lib.rs @@ -131,7 +131,7 @@ impl WsTransport { #[async_trait] impl Transport for WsTransport { - async fn send(&self, data: Vec) -> Result<(), anyhow::Error> { + async fn send(&self, data: bytes::Bytes) -> Result<(), anyhow::Error> { let mut guard = self.sink.lock().await; let sink = guard .as_mut() diff --git a/wacore/src/iq/props.rs b/wacore/src/iq/props.rs index 28181449a..7b99f8800 100644 --- a/wacore/src/iq/props.rs +++ b/wacore/src/iq/props.rs @@ -26,6 +26,7 @@ use crate::iq::spec::IqSpec; use crate::protocol::ProtocolNode; use crate::request::InfoQuery; +use wacore_binary::CompactString; use wacore_binary::builder::NodeBuilder; use wacore_binary::{Jid, Server}; use wacore_binary::{Node, NodeContent, NodeRef}; @@ -68,8 +69,8 @@ pub const PROPS_PROTOCOL_VERSION: &str = "1"; pub struct AbProp { /// The config code (property identifier). pub config_code: u32, - /// The config value. - pub config_value: String, + /// The config value. CompactString inlines values <=24 bytes (covers most props). + pub config_value: CompactString, /// Optional experiment exposure key. pub config_expo_key: Option, } @@ -82,7 +83,7 @@ impl crate::protocol::ProtocolNode for AbProp { fn into_node(self) -> Node { let mut builder = NodeBuilder::new("prop") .attr("config_code", self.config_code.to_string()) - .attr("config_value", &self.config_value); + .attr("config_value", &*self.config_value); if let Some(expo_key) = self.config_expo_key { builder = builder.attr("config_expo_key", expo_key.to_string()); @@ -105,8 +106,8 @@ impl crate::protocol::ProtocolNode for AbProp { return Err(anyhow::anyhow!("config_code must be >= 1")); } let config_value = optional_attr(node, "config_value") - .ok_or_else(|| anyhow::anyhow!("missing config_value in prop"))? - .into_owned(); + .ok_or_else(|| anyhow::anyhow!("missing config_value in prop"))?; + let config_value = CompactString::from(config_value.as_ref()); let config_expo_key = optional_attr(node, "config_expo_key").and_then(|s| s.parse().ok()); Ok(Self { @@ -483,7 +484,7 @@ mod tests { fn test_ab_prop_protocol_node_round_trip() { let prop = AbProp { config_code: 123, - config_value: "test_value".to_string(), + config_value: "test_value".into(), config_expo_key: Some(456), }; @@ -499,7 +500,7 @@ mod tests { fn test_ab_prop_protocol_node_no_expo_key() { let prop = AbProp { config_code: 789, - config_value: "another_value".to_string(), + config_value: "another_value".into(), config_expo_key: None, }; @@ -522,7 +523,7 @@ mod tests { props: vec![ AbPropConfig::Experiment(AbProp { config_code: 100, - config_value: "value1".to_string(), + config_value: "value1".into(), config_expo_key: None, }), AbPropConfig::Sampling(SamplingProp { @@ -531,7 +532,7 @@ mod tests { }), AbPropConfig::Experiment(AbProp { config_code: 200, - config_value: "value2".to_string(), + config_value: "value2".into(), config_expo_key: Some(99), }), ], diff --git a/wacore/src/net.rs b/wacore/src/net.rs index 2859cf524..83fd28b1a 100644 --- a/wacore/src/net.rs +++ b/wacore/src/net.rs @@ -24,7 +24,7 @@ pub enum TransportEvent { #[cfg_attr(not(target_arch = "wasm32"), async_trait)] pub trait Transport: Send + Sync { /// Sends raw data to the server. - async fn send(&self, data: Vec) -> Result<(), anyhow::Error>; + async fn send(&self, data: Bytes) -> Result<(), anyhow::Error>; /// Closes the connection. async fn disconnect(&self); diff --git a/wacore/src/store/ab_props.rs b/wacore/src/store/ab_props.rs index 59eafdb60..b6afd0bcb 100644 --- a/wacore/src/store/ab_props.rs +++ b/wacore/src/store/ab_props.rs @@ -10,12 +10,13 @@ use std::collections::HashMap; use std::sync::atomic::{AtomicBool, Ordering}; use async_lock::RwLock; +use wacore_binary::CompactString; use crate::iq::props::{AbPropConfig, PropsResponse}; /// In-memory cache of AB experiment properties, populated on connect. pub struct AbPropsCache { - props: RwLock>, + props: RwLock>, /// Guards against applying a delta into an empty cache on cold start. seeded: AtomicBool, } @@ -50,7 +51,7 @@ impl AbPropsCache { } } - pub async fn get(&self, config_code: u32) -> Option { + pub async fn get(&self, config_code: u32) -> Option { self.props.read().await.get(&config_code).cloned() } @@ -103,7 +104,7 @@ mod tests { fn experiment(code: u32, value: &str) -> AbPropConfig { AbPropConfig::Experiment(AbProp { config_code: code, - config_value: value.to_string(), + config_value: value.into(), config_expo_key: None, }) } From 5e20548d95cf226ce2a6e83f70c12f95cd9bbebd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Wed, 15 Apr 2026 11:18:38 -0300 Subject: [PATCH 07/17] perf: shared prekey encode buffer + reusable sender encryption buffer Prekey encoding: Replace 812 individual encode_to_vec() calls with a single contiguous buffer + Bytes::slice() sub-views. All prekey records are encoded into one pre-sized Vec, then sliced into zero-copy Bytes views for the store and upload paths. Eliminates 4,872 small Vec allocs. Sender task: encrypt_and_send now accepts Bytes instead of Vec. The sender task owns a reusable enc_buf for in-place encryption, avoiding per-send buffer allocations for the common small-message path. Measured impact: - connect_to_ready allocs: 17,172 -> 15,597 (-1,575 allocs, -9.2%) - DHAT total blocks: 131,823 -> 123,528 (-8,295 blocks, -6.3%) - DHAT peak live (t-gmax): 20,951 blocks -> 9,285 blocks (-55.7%) --- src/client.rs | 4 +++- src/prekeys.rs | 22 +++++++++++++---- src/socket/noise_socket.rs | 48 +++++++++++++++++++++----------------- 3 files changed, 46 insertions(+), 28 deletions(-) diff --git a/src/client.rs b/src/client.rs index 2e055e5d7..5c0dec966 100644 --- a/src/client.rs +++ b/src/client.rs @@ -893,7 +893,9 @@ impl Client { /// [`send_node`](Client::send_node) for normal stanza sending. pub async fn send_raw_bytes(&self, plaintext: Vec) -> Result<(), ClientError> { let noise_socket = self.get_noise_socket().await?; - noise_socket.encrypt_and_send(plaintext).await?; + noise_socket + .encrypt_and_send(bytes::Bytes::from(plaintext)) + .await?; self.last_data_sent_ms .store(wacore::time::now_millis().max(0) as u64, Ordering::Relaxed); Ok(()) diff --git a/src/prekeys.rs b/src/prekeys.rs index 88f08924f..809d71492 100644 --- a/src/prekeys.rs +++ b/src/prekeys.rs @@ -143,13 +143,25 @@ impl Client { key_pairs_to_upload.push((pre_key_id, key_pair)); } - // Encode once — reused for both pre-upload store and post-upload mark. - // Uses Bytes so store_prekeys_batch clones are O(1) refcount bumps. + // Encode all prekey records into a single contiguous buffer, then slice + // into Bytes sub-views. This replaces 812 individual encode_to_vec() allocs + // with one large allocation + zero-copy slicing. let encoded_batch: Vec<(u32, bytes::Bytes)> = { use prost::Message; - keys_to_upload - .iter() - .map(|(id, record)| (*id, bytes::Bytes::from(record.encode_to_vec()))) + let total_len: usize = keys_to_upload.iter().map(|(_, r)| r.encoded_len()).sum(); + let mut buf = Vec::with_capacity(total_len); + let mut offsets = Vec::with_capacity(keys_to_upload.len()); + for (id, record) in &keys_to_upload { + let start = buf.len(); + record + .encode(&mut buf) + .expect("prost encode into pre-sized Vec"); + offsets.push((*id, start..buf.len())); + } + let shared = bytes::Bytes::from(buf); + offsets + .into_iter() + .map(|(id, range)| (id, shared.slice(range))) .collect() }; diff --git a/src/socket/noise_socket.rs b/src/socket/noise_socket.rs index 3ed2ad00e..90b4b9575 100644 --- a/src/socket/noise_socket.rs +++ b/src/socket/noise_socket.rs @@ -15,7 +15,7 @@ type SendResult = std::result::Result<(), EncryptSendError>; /// A job sent to the dedicated sender task. struct SendJob { - plaintext_buf: Vec, + plaintext: bytes::Bytes, response_tx: oneshot::Sender, } @@ -78,6 +78,8 @@ impl NoiseSocket { send_job_rx: async_channel::Receiver, ) { let mut write_counter: u32 = 0; + // Reusable buffers -- capacity stays allocated between sends + let mut enc_buf = Vec::with_capacity(4096); let mut out_buf = Vec::with_capacity(4096); while let Ok(job) = send_job_rx.recv().await { @@ -86,7 +88,8 @@ impl NoiseSocket { &transport, &write_key, &mut write_counter, - job.plaintext_buf, + &job.plaintext, + &mut enc_buf, &mut out_buf, ) .await; @@ -101,32 +104,34 @@ impl NoiseSocket { transport: &Arc, write_key: &Arc, write_counter: &mut u32, - mut plaintext_buf: Vec, + plaintext: &[u8], + enc_buf: &mut Vec, out_buf: &mut Vec, ) -> SendResult { let counter = *write_counter; - if plaintext_buf.len() <= INLINE_ENCRYPT_THRESHOLD { - if let Err(e) = write_key.encrypt_in_place_with_counter(counter, &mut plaintext_buf) { + if plaintext.len() <= INLINE_ENCRYPT_THRESHOLD { + // Copy into reusable enc_buf, encrypt in place + enc_buf.clear(); + enc_buf.extend_from_slice(plaintext); + if let Err(e) = write_key.encrypt_in_place_with_counter(counter, enc_buf) { return Err(EncryptSendError::crypto(anyhow::anyhow!(e.to_string()))); } out_buf.clear(); - if let Err(e) = wacore::framing::encode_frame_into(&plaintext_buf, None, out_buf) { + if let Err(e) = wacore::framing::encode_frame_into(enc_buf, None, out_buf) { return Err(EncryptSendError::framing(e)); } } else { + // Large messages: encrypt on blocking thread (reads plaintext, returns new ciphertext) let write_key = write_key.clone(); - let plaintext_arc = Arc::new(plaintext_buf); - let plaintext_arc_for_task = plaintext_arc.clone(); + let plaintext_owned = plaintext.to_vec(); let encrypt_result = wacore::runtime::blocking(&**runtime, move || { - write_key.encrypt_with_counter(counter, &plaintext_arc_for_task[..]) + write_key.encrypt_with_counter(counter, &plaintext_owned) }) .await; - drop(plaintext_arc); - let ciphertext = match encrypt_result { Ok(c) => c, Err(e) => { @@ -140,7 +145,6 @@ impl NoiseSocket { } } - // Swap out the framed data; the Vec retains capacity after drain let frame = bytes::Bytes::from(std::mem::take(out_buf)); if let Err(e) = transport.send(frame).await { return Err(EncryptSendError::transport(e)); @@ -151,11 +155,11 @@ impl NoiseSocket { Ok(()) } - pub async fn encrypt_and_send(&self, plaintext_buf: Vec) -> SendResult { + pub async fn encrypt_and_send(&self, plaintext: bytes::Bytes) -> SendResult { let (response_tx, response_rx) = oneshot::channel(); let job = SendJob { - plaintext_buf, + plaintext, response_tx, }; @@ -205,9 +209,7 @@ mod tests { read_key, ); - let plaintext_buf = Vec::with_capacity(1024); - - let result = socket.encrypt_and_send(plaintext_buf).await; + let result = socket.encrypt_and_send(bytes::Bytes::new()).await; assert!(result.is_ok(), "encrypt_and_send should succeed"); } @@ -280,7 +282,7 @@ mod tests { // Use index as the first byte of plaintext to identify this send let mut plaintext = vec![i as u8]; plaintext.extend_from_slice(&[0u8; 99]); - socket.encrypt_and_send(plaintext).await + socket.encrypt_and_send(bytes::Bytes::from(plaintext)).await })); } @@ -340,7 +342,9 @@ mod tests { for size in test_sizes { let plaintext = vec![0xABu8; size]; - let result = socket.encrypt_and_send(plaintext.clone()).await; + let result = socket + .encrypt_and_send(bytes::Bytes::from(plaintext.clone())) + .await; assert!( result.is_ok(), @@ -391,11 +395,11 @@ mod tests { ); // Test empty payload - let result = socket.encrypt_and_send(vec![]).await; + let result = socket.encrypt_and_send(bytes::Bytes::new()).await; assert!(result.is_ok(), "Empty payload should encrypt successfully"); // Test payload at inline threshold boundary (16KB) - let at_threshold = vec![0u8; 16 * 1024]; + let at_threshold = bytes::Bytes::from(vec![0u8; 16 * 1024]); let result = socket.encrypt_and_send(at_threshold).await; assert!( result.is_ok(), @@ -403,7 +407,7 @@ mod tests { ); // Test payload just above inline threshold - let above_threshold = vec![0u8; 16 * 1024 + 1]; + let above_threshold = bytes::Bytes::from(vec![0u8; 16 * 1024 + 1]); let result = socket.encrypt_and_send(above_threshold).await; assert!( result.is_ok(), From 140a487491e17338f8987d817743a981cdd9437e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Wed, 15 Apr 2026 12:05:52 -0300 Subject: [PATCH 08/17] perf: zero-copy digest validation, smarter marshal sizing, smaller handler future Three optimizations targeting remaining allocation hotspots: 1. Zero-copy prekey digest validation: extract_prekey_public_key() reads the publicKey field directly from stored protobuf bytes without full prost decode. Eliminates all PreKeyRecordStructure decode allocs in validate_digest_key (was ~4,872 blocks per session). 2. Smarter marshal_auto: detect large nested child lists (e.g. -> -> 812 prekey nodes) to trigger capacity pre-estimation. The prekey upload IQ previously started at 256B and grew repeatedly to ~40KB; now it pre-allocates near the correct size. 3. Smaller MessageHandler future: replace get_with_by_ref (which captures the entire worker-spawning closure in the async state machine) with get + conditional insert. The handler future no longer carries the large init closure across the .await on cache hits (~29KB -> smaller). Measured impact: - connect_to_ready allocs: 15,597 -> 12,316 (-3,281 allocs, -21.0%) - connect_to_ready bytes: 4.03 MB -> 3.87 MB (-163 KB) - DHAT total: 30.4 MB -> 29.3 MB (-1.1 MB, -3.6%) - DHAT blocks: 123,528 -> 103,097 (-20,431 blocks, -16.5%) --- src/handlers/message.rs | 89 +++++++++++++++++++----------------- src/prekeys.rs | 24 +++------- wacore/binary/src/marshal.rs | 10 +++- wacore/src/prekeys.rs | 53 ++++++++++++++++++++- 4 files changed, 113 insertions(+), 63 deletions(-) diff --git a/src/handlers/message.rs b/src/handlers/message.rs index cb584c85d..fabd0a53c 100644 --- a/src/handlers/message.rs +++ b/src/handlers/message.rs @@ -38,52 +38,57 @@ impl StanzaHandler for MessageHandler { } }; - // Single cache lookup: get or create the lane (lock + queue + worker). - let lane = client - .chat_lanes - .get_with_by_ref(&chat_jid, async { - let (tx, rx) = async_channel::unbounded::>(); + // Fast path: cache hit avoids building the init future entirely. + // This shrinks the handler's async state machine because the large + // worker-spawning closure isn't captured across the .await. + let lane = if let Some(existing) = client.chat_lanes.get(&chat_jid).await { + existing + } else { + let (tx, rx) = async_channel::unbounded::>(); - let client_for_worker = client.clone(); - let spawn_generation = client - .connection_generation - .load(std::sync::atomic::Ordering::Acquire); + let client_for_worker = client.clone(); + let spawn_generation = client + .connection_generation + .load(std::sync::atomic::Ordering::Acquire); - client - .runtime - .spawn(Box::pin(async move { - while let Ok(msg_node) = rx.recv().await { - if client_for_worker - .connection_generation - .load(std::sync::atomic::Ordering::Acquire) - != spawn_generation - { - log::debug!(target: "MessageQueue", "Stale worker exiting; remaining messages will be redelivered by server"); - break; - } - let start = wacore::time::now_millis() as u64; - let client = client_for_worker.clone(); - Box::pin(client.handle_incoming_message(msg_node)).await; - let elapsed = - (wacore::time::now_millis() as u64).saturating_sub(start); - if elapsed > MAX_MESSAGE_DELAY_MS { - warn!( - target: "MessageQueue", - "Message processing took {:.1}s (MAX_MESSAGE_DELAY is {}s)", - elapsed as f64 / 1000.0, - MAX_MESSAGE_DELAY_MS / 1000 - ); - } + client + .runtime + .spawn(Box::pin(async move { + while let Ok(msg_node) = rx.recv().await { + if client_for_worker + .connection_generation + .load(std::sync::atomic::Ordering::Acquire) + != spawn_generation + { + log::debug!(target: "MessageQueue", "Stale worker exiting; remaining messages will be redelivered by server"); + break; } - })) - .detach(); + let start = wacore::time::now_millis() as u64; + let client = client_for_worker.clone(); + Box::pin(client.handle_incoming_message(msg_node)).await; + let elapsed = (wacore::time::now_millis() as u64).saturating_sub(start); + if elapsed > MAX_MESSAGE_DELAY_MS { + warn!( + target: "MessageQueue", + "Message processing took {:.1}s (MAX_MESSAGE_DELAY is {}s)", + elapsed as f64 / 1000.0, + MAX_MESSAGE_DELAY_MS / 1000 + ); + } + } + })) + .detach(); - ChatLane { - enqueue_lock: Arc::new(async_lock::Mutex::new(())), - queue_tx: tx, - } - }) - .await; + let lane = ChatLane { + enqueue_lock: Arc::new(async_lock::Mutex::new(())), + queue_tx: tx, + }; + client + .chat_lanes + .insert(chat_jid.clone(), lane.clone()) + .await; + lane + }; // Lock serializes enqueue order for this chat let _guard = lane.enqueue_lock.lock().await; diff --git a/src/prekeys.rs b/src/prekeys.rs index 809d71492..5c5480e01 100644 --- a/src/prekeys.rs +++ b/src/prekeys.rs @@ -364,37 +364,25 @@ impl Client { return Ok(()); } + // Extract public keys directly from stored protobuf bytes without full decode let mut prekey_pubkeys = Vec::with_capacity(response.prekey_ids.len()); for prekey_id in &response.prekey_ids { let Some(record_bytes) = loaded_map.get(prekey_id) else { log::warn!("digestKey: missing local prekey {}, skipping", prekey_id); return Ok(()); }; - use prost::Message; - match waproto::whatsapp::PreKeyRecordStructure::decode(record_bytes.as_ref()) { - Ok(record) => { - if let Some(pk) = record.public_key { - prekey_pubkeys.push(pk); - } else { - log::warn!( - "digestKey: prekey {} has no public key, skipping", - prekey_id - ); - return Ok(()); - } - } - Err(e) => { + match wacore::prekeys::extract_prekey_public_key(record_bytes) { + Some(pk) => prekey_pubkeys.push(pk), + None => { log::warn!( - "digestKey: failed to decode prekey {}: {}, skipping", - prekey_id, - e + "digestKey: prekey {} has no public key, skipping", + prekey_id ); return Ok(()); } } } - // Compute local SHA-1 digest matching WA Web's validateLocalKeyBundle let local_hash = wacore::prekeys::compute_key_bundle_digest( identity_bytes, skey_pub_bytes, diff --git a/wacore/binary/src/marshal.rs b/wacore/binary/src/marshal.rs index 26f9ba03d..202a43b7f 100644 --- a/wacore/binary/src/marshal.rs +++ b/wacore/binary/src/marshal.rs @@ -144,7 +144,15 @@ fn should_auto_reserve_node(node: &Node) -> bool { match &node.content { Some(NodeContent::Bytes(bytes)) => bytes.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, Some(NodeContent::String(text)) => text.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, - Some(NodeContent::Nodes(children)) => children.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD, + Some(NodeContent::Nodes(children)) => { + if children.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD { + return true; + } + // Check one level deeper for large nested lists (e.g., -> -> 812 keys) + children.iter().any(|child| { + matches!(&child.content, Some(NodeContent::Nodes(gc)) if gc.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD) + }) + } None => false, } } diff --git a/wacore/src/prekeys.rs b/wacore/src/prekeys.rs index 2112c628d..59feade4a 100644 --- a/wacore/src/prekeys.rs +++ b/wacore/src/prekeys.rs @@ -10,12 +10,12 @@ pub struct PreKeyUtils; /// Compute SHA-1 digest of a key bundle for validation against server. /// /// Matches WA Web's `validateLocalKeyBundle` hash computation: -/// SHA-1(identity_pub_key || signed_prekey_pub || signed_prekey_signature || prekey_pub_1 || prekey_pub_2 || ...) +/// SHA-1(identity_pub_key || signed_prekey_pub || signed_prekey_signature || prekey_pub_1 || ...) pub fn compute_key_bundle_digest( identity_pub_key: &[u8], signed_prekey_pub: &[u8], signed_prekey_signature: &[u8], - prekey_pubkeys: &[Vec], + prekey_pubkeys: &[&[u8]], ) -> Vec { use sha1::Digest; let mut hasher = sha1::Sha1::new(); @@ -28,6 +28,55 @@ pub fn compute_key_bundle_digest( hasher.finalize().to_vec() } +/// Extract the `publicKey` field (tag 2) from a protobuf-encoded PreKeyRecordStructure +/// without full prost decode. Returns None if the field is missing. +pub fn extract_prekey_public_key(record: &[u8]) -> Option<&[u8]> { + let mut pos = 0; + while pos < record.len() { + let (tag_byte, consumed) = decode_varint(&record[pos..])?; + pos += consumed; + let field_number = (tag_byte >> 3) as u32; + let wire_type = (tag_byte & 0x7) as u32; + match wire_type { + // varint + 0 => { + let (_, c) = decode_varint(&record[pos..])?; + pos += c; + } + // length-delimited + 2 => { + let (len, c) = decode_varint(&record[pos..])?; + pos += c; + let len = len as usize; + if pos + len > record.len() { + return None; + } + if field_number == 2 { + return Some(&record[pos..pos + len]); + } + pos += len; + } + // fixed64 + 1 => pos += 8, + // fixed32 + 5 => pos += 4, + _ => return None, + } + } + None +} + +fn decode_varint(buf: &[u8]) -> Option<(u64, usize)> { + let mut result: u64 = 0; + for (i, &byte) in buf.iter().enumerate().take(10) { + result |= ((byte & 0x7F) as u64) << (i * 7); + if byte & 0x80 == 0 { + return Some((result, i + 1)); + } + } + None +} + impl PreKeyUtils { pub fn build_fetch_prekeys_request(jids: &[Jid], reason: Option<&str>) -> Node { let user_nodes = jids.iter().map(|jid| { From e801322b2f866560b4973db22dfa2765097b120d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Wed, 15 Apr 2026 12:31:03 -0300 Subject: [PATCH 09/17] refactor: split handle_incoming_message into classify + process phases Split the large monolithic handle_incoming_message async function into two phases: 1. classify_incoming_message: borrows the node tree, extracts owned EncPayload structs (ciphertext Bytes + enc_type + padding_version), returns ClassifiedMessage with all owned data. 2. process_classified_message: holds no node borrows across the heavy decrypt .await points. The async state machine is smaller because it only carries owned data (Arc, Vec, Jid) instead of Vec<&NodeRef> + the entire node tree. Also changes process_session_enc_batch and process_group_enc_batch to accept &[EncPayload] instead of &[&NodeRef], completing the decoupling of decrypt phase from node tree lifetime. The benefit scales with message volume -- each spawned message worker task has a smaller future, reducing per-message heap allocation. --- src/message.rs | 219 ++++++++++++++++++++++++++++++------------------- 1 file changed, 136 insertions(+), 83 deletions(-) diff --git a/src/message.rs b/src/message.rs index 6eea83e18..aff79bfa7 100644 --- a/src/message.rs +++ b/src/message.rs @@ -26,6 +26,40 @@ use waproto::whatsapp::{self as wa}; /// After this many retries, we stop sending retry receipts and rely solely on PDO. const MAX_DECRYPT_RETRIES: u8 = 5; +/// Pre-extracted enc node payload. Holds owned copies of the fields needed for +/// decryption so the async decrypt phase doesn't borrow the original NodeRef tree. +/// This shrinks the async state machine for handle_incoming_message significantly. +pub(crate) struct EncPayload { + pub ciphertext: bytes::Bytes, + pub enc_type: String, + pub padding_version: u8, +} + +impl EncPayload { + /// Extract payload from a NodeRef (used in tests and classify_incoming_message). + pub(crate) fn from_node_ref(node: &NodeRef<'_>) -> Option { + let ciphertext = bytes::Bytes::copy_from_slice(node.content_bytes()?); + let enc_type = node.attrs().optional_string("type")?.to_string(); + let padding_version = node.attrs().optional_u64("v").unwrap_or(2) as u8; + Some(Self { + ciphertext, + enc_type, + padding_version, + }) + } +} + +/// Parsed and classified message ready for decryption. All data is owned -- +/// the original node tree is no longer borrowed. +pub(crate) struct ClassifiedMessage { + pub info: Arc, + pub sender_encryption_jid: Jid, + pub session_payloads: Vec, + pub group_payloads: Vec, + pub max_sender_retry_count: u8, + pub decrypt_fail_mode: crate::types::events::DecryptFailMode, +} + /// Retry count threshold for logging high retry warnings. /// WhatsApp Web logs metrics when retry count exceeds this value. const HIGH_RETRY_COUNT_THRESHOLD: u8 = 3; @@ -262,6 +296,22 @@ impl Client { } pub(crate) async fn handle_incoming_message(self: Arc, node: Arc) { + // Phase 1: classify borrows the node tree, extracts owned payloads, returns quickly. + // Phase 2: process_classified_message holds no node borrows across heavy .await points, + // keeping the async state machine small. + let classified = match self.classify_incoming_message(&node).await { + Some(c) => c, + None => return, + }; + // node is no longer borrowed here -- drop it before the heavy phase + drop(node); + self.process_classified_message(classified).await; + } + + async fn classify_incoming_message( + self: &Arc, + node: &OwnedNodeRef, + ) -> Option { let nr = node.get(); let info = match self.parse_message_info(nr).await { Ok(info) => Arc::new(info), @@ -269,14 +319,14 @@ impl Client { let id = nr.get_attr("id").map(|v| v.as_str()); let from = nr.get_attr("from").map(|v| v.as_str()); log::warn!("Failed to parse message info (id={id:?}, from={from:?}): {e:?}"); - return; + return None; } }; // Newsletters use instead of <enc> because they are not E2E encrypted. if info.source.chat.is_newsletter() { self.handle_newsletter_message(nr, &info).await; - return; + return None; } // Warm LID-PN cache before resolution so resolve_encryption_jid() finds the mapping @@ -313,7 +363,7 @@ impl Client { info.id, nr.tag ); - return; + return None; } if let Some(unavailable) = unavailable_node { @@ -337,11 +387,11 @@ impl Client { decrypt_fail_mode: crate::types::events::DecryptFailMode::Show, }, )); - return; + return None; } - let mut session_enc_nodes = Vec::with_capacity(all_enc_nodes.len()); - let mut group_content_enc_nodes = Vec::with_capacity(all_enc_nodes.len()); + let mut session_payloads = Vec::with_capacity(all_enc_nodes.len()); + let mut group_payloads = Vec::with_capacity(all_enc_nodes.len()); let mut max_sender_retry_count: u8 = 0; let mut has_hide_fail = false; @@ -402,19 +452,31 @@ impl Client { continue; } - // Fall back to built-in handlers + // Extract owned payload so the node doesn't need to be borrowed later + let ct = match enc_node.content_bytes() { + Some(b) => bytes::Bytes::copy_from_slice(b), + None => { + log::warn!("Enc node has no byte content"); + continue; + } + }; + let pv = enc_node.attrs().optional_u64("v").unwrap_or(2) as u8; + let payload = EncPayload { + ciphertext: ct, + enc_type: enc_type.to_string(), + padding_version: pv, + }; + match EncType::from_wire(enc_type.as_ref()) { - Some(et) if et.is_session() => session_enc_nodes.push(*enc_node), - Some(EncType::SenderKey) => group_content_enc_nodes.push(*enc_node), + Some(et) if et.is_session() => session_payloads.push(payload), + Some(EncType::SenderKey) => group_payloads.push(payload), _ => log::warn!("Unknown enc type: {enc_type}"), } } // WA Web diagnostic: validate skmsg is not first in multi-enc messages. - // If skmsg comes first, the SKDM (carried in pkmsg/msg) hasn't been processed yet, - // so the skmsg decryption would fail with NoSenderKey. - if !session_enc_nodes.is_empty() - && !group_content_enc_nodes.is_empty() + if !session_payloads.is_empty() + && !group_payloads.is_empty() && all_enc_nodes.first().is_some_and(|n| { n.get_attr("type") .map(|v| v.as_str()) @@ -429,16 +491,31 @@ impl Client { ); } - // Determine decrypt fail mode from enc nodes (WA Web: hideFail) - let decrypt_fail_mode = if has_hide_fail { - crate::types::events::DecryptFailMode::Hide - } else { - crate::types::events::DecryptFailMode::Show - }; + Some(ClassifiedMessage { + info, + sender_encryption_jid, + session_payloads, + group_payloads, + max_sender_retry_count, + decrypt_fail_mode: if has_hide_fail { + crate::types::events::DecryptFailMode::Hide + } else { + crate::types::events::DecryptFailMode::Show + }, + }) + } + + /// Phase 2: acquire permit, decrypt payloads, flush. No node borrows. + async fn process_classified_message(self: Arc<Self>, msg: ClassifiedMessage) { + let ClassifiedMessage { + info, + sender_encryption_jid, + session_payloads, + group_payloads, + max_sender_retry_count, + decrypt_fail_mode, + } = msg; - // Pre-seed retry cache with sender's retry count to avoid redundant retries. - // Uses max(existing, incoming) so redeliveries with higher counts update the cache, - // but lower counts don't reset our local counter. if max_sender_retry_count > 0 { let cache_key = self .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) @@ -449,11 +526,6 @@ impl Client { .insert(cache_key, max_sender_retry_count) .await; } - log::debug!( - "[msg:{}] Sender retry count {} pre-seeded into cache", - info.id, - max_sender_retry_count - ); } // Acquire global processing permit (1 during offline sync, N after). @@ -486,7 +558,7 @@ impl Client { log::debug!( "Starting PASS 1: Processing {} session establishment messages (pkmsg/msg)", - session_enc_nodes.len() + session_payloads.len() ); // Skip session processing for group/broadcast JIDs — they use sender keys, not 1:1 sessions. @@ -498,20 +570,20 @@ impl Client { session_decrypted_successfully, session_had_duplicates, session_dispatched_undecryptable, - ) = if !is_group_sender && !session_enc_nodes.is_empty() { + ) = if !is_group_sender && !session_payloads.is_empty() { self.clone() .process_session_enc_batch( - &session_enc_nodes, + &session_payloads, &info, &sender_encryption_jid, decrypt_fail_mode, ) .await } else { - if is_group_sender && !session_enc_nodes.is_empty() { + if is_group_sender && !session_payloads.is_empty() { log::debug!( "Skipping {} session messages from group sender {}", - session_enc_nodes.len(), + session_payloads.len(), sender_encryption_jid ); } @@ -520,7 +592,7 @@ impl Client { log::debug!( "Starting PASS 2: Processing {} group content messages (skmsg)", - group_content_enc_nodes.len() + group_payloads.len() ); // Only process group content if: @@ -532,8 +604,8 @@ impl Client { // the SKDM it carried is lost, so skmsg will always fail with NoSenderKey — skip it // to avoid unnecessary retry receipts. The retry for the pkmsg will cause the sender // to resend the entire message including SKDM. - if !group_content_enc_nodes.is_empty() { - let should_process_skmsg = session_enc_nodes.is_empty() + if !group_payloads.is_empty() { + let should_process_skmsg = session_payloads.is_empty() || session_decrypted_successfully || session_had_duplicates; @@ -541,7 +613,7 @@ impl Client { match self .clone() .process_group_enc_batch( - &group_content_enc_nodes, + &group_payloads, &info, &sender_encryption_jid, decrypt_fail_mode, @@ -590,7 +662,7 @@ impl Client { } } else if !session_decrypted_successfully && !session_had_duplicates - && !session_enc_nodes.is_empty() + && !session_payloads.is_empty() { // Edge case: message with only msg/pkmsg that failed to decrypt, no skmsg warn!( @@ -609,16 +681,15 @@ impl Client { .await; } - async fn process_session_enc_batch<'n>( + async fn process_session_enc_batch( self: Arc<Self>, - enc_nodes: &[&'n NodeRef<'n>], + payloads: &[EncPayload], info: &Arc<MessageInfo>, sender_encryption_jid: &Jid, decrypt_fail_mode: crate::types::events::DecryptFailMode, ) -> (bool, bool, bool) { - // Returns (any_success, any_duplicate, dispatched_undecryptable) use wacore::libsignal::protocol::CiphertextMessage; - if enc_nodes.is_empty() { + if payloads.is_empty() { return (false, false, false); } @@ -637,23 +708,11 @@ impl Client { let mut any_duplicate = false; let mut dispatched_undecryptable = false; - for enc_node in enc_nodes { - let ciphertext: &[u8] = match enc_node.content_bytes() { - Some(b) => b, - None => { - log::warn!("Enc node has no byte content (batch session)"); - continue; - } - }; - let enc_type = match enc_node.attrs().optional_string("type") { - Some(t) => t, - None => { - log::warn!("Enc node missing 'type' attribute (batch session)"); - continue; - } - }; - let padding_version = enc_node.attrs().optional_u64("v").unwrap_or(2) as u8; - let enc_type_enum = EncType::from_wire(enc_type.as_ref()); + for payload in payloads { + let ciphertext = &payload.ciphertext[..]; + let enc_type = &payload.enc_type; + let padding_version = payload.padding_version; + let enc_type_enum = EncType::from_wire(enc_type); let parsed_message = if enc_type_enum == Some(EncType::PreKeyMessage) { match PreKeySignalMessage::try_from(ciphertext) { @@ -1015,29 +1074,21 @@ impl Client { (any_success, any_duplicate, dispatched_undecryptable) } - async fn process_group_enc_batch<'n>( + async fn process_group_enc_batch( self: Arc<Self>, - enc_nodes: &[&'n NodeRef<'n>], + payloads: &[EncPayload], info: &Arc<MessageInfo>, _sender_encryption_jid: &Jid, decrypt_fail_mode: crate::types::events::DecryptFailMode, ) -> Result<(), DecryptionError> { - if enc_nodes.is_empty() { + if payloads.is_empty() { return Ok(()); } - // Use the signal cache adapter for group decryption so sender keys are read/written - // through the cache, keeping it consistent with SKDM processing. let mut adapter = self.signal_adapter().await; - for enc_node in enc_nodes { - let ciphertext: &[u8] = match enc_node.content_bytes() { - Some(b) => b, - None => { - log::warn!("Enc node has no byte content (batch group)"); - continue; - } - }; - let padding_version = enc_node.attrs().optional_u64("v").unwrap_or(2) as u8; + for payload in payloads { + let ciphertext = &payload.ciphertext[..]; + let padding_version = payload.padding_version; // Always use bare sender for sender key operations. Real WA delivers // skmsg with bare participant but pkmsg (SKDM) with device-qualified @@ -1725,12 +1776,12 @@ mod tests { .bytes(signal_message.serialized().to_vec()) .build(); let enc_node_ref = enc_node.as_node_ref(); - let enc_nodes = vec![&enc_node_ref]; + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; // With SessionNotFound, should return (false, false, true) - no success, no dupe, dispatched event let (success, had_duplicates, dispatched) = client .process_session_enc_batch( - &enc_nodes, + &payloads, &info, &sender_jid, crate::types::events::DecryptFailMode::Show, @@ -1814,12 +1865,12 @@ mod tests { .bytes(signal_message.serialized().to_vec()) .build(); let enc_node_ref = enc_node.as_node_ref(); - let enc_nodes = vec![&enc_node_ref]; + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; let (success, had_duplicates, dispatched) = client .clone() .process_session_enc_batch( - &enc_nodes, + &payloads, &info, &sender_jid, crate::types::events::DecryptFailMode::Show, @@ -2765,13 +2816,13 @@ mod tests { .build(); let enc_node_ref = enc_node.as_node_ref(); - let enc_nodes = vec![&enc_node_ref]; + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; // Call process_session_enc_batch // This should handle any errors gracefully without panicking let (success, _had_duplicates, _dispatched) = client .process_session_enc_batch( - &enc_nodes, + &payloads, &info, &sender_jid, crate::types::events::DecryptFailMode::Show, @@ -2851,14 +2902,16 @@ mod tests { log::info!("Test: Created batch of 2 messages with invalid data"); - let enc_node_refs_owned: Vec<_> = enc_nodes.iter().map(|n| n.as_node_ref()).collect(); - let enc_node_refs: Vec<&NodeRef<'_>> = enc_node_refs_owned.iter().collect(); + let payloads: Vec<EncPayload> = enc_nodes + .iter() + .filter_map(|n| EncPayload::from_node_ref(&n.as_node_ref())) + .collect(); // Process the batch // Should handle all errors gracefully without stopping at first error let (success, _had_duplicates, _dispatched) = client .process_session_enc_batch( - &enc_node_refs, + &payloads, &info, &sender_jid, crate::types::events::DecryptFailMode::Show, @@ -2924,13 +2977,13 @@ mod tests { .build(); let enc_node_ref = enc_node.as_node_ref(); - let enc_nodes = vec![&enc_node_ref]; + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; // Process the message // Should handle errors gracefully in group context let (success, _had_duplicates, _dispatched) = client .process_session_enc_batch( - &enc_nodes, + &payloads, &info, &sender_phone, crate::types::events::DecryptFailMode::Show, From 79a282cc1060b7abec564b13faa42c04280d913a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <jlucaso@hotmail.com> Date: Wed, 15 Apr 2026 12:39:14 -0300 Subject: [PATCH 10/17] perf: zero-copy EncPayload extraction via OwnedNodeRef::slice_bytes Add OwnedNodeRef::slice_bytes() which returns a Bytes sub-view into the node's backing buffer using pointer arithmetic -- no memcpy. EncPayload::from_owned_node() uses this to extract ciphertext bytes as a zero-copy Bytes view instead of Bytes::copy_from_slice(). Measured impact vs previous commit: - DHAT total: 29.47 MB -> 29.15 MB (-328 KB, -1.1%) - DHAT peak bytes: 5.68 MB -> 5.46 MB (-221 KB, -3.9%) - connect_to_ready bytes: -58 KB (-1.5%) --- src/message.rs | 30 ++++++++++++++++++++---------- wacore/binary/src/node.rs | 15 +++++++++++++++ 2 files changed, 35 insertions(+), 10 deletions(-) diff --git a/src/message.rs b/src/message.rs index aff79bfa7..7ad88334f 100644 --- a/src/message.rs +++ b/src/message.rs @@ -36,7 +36,22 @@ pub(crate) struct EncPayload { } impl EncPayload { - /// Extract payload from a NodeRef (used in tests and classify_incoming_message). + /// Zero-copy extraction from an OwnedNodeRef. The ciphertext Bytes is a + /// sub-view into the node's backing buffer (no memcpy). + pub(crate) fn from_owned_node(owner: &OwnedNodeRef, enc_node: &NodeRef<'_>) -> Option<Self> { + let raw = enc_node.content_bytes()?; + let ciphertext = owner.slice_bytes(raw); + let enc_type = enc_node.attrs().optional_string("type")?.to_string(); + let padding_version = enc_node.attrs().optional_u64("v").unwrap_or(2) as u8; + Some(Self { + ciphertext, + enc_type, + padding_version, + }) + } + + /// Copying extraction from a NodeRef (used in tests where there's no OwnedNodeRef). + #[cfg(test)] pub(crate) fn from_node_ref(node: &NodeRef<'_>) -> Option<Self> { let ciphertext = bytes::Bytes::copy_from_slice(node.content_bytes()?); let enc_type = node.attrs().optional_string("type")?.to_string(); @@ -452,20 +467,15 @@ impl Client { continue; } - // Extract owned payload so the node doesn't need to be borrowed later - let ct = match enc_node.content_bytes() { - Some(b) => bytes::Bytes::copy_from_slice(b), + // Zero-copy: slice_bytes returns a Bytes sub-view into the + // node's backing buffer without memcpy + let payload = match EncPayload::from_owned_node(node, enc_node) { + Some(p) => p, None => { log::warn!("Enc node has no byte content"); continue; } }; - let pv = enc_node.attrs().optional_u64("v").unwrap_or(2) as u8; - let payload = EncPayload { - ciphertext: ct, - enc_type: enc_type.to_string(), - padding_version: pv, - }; match EncType::from_wire(enc_type.as_ref()) { Some(et) if et.is_session() => session_payloads.push(payload), diff --git a/wacore/binary/src/node.rs b/wacore/binary/src/node.rs index 52fbccb38..e95e5b5b8 100644 --- a/wacore/binary/src/node.rs +++ b/wacore/binary/src/node.rs @@ -802,6 +802,21 @@ impl OwnedNodeRef { self.inner.get().to_owned() } + /// Return a zero-copy `Bytes` sub-view for a slice that borrows from this + /// node's backing buffer. Panics if `slice` does not point within the buffer. + pub fn slice_bytes(&self, slice: &[u8]) -> Bytes { + let cart = &self.inner.backing_cart().0; + let base = cart.as_ptr() as usize; + let end = base + cart.len(); + let ptr = slice.as_ptr() as usize; + assert!( + ptr >= base && ptr + slice.len() <= end, + "slice is not within the backing buffer" + ); + let offset = ptr - base; + cart.slice(offset..offset + slice.len()) + } + /// The tag name of this node. #[inline] pub fn tag(&self) -> &str { From 9e1adbc661ffec0c4cc88771110a3a517ae0b4f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <jlucaso@hotmail.com> Date: Wed, 15 Apr 2026 12:47:28 -0300 Subject: [PATCH 11/17] fix: clippy --all --tests errors and warnings - e2e digest_key test: use extract_prekey_public_key() instead of full prost decode, fix Bytes::as_slice() -> as_ref(), adapt to new compute_key_bundle_digest(&[&[u8]]) signature - message.rs: remove 5 needless_borrow warnings (&enc_type -> enc_type) --- src/message.rs | 10 +++++----- tests/e2e/tests/digest_key.rs | 19 +++++-------------- 2 files changed, 10 insertions(+), 19 deletions(-) diff --git a/src/message.rs b/src/message.rs index 7ad88334f..de9ad46e8 100644 --- a/src/message.rs +++ b/src/message.rs @@ -789,7 +789,7 @@ impl Client { if let Err(e) = self .clone() .handle_decrypted_plaintext( - &enc_type, + enc_type, &padded_plaintext, padding_version, info, @@ -875,7 +875,7 @@ impl Client { if let Err(e) = self .clone() .handle_decrypted_plaintext( - &enc_type, + enc_type, &padded_plaintext, padding_version, info, @@ -913,7 +913,7 @@ impl Client { &parsed_message, &mut adapter, &mut rng, - &enc_type, + enc_type, padding_version, info, ) @@ -975,7 +975,7 @@ impl Client { &parsed_message, &mut adapter, &mut rng, - &enc_type, + enc_type, padding_version, info, ) @@ -1041,7 +1041,7 @@ impl Client { &parsed_message, &mut adapter, &mut rng, - &enc_type, + enc_type, padding_version, info, ) diff --git a/tests/e2e/tests/digest_key.rs b/tests/e2e/tests/digest_key.rs index 595e49106..22ca59a7d 100644 --- a/tests/e2e/tests/digest_key.rs +++ b/tests/e2e/tests/digest_key.rs @@ -68,23 +68,14 @@ async fn test_digest_key_hash_matches_server() -> anyhow::Result<()> { ) }); - use prost::Message; - let record = whatsapp_rust::waproto::whatsapp::PreKeyRecordStructure::decode( - record_bytes.as_slice(), - )?; - let pk = record - .public_key + let pk = wacore::prekeys::extract_prekey_public_key(&record_bytes) .unwrap_or_else(|| panic!("Prekey {} has no public_key field", prekey_id)); - prekey_pubkeys.push(pk); + prekey_pubkeys.push(pk.to_vec()); } - // Compute SHA-1 digest the same way the server does - let local_hash = wacore::prekeys::compute_key_bundle_digest( - identity_pub, - skey_pub, - skey_sig, - &prekey_pubkeys, - ); + let pubkey_refs: Vec<&[u8]> = prekey_pubkeys.iter().map(|v| v.as_slice()).collect(); + let local_hash = + wacore::prekeys::compute_key_bundle_digest(identity_pub, skey_pub, skey_sig, &pubkey_refs); assert_eq!( local_hash, From f47371abecaa5820e805ae27d1bfb1c2fbd7ad15 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <jlucaso@hotmail.com> Date: Wed, 15 Apr 2026 12:49:49 -0300 Subject: [PATCH 12/17] fix: zlib pool BufError spin + scratch buffer retention - Detect truncated/corrupt zlib streams by checking if total_in and total_out made no progress on BufError, and return an error instead of spinning forever growing the output buffer. - Use scratch.clone() instead of mem::swap to return data, so the pooled scratch Vec retains its capacity for subsequent calls. --- wacore/binary/src/zlib_pool.rs | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/wacore/binary/src/zlib_pool.rs b/wacore/binary/src/zlib_pool.rs index 3653adf30..5d4e65208 100644 --- a/wacore/binary/src/zlib_pool.rs +++ b/wacore/binary/src/zlib_pool.rs @@ -26,6 +26,9 @@ pub fn decompress_zlib_pooled(compressed: &[u8], max_size: u64) -> io::Result<Ve let mut input_offset = 0; loop { + let prev_in = decompressor.total_in(); + let prev_out = decompressor.total_out(); + let status = decompressor .decompress_vec( &compressed[input_offset..], @@ -45,17 +48,23 @@ pub fn decompress_zlib_pooled(compressed: &[u8], max_size: u64) -> io::Result<Ve match status { Status::StreamEnd => break, - Status::Ok | Status::BufError => { - // Need more output space + Status::Ok => { + scratch.reserve(scratch.capacity().max(4096)); + } + Status::BufError => { + // BufError with no progress means the stream is truncated/corrupt + if decompressor.total_in() == prev_in && decompressor.total_out() == prev_out { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "zlib stream truncated (no progress)", + )); + } scratch.reserve(scratch.capacity().max(4096)); } } } - // Return the data by swapping out to avoid cloning. - // The scratch buffer keeps its capacity for the next call. - let mut result = Vec::new(); - std::mem::swap(scratch, &mut result); - Ok(result) + // Clone data out so scratch retains its capacity for the next call + Ok(scratch.clone()) }) } From 1563f49ba77583a1ad6fbad57c2acbc65b641490 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <jlucaso@hotmail.com> Date: Wed, 15 Apr 2026 12:59:26 -0300 Subject: [PATCH 13/17] fix: address review findings - race condition, buffer reuse, EncType enum - Restore get_with_by_ref for chat lane creation to fix TOCTOU race that could create duplicate workers for the same chat. Extract init logic to create_chat_lane() helper to keep the future small. - Fix out_buf reuse in NoiseSocket: use Bytes::copy_from_slice + clear instead of mem::take so the Vec retains capacity between sends. - Change EncPayload.enc_type from String to EncType enum, eliminating one heap allocation per enc node and redundant from_wire() parsing. - DRY EncPayload constructors via shared from_parts(). - Restore removed debug log for retry count pre-seeding. --- src/handlers/message.rs | 101 ++++++++++++++++++------------------- src/message.rs | 56 ++++++++++---------- src/socket/noise_socket.rs | 4 +- 3 files changed, 82 insertions(+), 79 deletions(-) diff --git a/src/handlers/message.rs b/src/handlers/message.rs index fabd0a53c..5d3854559 100644 --- a/src/handlers/message.rs +++ b/src/handlers/message.rs @@ -38,57 +38,12 @@ impl StanzaHandler for MessageHandler { } }; - // Fast path: cache hit avoids building the init future entirely. - // This shrinks the handler's async state machine because the large - // worker-spawning closure isn't captured across the .await. - let lane = if let Some(existing) = client.chat_lanes.get(&chat_jid).await { - existing - } else { - let (tx, rx) = async_channel::unbounded::<Arc<wacore_binary::OwnedNodeRef>>(); - - let client_for_worker = client.clone(); - let spawn_generation = client - .connection_generation - .load(std::sync::atomic::Ordering::Acquire); - - client - .runtime - .spawn(Box::pin(async move { - while let Ok(msg_node) = rx.recv().await { - if client_for_worker - .connection_generation - .load(std::sync::atomic::Ordering::Acquire) - != spawn_generation - { - log::debug!(target: "MessageQueue", "Stale worker exiting; remaining messages will be redelivered by server"); - break; - } - let start = wacore::time::now_millis() as u64; - let client = client_for_worker.clone(); - Box::pin(client.handle_incoming_message(msg_node)).await; - let elapsed = (wacore::time::now_millis() as u64).saturating_sub(start); - if elapsed > MAX_MESSAGE_DELAY_MS { - warn!( - target: "MessageQueue", - "Message processing took {:.1}s (MAX_MESSAGE_DELAY is {}s)", - elapsed as f64 / 1000.0, - MAX_MESSAGE_DELAY_MS / 1000 - ); - } - } - })) - .detach(); - - let lane = ChatLane { - enqueue_lock: Arc::new(async_lock::Mutex::new(())), - queue_tx: tx, - }; - client - .chat_lanes - .insert(chat_jid.clone(), lane.clone()) - .await; - lane - }; + // Single-flight: get_with_by_ref guarantees exactly one init runs per key, + // preventing duplicate workers for the same chat (TOCTOU race). + let lane = client + .chat_lanes + .get_with_by_ref(&chat_jid, async { create_chat_lane(&client) }) + .await; // Lock serializes enqueue order for this chat let _guard = lane.enqueue_lock.lock().await; @@ -102,3 +57,47 @@ impl StanzaHandler for MessageHandler { true } } + +/// Construct a ChatLane with a spawned worker task. Extracted to keep the +/// init closure passed to `get_with_by_ref` small. +fn create_chat_lane(client: &Arc<Client>) -> ChatLane { + let (tx, rx) = async_channel::unbounded::<Arc<wacore_binary::OwnedNodeRef>>(); + + let client_for_worker = client.clone(); + let spawn_generation = client + .connection_generation + .load(std::sync::atomic::Ordering::Acquire); + + client + .runtime + .spawn(Box::pin(async move { + while let Ok(msg_node) = rx.recv().await { + if client_for_worker + .connection_generation + .load(std::sync::atomic::Ordering::Acquire) + != spawn_generation + { + log::debug!(target: "MessageQueue", "Stale worker exiting; remaining messages will be redelivered by server"); + break; + } + let start = wacore::time::now_millis() as u64; + let client = client_for_worker.clone(); + Box::pin(client.handle_incoming_message(msg_node)).await; + let elapsed = (wacore::time::now_millis() as u64).saturating_sub(start); + if elapsed > MAX_MESSAGE_DELAY_MS { + warn!( + target: "MessageQueue", + "Message processing took {:.1}s (MAX_MESSAGE_DELAY is {}s)", + elapsed as f64 / 1000.0, + MAX_MESSAGE_DELAY_MS / 1000 + ); + } + } + })) + .detach(); + + ChatLane { + enqueue_lock: Arc::new(async_lock::Mutex::new(())), + queue_tx: tx, + } +} diff --git a/src/message.rs b/src/message.rs index de9ad46e8..4e4657a6f 100644 --- a/src/message.rs +++ b/src/message.rs @@ -28,20 +28,15 @@ const MAX_DECRYPT_RETRIES: u8 = 5; /// Pre-extracted enc node payload. Holds owned copies of the fields needed for /// decryption so the async decrypt phase doesn't borrow the original NodeRef tree. -/// This shrinks the async state machine for handle_incoming_message significantly. pub(crate) struct EncPayload { pub ciphertext: bytes::Bytes, - pub enc_type: String, + pub enc_type: EncType, pub padding_version: u8, } impl EncPayload { - /// Zero-copy extraction from an OwnedNodeRef. The ciphertext Bytes is a - /// sub-view into the node's backing buffer (no memcpy). - pub(crate) fn from_owned_node(owner: &OwnedNodeRef, enc_node: &NodeRef<'_>) -> Option<Self> { - let raw = enc_node.content_bytes()?; - let ciphertext = owner.slice_bytes(raw); - let enc_type = enc_node.attrs().optional_string("type")?.to_string(); + fn from_parts(ciphertext: bytes::Bytes, enc_node: &NodeRef<'_>) -> Option<Self> { + let enc_type = EncType::from_wire(enc_node.attrs().optional_string("type")?.as_ref())?; let padding_version = enc_node.attrs().optional_u64("v").unwrap_or(2) as u8; Some(Self { ciphertext, @@ -50,17 +45,15 @@ impl EncPayload { }) } + /// Zero-copy extraction from an OwnedNodeRef. + pub(crate) fn from_owned_node(owner: &OwnedNodeRef, enc_node: &NodeRef<'_>) -> Option<Self> { + Self::from_parts(owner.slice_bytes(enc_node.content_bytes()?), enc_node) + } + /// Copying extraction from a NodeRef (used in tests where there's no OwnedNodeRef). #[cfg(test)] pub(crate) fn from_node_ref(node: &NodeRef<'_>) -> Option<Self> { - let ciphertext = bytes::Bytes::copy_from_slice(node.content_bytes()?); - let enc_type = node.attrs().optional_string("type")?.to_string(); - let padding_version = node.attrs().optional_u64("v").unwrap_or(2) as u8; - Some(Self { - ciphertext, - enc_type, - padding_version, - }) + Self::from_parts(bytes::Bytes::copy_from_slice(node.content_bytes()?), node) } } @@ -468,19 +461,20 @@ impl Client { } // Zero-copy: slice_bytes returns a Bytes sub-view into the - // node's backing buffer without memcpy + // node's backing buffer without memcpy. + // from_owned_node returns None for unknown enc types or missing content. let payload = match EncPayload::from_owned_node(node, enc_node) { Some(p) => p, None => { - log::warn!("Enc node has no byte content"); + log::warn!("Enc node has no content or unknown type: {enc_type}"); continue; } }; - match EncType::from_wire(enc_type.as_ref()) { - Some(et) if et.is_session() => session_payloads.push(payload), - Some(EncType::SenderKey) => group_payloads.push(payload), - _ => log::warn!("Unknown enc type: {enc_type}"), + if payload.enc_type.is_session() { + session_payloads.push(payload); + } else { + group_payloads.push(payload); } } @@ -536,6 +530,11 @@ impl Client { .insert(cache_key, max_sender_retry_count) .await; } + log::debug!( + "[msg:{}] Sender retry count {} pre-seeded into cache", + info.id, + max_sender_retry_count + ); } // Acquire global processing permit (1 during offline sync, N after). @@ -720,11 +719,11 @@ impl Client { for payload in payloads { let ciphertext = &payload.ciphertext[..]; - let enc_type = &payload.enc_type; + let enc_type = payload.enc_type; + let enc_type_str = enc_type.as_wire_str(); let padding_version = payload.padding_version; - let enc_type_enum = EncType::from_wire(enc_type); - let parsed_message = if enc_type_enum == Some(EncType::PreKeyMessage) { + let parsed_message = if enc_type == EncType::PreKeyMessage { match PreKeySignalMessage::try_from(ciphertext) { Ok(m) => CiphertextMessage::PreKeySignalMessage(m), Err(e) => { @@ -742,7 +741,7 @@ impl Client { } }; - if enc_type_enum == Some(EncType::PreKeyMessage) { + if enc_type == EncType::PreKeyMessage { // FLAGGED FOR DEBUGGING: "Bad Mac" Reproducibility #[cfg(feature = "debug-snapshots")] { @@ -751,7 +750,7 @@ impl Client { "id": info.id, "sender_jid": sender_encryption_jid.to_string(), "timestamp": info.timestamp, - "enc_type": enc_type, + "enc_type": enc_type_str, "payload_base64": BASE64_STANDARD.encode(ciphertext), }); @@ -771,6 +770,9 @@ impl Client { } } + // Shadow with wire string for all downstream usage (logging, handlers) + let enc_type = enc_type_str; + let decrypt_res = message_decrypt( &parsed_message, &signal_address, diff --git a/src/socket/noise_socket.rs b/src/socket/noise_socket.rs index 90b4b9575..48ca7f35d 100644 --- a/src/socket/noise_socket.rs +++ b/src/socket/noise_socket.rs @@ -145,7 +145,9 @@ impl NoiseSocket { } } - let frame = bytes::Bytes::from(std::mem::take(out_buf)); + // copy_from_slice so out_buf retains its capacity for the next send + let frame = bytes::Bytes::copy_from_slice(out_buf); + out_buf.clear(); if let Err(e) = transport.send(frame).await { return Err(EncryptSendError::transport(e)); } From 2ba5f6e72e7f8f2d61165880d17b1d499b4aacba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <jlucaso@hotmail.com> Date: Wed, 15 Apr 2026 13:04:06 -0300 Subject: [PATCH 14/17] fix: address all review findings - zlib pool: enforce max_size cap before decompress_vec to prevent compressed bombs; zero-copy return via mem::take + re-reserve scratch - extract_prekey_public_key: use last-one-wins semantics per protobuf spec; gracefully skip unknown wire types instead of returning None - decode_varint: reject 10th byte with payload > 1 bit (overflow guard) - marshal_ref_auto: mirror nested child list check from marshal_auto so both heuristics stay in sync - CountingAlloc: implement realloc (adjusts bytes, not count) and alloc_zeroed (tracks both) for accurate measurement - bench-integration: precompute format! strings outside measured closure - sqlite store_prekeys_batch: keep Bytes through retry loop, only call as_ref() at diesel bind site - CI workflow: split into PR (read-only) and push (write) jobs with least-privilege permissions --- .github/workflows/bench-integration.yml | 101 +++++++++++++++--- storages/sqlite-storage/src/sqlite_store.rs | 9 +- tests/bench-integration/src/counting_alloc.rs | 14 +++ tests/bench-integration/src/main.rs | 11 +- wacore/binary/src/marshal.rs | 10 +- wacore/binary/src/zlib_pool.rs | 29 ++++- wacore/src/prekeys.rs | 32 ++++-- 7 files changed, 167 insertions(+), 39 deletions(-) diff --git a/.github/workflows/bench-integration.yml b/.github/workflows/bench-integration.yml index d64b8df34..c51263ae8 100644 --- a/.github/workflows/bench-integration.yml +++ b/.github/workflows/bench-integration.yml @@ -6,11 +6,6 @@ on: push: branches: [main] -permissions: - contents: write - pull-requests: write - deployments: write - concurrency: group: bench-integration-${{ github.head_ref || github.run_id }} cancel-in-progress: true @@ -22,9 +17,13 @@ env: RUSTC_WRAPPER: "sccache" jobs: - bench-integration: - name: Integration Benchmark + bench-integration-pr: + if: github.event_name == 'pull_request' + name: Integration Benchmark (PR) runs-on: ubuntu-latest + permissions: + contents: read + pull-requests: write services: mock-server: image: ghcr.io/whiskeysockets-devtools/bartender:latest @@ -90,30 +89,98 @@ jobs: bench_integration.log retention-days: 30 - # Push to main: store baseline - - name: Store baseline (push to main) - if: github.event_name == 'push' + - name: Compare against baseline uses: benchmark-action/github-action-benchmark@v1 with: name: "whatsapp-rust integration benchmarks" tool: "customSmallerIsBetter" output-file-path: bench_integration_results.json github-token: ${{ secrets.GITHUB_TOKEN }} - auto-push: true + auto-push: false + save-data-file: false benchmark-data-dir-path: dev/bench-integration - max-items-in-chart: 100 summary-always: true - # PR: compare against baseline - - name: Compare against baseline (pull request) - if: github.event_name == 'pull_request' + bench-integration-push: + if: github.event_name == 'push' + name: Integration Benchmark (push) + runs-on: ubuntu-latest + permissions: + contents: write + deployments: write + services: + mock-server: + image: ghcr.io/whiskeysockets-devtools/bartender:latest + credentials: + username: ${{ github.actor }} + password: ${{ secrets.BARTENDER_GHCR_TOKEN }} + ports: + - 8080:8080 + env: + CHATSTATE_TTL_SECS: "3" + options: --log-driver none + steps: + - uses: actions/checkout@v6 + + - uses: dtolnay/rust-toolchain@master + with: + toolchain: nightly-2026-04-05 + + - name: Install protoc + uses: taiki-e/install-action@v2 + with: + tool: protoc@${{ env.PROTOC_VERSION }} + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Cache Rust registry + uses: Swatinem/rust-cache@v2 + with: + prefix-key: ${{ runner.os }}-cargo-bench-int + cache-targets: "false" + + - name: Wait for mock server + run: | + for i in $(seq 1 30); do + if curl -sk https://localhost:8080/ > /dev/null 2>&1; then + echo "Mock server is ready" + exit 0 + fi + sleep 1 + done + echo "Mock server failed to become ready" + exit 1 + + - name: Run integration benchmarks + env: + MOCK_SERVER_URL: "wss://127.0.0.1:8080/ws/chat" + RUST_LOG: warn + run: | + cargo run -p bench-integration --release \ + > bench_integration_results.json \ + 2> bench_integration.log + cat bench_integration.log >&2 + echo "Results:" + cat bench_integration_results.json + + - name: Upload benchmark artifacts + uses: actions/upload-artifact@v4 + with: + name: bench-integration-results + path: | + bench_integration_results.json + bench_integration.log + retention-days: 30 + + - name: Store baseline uses: benchmark-action/github-action-benchmark@v1 with: name: "whatsapp-rust integration benchmarks" tool: "customSmallerIsBetter" output-file-path: bench_integration_results.json github-token: ${{ secrets.GITHUB_TOKEN }} - auto-push: false - save-data-file: false + auto-push: true benchmark-data-dir-path: dev/bench-integration + max-items-in-chart: 100 summary-always: true diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index 4748a00ef..66e31a141 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -1288,7 +1288,7 @@ impl SignalStore for SqliteStore { let pool = self.pool.clone(); let db_semaphore = self.db_semaphore.clone(); let device_id = self.device_id; - let keys: Vec<(u32, Vec<u8>)> = keys.iter().map(|(id, b)| (*id, b.to_vec())).collect(); + let keys: Vec<(u32, Bytes)> = keys.to_vec(); const MAX_RETRIES: u32 = 5; @@ -1312,13 +1312,16 @@ impl SignalStore for SqliteStore { diesel::insert_into(prekeys::table) .values(( prekeys::id.eq(*id as i32), - prekeys::key.eq(record), + prekeys::key.eq(record.as_ref()), prekeys::uploaded.eq(uploaded), prekeys::device_id.eq(device_id), )) .on_conflict((prekeys::id, prekeys::device_id)) .do_update() - .set((prekeys::key.eq(record), prekeys::uploaded.eq(uploaded))) + .set(( + prekeys::key.eq(record.as_ref()), + prekeys::uploaded.eq(uploaded), + )) .execute(conn)?; } Ok::<(), diesel::result::Error>(()) diff --git a/tests/bench-integration/src/counting_alloc.rs b/tests/bench-integration/src/counting_alloc.rs index 130049772..74d683e27 100644 --- a/tests/bench-integration/src/counting_alloc.rs +++ b/tests/bench-integration/src/counting_alloc.rs @@ -16,6 +16,20 @@ unsafe impl GlobalAlloc for CountingAlloc { unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { unsafe { System.dealloc(ptr, layout) } } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + // Track the size difference but not as a new allocation + if new_size > layout.size() { + ALLOC_BYTES.fetch_add((new_size - layout.size()) as u64, Ordering::Relaxed); + } + unsafe { System.realloc(ptr, layout, new_size) } + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + ALLOC_COUNT.fetch_add(1, Ordering::Relaxed); + ALLOC_BYTES.fetch_add(layout.size() as u64, Ordering::Relaxed); + unsafe { System.alloc_zeroed(layout) } + } } #[derive(Debug, Clone, Copy)] diff --git a/tests/bench-integration/src/main.rs b/tests/bench-integration/src/main.rs index 82c6c30f2..71481f39f 100644 --- a/tests/bench-integration/src/main.rs +++ b/tests/bench-integration/src/main.rs @@ -181,11 +181,12 @@ async fn bench_send_message(results: &mut BenchResults) -> anyhow::Result<()> { // -- Amortized: send N messages -- const N: u64 = 20; + let send_texts: Vec<String> = (0..N).map(|i| format!("bench-send-{i}")).collect(); let m = measure(async || { - for i in 0..N { + for text in &send_texts { client_a .client - .send_message(jid_b.clone(), text_msg(&format!("bench-send-{i}"))) + .send_message(jid_b.clone(), text_msg(text)) .await?; } Ok(()) @@ -208,10 +209,8 @@ async fn bench_send_message(results: &mut BenchResults) -> anyhow::Result<()> { ); } - for i in 0..N { - client_b - .wait_for_text(&format!("bench-send-{i}"), 30) - .await?; + for text in &send_texts { + client_b.wait_for_text(text, 30).await?; } client_a.disconnect().await; diff --git a/wacore/binary/src/marshal.rs b/wacore/binary/src/marshal.rs index 202a43b7f..1c21b9221 100644 --- a/wacore/binary/src/marshal.rs +++ b/wacore/binary/src/marshal.rs @@ -166,7 +166,15 @@ fn should_auto_reserve_node_ref(node: &NodeRef<'_>) -> bool { match node.content.as_deref() { Some(NodeContentRef::Bytes(bytes)) => bytes.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, Some(NodeContentRef::String(text)) => text.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, - Some(NodeContentRef::Nodes(children)) => children.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD, + Some(NodeContentRef::Nodes(children)) => { + if children.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD { + return true; + } + // Check one level deeper for large nested lists (e.g., <iq> -> <list> -> 812 keys) + children.iter().any(|child| { + matches!(child.content.as_deref(), Some(NodeContentRef::Nodes(gc)) if gc.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD) + }) + } None => false, } } diff --git a/wacore/binary/src/zlib_pool.rs b/wacore/binary/src/zlib_pool.rs index 5d4e65208..ce7b9711e 100644 --- a/wacore/binary/src/zlib_pool.rs +++ b/wacore/binary/src/zlib_pool.rs @@ -24,8 +24,20 @@ pub fn decompress_zlib_pooled(compressed: &[u8], max_size: u64) -> io::Result<Ve scratch.reserve(estimated - scratch.capacity()); } + // Cap output growth to max_size + 1 so we detect oversized payloads + // without allocating unbounded memory from a compressed bomb. + let cap = (max_size as usize).saturating_add(1); + let mut input_offset = 0; loop { + // Enforce cap before decompress_vec can grow the buffer + if scratch.len() >= cap { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("decompressed payload exceeds {max_size} bytes"), + )); + } + let prev_in = decompressor.total_in(); let prev_out = decompressor.total_out(); @@ -49,22 +61,29 @@ pub fn decompress_zlib_pooled(compressed: &[u8], max_size: u64) -> io::Result<Ve match status { Status::StreamEnd => break, Status::Ok => { - scratch.reserve(scratch.capacity().max(4096)); + // Grow but never past the cap + let want = scratch.capacity().max(4096).min(cap - scratch.len()); + scratch.reserve(want); } Status::BufError => { - // BufError with no progress means the stream is truncated/corrupt if decompressor.total_in() == prev_in && decompressor.total_out() == prev_out { return Err(io::Error::new( io::ErrorKind::InvalidData, "zlib stream truncated (no progress)", )); } - scratch.reserve(scratch.capacity().max(4096)); + let want = scratch.capacity().max(4096).min(cap - scratch.len()); + scratch.reserve(want); } } } - // Clone data out so scratch retains its capacity for the next call - Ok(scratch.clone()) + // Move the Vec out (zero-copy), then restore scratch with fresh capacity. + // Callers (unpack_bytes, history_sync) wrap in Bytes::from() which takes + // ownership of the Vec's allocation, so no extra copy occurs. + let result = std::mem::take(scratch); + // Pre-allocate for next call so the first decompress_vec doesn't start at 0 + scratch.reserve(4096); + Ok(result) }) } diff --git a/wacore/src/prekeys.rs b/wacore/src/prekeys.rs index 59feade4a..d03132717 100644 --- a/wacore/src/prekeys.rs +++ b/wacore/src/prekeys.rs @@ -29,9 +29,11 @@ pub fn compute_key_bundle_digest( } /// Extract the `publicKey` field (tag 2) from a protobuf-encoded PreKeyRecordStructure -/// without full prost decode. Returns None if the field is missing. +/// without full prost decode. Uses last-one-wins semantics per protobuf spec. +/// Skips unknown fields gracefully. pub fn extract_prekey_public_key(record: &[u8]) -> Option<&[u8]> { let mut pos = 0; + let mut result: Option<&[u8]> = None; while pos < record.len() { let (tag_byte, consumed) = decode_varint(&record[pos..])?; pos += consumed; @@ -49,26 +51,42 @@ pub fn extract_prekey_public_key(record: &[u8]) -> Option<&[u8]> { pos += c; let len = len as usize; if pos + len > record.len() { - return None; + return result; } if field_number == 2 { - return Some(&record[pos..pos + len]); + result = Some(&record[pos..pos + len]); } pos += len; } // fixed64 - 1 => pos += 8, + 1 => { + if pos + 8 > record.len() { + return result; + } + pos += 8; + } // fixed32 - 5 => pos += 4, - _ => return None, + 5 => { + if pos + 4 > record.len() { + return result; + } + pos += 4; + } + // Unknown wire type -- skip gracefully + _ => return result, } } - None + result } fn decode_varint(buf: &[u8]) -> Option<(u64, usize)> { let mut result: u64 = 0; for (i, &byte) in buf.iter().enumerate().take(10) { + // The 10th byte (i==9) carries the highest bits; only the low bit + // is valid payload (64 - 9*7 = 1). Reject if more bits are set. + if i == 9 && (byte & 0x7F) > 1 { + return None; + } result |= ((byte & 0x7F) as u64) << (i * 7); if byte & 0x80 == 0 { return Some((result, i + 1)); From 8c599aaf28276d6934c63fbc21d7fd0647b5151f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <jlucaso@hotmail.com> Date: Wed, 15 Apr 2026 13:26:55 -0300 Subject: [PATCH 15/17] fix: normalize retry stanza to JID to bare (no device suffix) The retry resend stanza to attribute was using the raw from JID from the retry receipt, which can include a device suffix (e.g. :33) depending on server behavior. The device-specific routing is already handled by the participants/to child node, so the outer to should be the bare user-level chat target. Normalizes original_from to the bare chat JID, matching the normal send path and preventing flaky failures when the server includes device suffixes in receipt JIDs. --- src/retry.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/src/retry.rs b/src/retry.rs index 35e158561..5b73717a9 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -158,9 +158,12 @@ fn resolve_retry_chat_info( }; RetryChatInfo { - chat, + chat: chat.clone(), requester, - original_from: from.clone(), + // Use bare JID for stanza `to` -- device routing is via <participants><to jid=requester>. + // Normalizing here ensures consistent behavior regardless of storage backend or + // whether the server includes a device suffix in the receipt. + original_from: chat, is_bot, } } @@ -2543,17 +2546,18 @@ mod tests { } #[test] - fn resolve_retry_chat_info_preserves_original_from() { + fn resolve_retry_chat_info_normalizes_original_from() { use wacore_binary::builder::NodeBuilder; - // DM with device suffix — original_from should preserve it + // DM with device suffix — original_from should be normalized to bare + // (device routing is handled by <participants><to jid=requester>) let node = NodeBuilder::new("receipt").build(); let receipt = make_test_receipt("5511999999999:33@s.whatsapp.net"); let info = resolve_retry_chat_info(&receipt, &node.as_node_ref(), None, None); - // original_from keeps the full JID including device - assert_eq!(info.original_from.device(), 33); + // original_from is bare (same as chat) + assert_eq!(info.original_from.device(), 0); assert_eq!(info.original_from.user, "5511999999999"); // chat is bare From c1abd3d262a5d8453582aa8730ec98a2feca74c3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <jlucaso@hotmail.com> Date: Wed, 15 Apr 2026 13:32:51 -0300 Subject: [PATCH 16/17] fix: relax retry e2e test to tolerate device-qualified JIDs WA Web uses the raw receipt from JID (with device suffix) as the retry stanza to attribute (confirmed in WAWebHandleRetryRequest.js line 30/97). The mock server may assign non-zero device IDs under CI load, making the to JID device-qualified. Changed assertion to check user part only via is_same_user_as instead of exact string match. --- tests/e2e/tests/retry_dm_multidevice.rs | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/tests/e2e/tests/retry_dm_multidevice.rs b/tests/e2e/tests/retry_dm_multidevice.rs index bcf613d0a..bfdd36571 100644 --- a/tests/e2e/tests/retry_dm_multidevice.rs +++ b/tests/e2e/tests/retry_dm_multidevice.rs @@ -3,6 +3,7 @@ use e2e_tests::{TestClient, send_and_expect_text, text_msg}; use log::info; use wacore::types::events::Event; +use wacore_binary::JidExt as _; use wacore_binary::node::Node; use whatsapp_rust::{NodeFilter, SendOptions}; @@ -107,11 +108,17 @@ async fn test_dm_retry_recovers_after_session_deletion() -> anyhow::Result<()> { Some("1"), "Retry resend should mark the payload with count=1" ); - let jid_b_str = jid_b.to_string(); - assert_eq!( - retry_node.attrs().optional_string("to").as_deref(), - Some(jid_b_str.as_str()), - "Retry resend should keep the user-level chat target" + // WA Web uses the raw receipt `from` (with device suffix) as the stanza `to`. + // The mock server may or may not include a device suffix depending on timing, + // so check user part only. + let retry_to = retry_node + .attrs() + .optional_jid("to") + .expect("Retry resend should have a 'to' attribute"); + assert!( + retry_to.is_same_user_as(&jid_b), + "Retry resend 'to' should target the same user (got {retry_to}, expected user {})", + jid_b ); info!("Retry recovered after B deleted its session with A"); From 36224a8cc2f6380a606cdc3099db6d2879f8cd4c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <jlucaso@hotmail.com> Date: Wed, 15 Apr 2026 13:43:00 -0300 Subject: [PATCH 17/17] fix: restore WA Web-compliant original_from (raw receipt JID with device suffix) Reverts the normalization from 8c599aa. WA Web's HandleRetryRequest.js captures m = e.from (raw, with device suffix) at line 30 BEFORE any normalization, then uses to: m at line 97 for the retry stanza. The e2e test flake was already fixed in c1abd3d by relaxing the assertion to is_same_user_as instead of exact string match. --- src/retry.rs | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/src/retry.rs b/src/retry.rs index 5b73717a9..6f871433b 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -158,12 +158,9 @@ fn resolve_retry_chat_info( }; RetryChatInfo { - chat: chat.clone(), + chat, requester, - // Use bare JID for stanza `to` -- device routing is via <participants><to jid=requester>. - // Normalizing here ensures consistent behavior regardless of storage backend or - // whether the server includes a device suffix in the receipt. - original_from: chat, + original_from: from.clone(), is_bot, } } @@ -2546,18 +2543,18 @@ mod tests { } #[test] - fn resolve_retry_chat_info_normalizes_original_from() { + fn resolve_retry_chat_info_preserves_original_from() { use wacore_binary::builder::NodeBuilder; - // DM with device suffix — original_from should be normalized to bare - // (device routing is handled by <participants><to jid=requester>) + // DM with device suffix — original_from preserves the raw receipt from + // (WA Web: variable m = e.from, used as-is for stanza to) let node = NodeBuilder::new("receipt").build(); let receipt = make_test_receipt("5511999999999:33@s.whatsapp.net"); let info = resolve_retry_chat_info(&receipt, &node.as_node_ref(), None, None); - // original_from is bare (same as chat) - assert_eq!(info.original_from.device(), 0); + // original_from keeps the full JID including device + assert_eq!(info.original_from.device(), 33); assert_eq!(info.original_from.user, "5511999999999"); // chat is bare