perf(send): parallelize group encrypt fan-out + adjacent wins - #610
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
📝 WalkthroughSummary by CodeRabbit
WalkthroughBatches device-list writes with per-record canonicalization and best-effort canonical-flip cleanup; threads a runtime into Signal encryption APIs to enable bounded concurrent per-device crypto tasks; moves identity backend load outside the identities mutex and makes trust-check always succeed. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Client
participant WACore as wacore::encrypt_for_devices
participant Runtime as Runtime (task pool)
participant SessionStore as SessionStore / IdentityStore
participant Backend as ProtocolStore
Client->>WACore: call(runtime, stores, devices, plaintext, meta)
WACore->>WACore: build device index / overrides
loop per-device (bounded concurrency)
WACore->>Runtime: spawn_oneshot(prekey + session setup)
Runtime->>SessionStore: process_prekey_bundle / establish session
SessionStore-->>Runtime: session established / identity saved
Runtime->>Runtime: spawn_oneshot(message_encrypt)
Runtime-->>WACore: EncryptOneResult (ciphertext / enc attrs)
end
WACore->>Client: assemble participant nodes & return stanza
Client->>Backend: update_device_lists(batch) -- canonicalize & upsert
Backend-->>Client: Ok / Err
Client->>Backend: best-effort delete stale rows (if canonical flip)
Backend-->>Client: log failures (non-fatal)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Review rate limit: 1/3 review remaining, refill in 33 minutes and 27 seconds. Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/store/signal_adapter.rs`:
- Around line 157-168: The current is_trusted_identity stub always returns
Ok(true), which bypasses the UntrustedIdentity path; change is_trusted_identity
to actually compare the provided IdentityKey against the stored device identity:
lookup the device record for the given ProtocolAddress, acquire the device
RwLock or appropriate store access used elsewhere in this module, read the
stored identity key, and if it differs return
Err(SignalProtocolError::UntrustedIdentity) so callers (e.g.,
process_prekey_bundle / send retry that calls save_identity) can surface and
handle identity changes; otherwise return Ok(true). Ensure you reference and use
the same device-store access patterns as other methods in this file to avoid
deadlock or contention.
In `@wacore/src/send.rs`:
- Around line 357-375: The current spawn_oneshot detaches the AbortHandle and
panics on a closed oneshot; instead keep the AbortHandle alive and cancel
spawned tasks when the caller's await is dropped or the oneshot closes, and
propagate a proper error instead of calling expect(...). Update spawn_oneshot
(and the similar block at the later occurrence) to return or store the
AbortHandle so you can call abort() when rx.await returns Err, and change the
async move to treat a closed oneshot as an error result (return a Result or map
to a send-path error) rather than panicking; ensure encrypt_for_devices callers
use the abort handle to cancel outstanding per-device tasks when the send
operation ends.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: babaacea-d291-4052-bde3-cbb1edb0bdbe
📒 Files selected for processing (9)
src/client/device_registry.rssrc/features/signal.rssrc/send.rssrc/store/signal_adapter.rssrc/usync.rsstorages/sqlite-storage/src/sqlite_store.rswacore/src/send.rswacore/src/store/signal_cache.rswacore/src/store/traits.rs
bb8d374 to
cdece1e
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cdece1ec4b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
cdece1e to
720c686
Compare
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
wacore/benches/send_receive_benchmark.rs (1)
84-89:⚠️ Potential issue | 🟠 MajorThese benchmark stores no longer model production clone semantics.
encrypt_for_devicesclones the session and identity stores per spawned device task. With#[derive(Clone)]on theseHashMap-backed fixtures, every task gets a deep copy of the store and writes into its private clone, while production uses cheap shared-handle clones into the same cache-backed state. That makes the reported fan-out numbers materially unrepresentative. Use anArc-backed fixture or the real adapter shape soClonebehaves like production.Also applies to: 164-165
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/benches/send_receive_benchmark.rs` around lines 84 - 89, The benchmark stores currently derive Clone which produces deep copies of their HashMap-backed state (e.g., MemIdentityStore with field identities) and causes each spawned task to write into private clones; change these fixtures to use shared handles instead (wrap the HashMap state in Arc<Mutex<HashMap<...>>> or Arc<RwLock<HashMap<...>>> or use the real adapter type) so that cloning the fixture is cheap and points at the same underlying cache; update MemIdentityStore (and the other fixture structs referenced around lines 164-165) to hold Arc-wrapped collections and rely on Arc's Clone semantics (or implement Clone to clone the Arcs) so benchmark behavior matches production.src/send.rs (1)
416-431:⚠️ Potential issue | 🔴 CriticalSerialize status/group sends before calling
prepare_group_stanza.These paths still run with no sender-key lock and no per-device
session_locks.prepare_group_stanzaalways advances the single sender-key chain for the chat/status, and when SKDM is needed it also mutates pairwise Signal sessions. Two concurrent sends can reuse sender-key state or racemessage_encrypt/process_prekey_bundleagainst the same recipients. The new comment at Lines 956-960 says this is safe, but it isn't.As per coding guidelines "Use
session_locksto serialize per-sender Signal encrypt/decrypt operations andmessage_enqueue_locksto serialize per-chat incoming message processing; outgoing sends are not per-chat locked".Also applies to: 458-473, 956-960, 1047-1062, 1092-1107
♻️ Duplicate comments (1)
wacore/src/send.rs (1)
357-375:⚠️ Potential issue | 🔴 CriticalDon't detach these fan-out tasks.
This still drops the
AbortHandleimmediately and then panics on a closed oneshot. Ifencrypt_for_devicesreturns early or the send is cancelled, the spawned work keeps mutating session/identity state after the lock window is gone, and a dropped task turns into a panic in the send path. Keep the abort handle alive until the await finishes and propagate a real error instead ofexpect(...).Also applies to: 378-392
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/src/send.rs` around lines 357 - 375, spawn_oneshot currently detaches the AbortHandle and panics on a closed oneshot (expect), which lets background work keep mutating state after callers return; instead keep the AbortHandle alive until the spawned future finishes and return a Result so failures are propagated. Change spawn_oneshot to capture the AbortHandle returned by rt.spawn (do not call detach), hold it alongside the oneshot receiver until rx.await completes, and when the receiver errors return an Err variant (propagate a meaningful error type instead of calling expect). Update the function signature/return type accordingly and adjust callers (e.g., encrypt_for_devices and the similar spawn site) to handle the Result and abort/dismiss the handle when the await completes.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@wacore/benches/send_receive_benchmark.rs`:
- Around line 84-89: The benchmark stores currently derive Clone which produces
deep copies of their HashMap-backed state (e.g., MemIdentityStore with field
identities) and causes each spawned task to write into private clones; change
these fixtures to use shared handles instead (wrap the HashMap state in
Arc<Mutex<HashMap<...>>> or Arc<RwLock<HashMap<...>>> or use the real adapter
type) so that cloning the fixture is cheap and points at the same underlying
cache; update MemIdentityStore (and the other fixture structs referenced around
lines 164-165) to hold Arc-wrapped collections and rely on Arc's Clone semantics
(or implement Clone to clone the Arcs) so benchmark behavior matches production.
---
Duplicate comments:
In `@wacore/src/send.rs`:
- Around line 357-375: spawn_oneshot currently detaches the AbortHandle and
panics on a closed oneshot (expect), which lets background work keep mutating
state after callers return; instead keep the AbortHandle alive until the spawned
future finishes and return a Result so failures are propagated. Change
spawn_oneshot to capture the AbortHandle returned by rt.spawn (do not call
detach), hold it alongside the oneshot receiver until rx.await completes, and
when the receiver errors return an Err variant (propagate a meaningful error
type instead of calling expect). Update the function signature/return type
accordingly and adjust callers (e.g., encrypt_for_devices and the similar spawn
site) to handle the Result and abort/dismiss the handle when the await
completes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: d5a1aacf-3cba-437a-aa3b-9fe75cfdabae
📒 Files selected for processing (5)
src/features/signal.rssrc/send.rswacore/Cargo.tomlwacore/benches/send_receive_benchmark.rswacore/src/send.rs
720c686 to
a1b01bc
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/send.rs`:
- Around line 956-960: prepare_group_stanza currently runs without holding
session-level locks which allows concurrent group sends to race on recipient
sessions and the sender-key chain; fix by acquiring the appropriate
session_locks for the entire fan-out device set (and the sender's group
sender-key chain) before calling prepare_group_stanza and hold them through the
sender-key mutation path and until stanza preparation is complete; use the
existing session_locks mechanism to serialize per-sender Signal encrypt/decrypt
operations (and keep message_enqueue_locks only for incoming per-chat
processing), updating the code paths around prepare_group_stanza, the group_info
usage, and the sender-key mutation sections so they acquire/release those locks
correctly.
In `@wacore/benches/send_receive_benchmark.rs`:
- Around line 679-681: The benchmark currently calls BenchRuntime::default()
inside each sample which constructs a new ThreadPool and contaminates the
measured work; instead create and reuse a single runtime/thread-pool in the
benchmark setup (e.g., move the BenchRuntime::default() call out of the
per-sample loop into the setup fixture or a static shared instance) and pass
that shared runtime to prepare_group_stanza so the measured code only runs
prepare_group_stanza rather than including executor construction.
In `@wacore/src/send.rs`:
- Around line 566-679: The closure make_session_task currently runs CPU‑heavy
X3DH work inside spawn_oneshot (executor tasks) by calling process_prekey_bundle
directly, which can block the async runtime; change it to run the CPU-bound
portions on a blocking worker (use runtime.spawn_blocking or a crypto pool) and
keep only async store/IO awaits on the executor: move the call(s) to
process_prekey_bundle (and any RNG/key derivation) into a spawn_blocking closure
that returns results, then await that future inside the existing async block and
continue to call identity_store.save_identity/other async store methods on the
executor; update spawn_oneshot usage in make_session_task (and the analogous
code around lines 719-807) to reflect this split.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: fa4e106b-d2ce-426d-8a59-d5aae96abc7a
📒 Files selected for processing (5)
src/features/signal.rssrc/send.rswacore/Cargo.tomlwacore/benches/send_receive_benchmark.rswacore/src/send.rs
There was a problem hiding this comment.
♻️ Duplicate comments (3)
wacore/benches/send_receive_benchmark.rs (1)
693-695:⚠️ Potential issue | 🟠 MajorDon’t rebuild the thread pool inside each measured send.
run_group_sendis on the hot path for the benchmark, andBenchRuntime::default()allocates a freshThreadPoolevery sample. That means the measurement now includes executor startup instead of justprepare_group_stanza. Reuse one runtime from the fixture.Suggested change
struct GrpSendData { alice: User, group_jid: Jid, participants: Vec<Jid>, force_skdm: bool, resolver: MockResolver, + runtime: BenchRuntime, msg: wa::Message, } fn setup_group_send(n: usize) -> GrpSendData { @@ GrpSendData { alice, group_jid, participants, force_skdm: false, resolver: MockResolver(devices), + runtime: BenchRuntime::default(), msg: text_msg(), } } fn run_group_send(d: &mut GrpSendData) { @@ - let runtime = BenchRuntime::default(); let result = futures::executor::block_on(prepare_group_stanza( - &runtime, + &d.runtime, &mut stores, &d.resolver,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/benches/send_receive_benchmark.rs` around lines 693 - 695, The benchmark currently constructs a fresh runtime for every measured send by calling BenchRuntime::default() inside run_group_send, which includes thread-pool startup in the timing; move the creation of the runtime out of the hot path (create a single BenchRuntime once in the test/fixture setup) and reuse that instance when calling prepare_group_stanza and other measured functions; update run_group_send to accept a &BenchRuntime (or use the fixture-scoped runtime variable) so you no longer call BenchRuntime::default() per sample and the measurement only covers prepare_group_stanza.src/send.rs (1)
416-430:⚠️ Potential issue | 🔴 CriticalReintroduce session serialization around
prepare_group_stanza.This is not safe. Both the status path and the group path call into
prepare_group_stanzawithout holdingsession_locks, even though that lower layer mutates per-device Signal sessions during SKDM fan-out and advances the sender-key chain for the outgoing group/status message. Two concurrent sends can race and corrupt that state. The new “no client-level lock needed” assumption is false.As per coding guidelines, "Use
session_locksto serialize per-sender Signal encrypt/decrypt operations andmessage_enqueue_locksto serialize per-chat incoming message processing; outgoing sends are not per-chat locked".Also applies to: 458-473, 956-960, 1047-1062, 1092-1107
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/send.rs` around lines 416 - 430, prepare_group_stanza mutates per-device Signal sessions and advances sender-key state, so wrap each call to wacore::send::prepare_group_stanza (the calls at the shown sites and the other occurrences you noted) inside the existing session serialization: acquire the session_locks (the mutex/guard used to serialize per-sender Signal operations) before calling prepare_group_stanza and release it after the call completes (including on error paths), so SKDM fan-out and sender-key chain updates cannot race; ensure you use the same session_locks mechanism used elsewhere for encrypt/decrypt serialization rather than assuming no client-level lock is needed.wacore/src/send.rs (1)
595-679:⚠️ Potential issue | 🟠 MajorMove the crypto hot path off the async executor.
This still runs the expensive X3DH / session-setup and
message_encryptwork insideRuntime::spawn. With fan-out at 16, that can pin the shared runtime and delay I/O/ack progress for the rest of the client under load. Split the async store calls from the CPU-bound section and run the derivation/encrypt step onspawn_blockingor a dedicated crypto pool.As per coding guidelines, "Wrap blocking I/O (e.g.,
ureq) and heavy CPU work intokio::task::spawn_blockingfor async code".Also applies to: 729-757
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/src/send.rs` around lines 595 - 679, The crypto hot path (process_prekey_bundle and any message_encrypt work) is currently running inside spawn_oneshot and can block the async runtime; refactor by separating async store access from CPU-bound processing: perform async reads/writes to session_store and identity_store and clone or extract the needed Bundle, addr, and identity data within the async task, then offload the heavy work (calling process_prekey_bundle and subsequent encrypt/derivation steps that use rand::make_rng and UsePQRatchet::No) to tokio::task::spawn_blocking (or a dedicated crypto thread pool) and await that result; ensure you still persist updated identity via identity_store.save_identity in the async context (or do a follow-up async save after the blocking section) and keep the same error handling and logging around process_prekey_bundle, using the function names process_prekey_bundle and identity_store.save_identity to locate the code to change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@src/send.rs`:
- Around line 416-430: prepare_group_stanza mutates per-device Signal sessions
and advances sender-key state, so wrap each call to
wacore::send::prepare_group_stanza (the calls at the shown sites and the other
occurrences you noted) inside the existing session serialization: acquire the
session_locks (the mutex/guard used to serialize per-sender Signal operations)
before calling prepare_group_stanza and release it after the call completes
(including on error paths), so SKDM fan-out and sender-key chain updates cannot
race; ensure you use the same session_locks mechanism used elsewhere for
encrypt/decrypt serialization rather than assuming no client-level lock is
needed.
In `@wacore/benches/send_receive_benchmark.rs`:
- Around line 693-695: The benchmark currently constructs a fresh runtime for
every measured send by calling BenchRuntime::default() inside run_group_send,
which includes thread-pool startup in the timing; move the creation of the
runtime out of the hot path (create a single BenchRuntime once in the
test/fixture setup) and reuse that instance when calling prepare_group_stanza
and other measured functions; update run_group_send to accept a &BenchRuntime
(or use the fixture-scoped runtime variable) so you no longer call
BenchRuntime::default() per sample and the measurement only covers
prepare_group_stanza.
In `@wacore/src/send.rs`:
- Around line 595-679: The crypto hot path (process_prekey_bundle and any
message_encrypt work) is currently running inside spawn_oneshot and can block
the async runtime; refactor by separating async store access from CPU-bound
processing: perform async reads/writes to session_store and identity_store and
clone or extract the needed Bundle, addr, and identity data within the async
task, then offload the heavy work (calling process_prekey_bundle and subsequent
encrypt/derivation steps that use rand::make_rng and UsePQRatchet::No) to
tokio::task::spawn_blocking (or a dedicated crypto thread pool) and await that
result; ensure you still persist updated identity via
identity_store.save_identity in the async context (or do a follow-up async save
after the blocking section) and keep the same error handling and logging around
process_prekey_bundle, using the function names process_prekey_bundle and
identity_store.save_identity to locate the code to change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 07dfdaf2-23e2-4573-a6c4-74b21a9e46d6
📒 Files selected for processing (5)
src/features/signal.rssrc/send.rswacore/Cargo.tomlwacore/benches/send_receive_benchmark.rswacore/src/send.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: faa0312f46
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
The IdentityKeyStore impl on Device is a documented WA Web stub that always returns Ok(true). The adapter was acquiring an async RwLock read on the Device just to delegate to it — once per encrypt, so 666× per large-group send. Return Ok(true) directly.
get_identity was holding the identities mutex across backend.load_identity().await, serializing every concurrent caller on the round-trip. get_session, peek_session and has_session already use the scoped-lock + double-check pattern. Apply it here too — required for the upcoming parallel encrypt fan-out, and benchmark shows ~200x reduction in total wall-time under 666 concurrent identity lookups.
The previous wording asserted a per-device lock inside prepare_group_stanza. None exists; the encrypt loop is a plain sequential .await iteration with no per-recipient lock. The no-client-lock-needed conclusion still holds, but for the actual reason: each recipient device has its own Signal session, so concurrent group sends to the same chat don't share mutable state on that path.
Per-user update_device_list inside the usync loop spent a spawn_blocking + SQLite commit per participant — 582 hops on a typical large-group send. Add update_device_lists(Vec<DeviceListRecord>) on the SignalStore trait (default impl loops, SQLite backend overrides with a single transaction) and a sibling Client::update_device_lists that resolves canonical keys, populates the in-memory cache synchronously per record, then issues one batched backend write. Empirically: 5.1ms → 870us for 582 rows in-memory; production WAL fsync amortization makes the speedup larger.
…Jid>> The encrypt fan-out built a HashMap mapping each device JID to its LID-upgraded encryption JID. For the common case where most devices need an upgrade, the map paid hash + alloc per insert and per lookup (roughly 2N hashes per group send). Replace with a Vec<Option<Jid>> indexed by position in `devices`: `None` means "use the original", `Some(jid)` is the upgrade. Direct indexing, contiguous memory, no hashing on the hot path. Also prepares the structure for the upcoming parallel encrypt loop, where indexed access is friendlier than a shared HashMap.
The encrypt fan-out was a single-threaded loop over recipient devices: 666 sequential X3DH derivations + AES-CBC + HMAC, all on one worker thread. Recipients have independent Signal sessions, so this trivially parallelizes. Refactor encrypt_for_devices to spawn per-device tasks via the existing `wacore::Runtime` abstraction (no new tokio main dep — the platform- agnostic split stays intact). Each task is dispatched through `Runtime::spawn` + a oneshot channel for result delivery; concurrency is bounded by FuturesUnordered with ENCRYPT_FANOUT_CONCURRENCY = 16. Two loops are parallelized: process_prekey_bundle (X3DH for first- contact devices) and message_encrypt (per-device crypto). Plaintext is shared across tasks via Arc<[u8]>; the SignalAdapter clones are cheap (Arc bumps under the hood) and the underlying SignalStoreCache is already interior-mutable, so each task can hold its own &mut handle without contention beyond brief mutex acquires. Adds Clone + 'static bounds on the S/I trait params; the existing adapters in whatsapp-rust already satisfy them. prepare_dm_stanza, prepare_group_stanza and encrypt_for_devices gain a `runtime` parameter threaded through from the call sites — the embedder's choice of runtime stays the source of truth. Wire-order: the previous code preserved input device order in the participant_nodes Vec. Output now follows completion order. WA Web (`WAWebPhashUtils.phashV2`) sorts before computing the participant hash, and our `participant_list_hash` does the same — server-side phash validation is order-independent. Failure semantics unchanged: per-device encrypt errors continue to log + skip rather than aborting the send. WA Web's GroupKeyDistributionMsg.js rejects the whole send only when a primary (device 0) recipient fails; whatsmeow and this codebase have always been more permissive.
faa0312 to
36b41ba
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
wacore/benches/send_receive_benchmark.rs (1)
120-134:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winDon’t add new
unwrap()s in non-test benchmark code.These lock acquisitions will panic the whole benchmark run on a poisoned mutex. The trait already returns
SigResult, so this path should surface a real error instead of aborting on Line 120, Line 134, Line 189, and Line 195.As per coding guidelines "Never use
.unwrap()outside of test code".Also applies to: 188-195
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/benches/send_receive_benchmark.rs` around lines 120 - 134, Replace all direct `.unwrap()` calls on `self.identities.lock()` with proper error handling that converts a poisoned mutex into a SigResult error instead of panicking: e.g. in the method that inserts/returns IdentityChange (the block that creates `guard`, compares and inserts), in `get_identity`, and in any other methods using `self.identities.lock()` (the occurrences around the `is_trusted_identity`/identity methods), call `self.identities.lock().map_err(|e| /* convert e into the crate's SigError */)?` or otherwise map the PoisonError into a SigResult::Err and propagate it; ensure you return an appropriate SigError value rather than using `.unwrap()`, keeping the rest of the logic (comparison, insert, cloning) unchanged and using the same symbols (`guard`, `identities`, `get_identity`, `IdentityChange::from_changed`) to locate the fixes.
♻️ Duplicate comments (3)
src/store/signal_adapter.rs (1)
157-168:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winDon't turn identity verification into a blanket allow.
process_prekey_bundleonly surfaces unexpected identity rotation through this hook. ReturningOk(true)here means theUntrustedIdentityrecovery path inwacore/src/send.rsnever runs, so new identities get silently accepted and cached instead of being surfaced.#!/bin/bash # Verify that identity-rotation handling depends on this hook and that the send path only retries on UntrustedIdentity. rg -n -C3 'is_trusted_identity\s*\(|UntrustedIdentity|save_identity\s*\(' src/store/signal_adapter.rs wacore/src/send.rs🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/store/signal_adapter.rs` around lines 157 - 168, The is_trusted_identity stub currently always returns Ok(true), which bypasses the UntrustedIdentity recovery in process_prekey_bundle and causes silent acceptance of rotated identities; change is_trusted_identity to perform a real check against stored device identity data instead of blanket-allow: look up the stored identity for the given ProtocolAddress/device (the same source that save_identity writes to), compare it to the provided IdentityKey, and return Ok(true) only if they match, otherwise return Err(SignalProtocolError::UntrustedIdentity) or Ok(false) as appropriate so the caller (e.g., the send path handling UntrustedIdentity) can trigger the recovery/verification flow. Ensure you reference and reuse the same storage accessors used by save_identity to avoid races and keep locking minimal while performing the comparison.src/send.rs (1)
416-431:⚠️ Potential issue | 🔴 Critical | 🏗️ Heavy liftGroup/status sends still need
session_locksaround stanza preparation.Both paths call
prepare_group_stanzawithout acquiring the per-device session locks thatencrypt_for_devicesrequires, and they advance the group sender-key chain in the same unlocked window. Two concurrent sends to the same group or status fan-out can ratchet the same recipient sessions or sender key at once and corrupt Signal state. Reusebuild_session_lock_keys/session_mutexes_forhere and hold them across the retry path too. Based on learnings:Use session_locks to serialize per-sender Signal encrypt/decrypt operations and message_enqueue_locks to serialize per-chat incoming message processing; outgoing sends are not per-chat locked.Also applies to: 458-473, 956-960, 1047-1062, 1092-1107
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/send.rs` around lines 416 - 431, prepare_group_stanza is being called without holding the per-device session locks required by encrypt_for_devices, which can allow concurrent send operations to ratchet sender keys/sessions and corrupt Signal state; fix by computing the necessary lock keys via build_session_lock_keys and acquiring the mutexes from session_mutexes_for (session_locks) before calling prepare_group_stanza and keep those locks held across the entire retry path that performs the prepare/send, ensuring they are released only after the send/retry completes; apply the same pattern wherever prepare_group_stanza (and similar group/status send paths) is invoked so the per-sender Signal encrypt/decrypt operations are serialized.wacore/src/send.rs (1)
566-679:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftMove the crypto fan-out off
Runtime::spawn.Both spawned branches still run
process_prekey_bundle/message_encryptdirectly on executor tasks. At fan-out 16 that is enough CPU to pin runtime workers and delay socket/ACK progress for the rest of the client. Split the async store access from the heavy X3DH/AES/HMAC work and run the CPU section onRuntime::spawn_blockingor a dedicated crypto pool. As per coding guidelines,Wrap blocking I/O (e.g., ureq) and heavy CPU work in tokio::task::spawn_blocking for async code.#!/bin/bash # Verify that the runtime exposes spawn_blocking and that the fan-out still uses spawn_oneshot around crypto-heavy operations. rg -n -C2 'spawn_blocking|spawn_oneshot|process_prekey_bundle\s*\(|message_encrypt\s*\(' wacore/src/send.rs wacore/src/runtime.rs src/runtime_impl.rsAlso applies to: 719-807
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/src/send.rs` around lines 566 - 679, The current make_session_task/spawn_oneshot path runs process_prekey_bundle (and elsewhere message_encrypt) directly on the async runtime, which can starve the executor at high fan-out; change the implementation so async store/IO work (loading/saving identity/session via session_store/identity_store) remains on the async task but the CPU-heavy X3DH/AES/HMAC work is executed inside runtime.spawn_blocking (or a dedicated crypto thread-pool) — specifically, wrap the call to process_prekey_bundle (and the retry call after save_identity) inside runtime.spawn_blocking (or offload to a crypto worker) and await its result from the async task, preserving the same error handling/log messages; do the same for message_encrypt paths referenced in the review. Ensure you still use spawn_oneshot for the task envelope but move the crypto-heavy per-device calls into spawn_blocking so the async executor is not pinned.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/client/device_registry.rs`:
- Around line 217-232: The cache is being warmed before the backend batch is
committed, causing a mismatch if update_device_lists (called via
self.persistence_manager.backend().update_device_lists(prepared)) fails; move
the cache insert and any to_delete handling to occur only after
update_device_lists completes successfully: perform
backend.update_device_lists(prepared).await and check for success, then iterate
the prepared records (using canonical_key/record_for_cache) to insert into
self.device_registry_cache and handle to_delete removals so the cache reflects
committed state.
---
Outside diff comments:
In `@wacore/benches/send_receive_benchmark.rs`:
- Around line 120-134: Replace all direct `.unwrap()` calls on
`self.identities.lock()` with proper error handling that converts a poisoned
mutex into a SigResult error instead of panicking: e.g. in the method that
inserts/returns IdentityChange (the block that creates `guard`, compares and
inserts), in `get_identity`, and in any other methods using
`self.identities.lock()` (the occurrences around the
`is_trusted_identity`/identity methods), call
`self.identities.lock().map_err(|e| /* convert e into the crate's SigError */)?`
or otherwise map the PoisonError into a SigResult::Err and propagate it; ensure
you return an appropriate SigError value rather than using `.unwrap()`, keeping
the rest of the logic (comparison, insert, cloning) unchanged and using the same
symbols (`guard`, `identities`, `get_identity`, `IdentityChange::from_changed`)
to locate the fixes.
---
Duplicate comments:
In `@src/send.rs`:
- Around line 416-431: prepare_group_stanza is being called without holding the
per-device session locks required by encrypt_for_devices, which can allow
concurrent send operations to ratchet sender keys/sessions and corrupt Signal
state; fix by computing the necessary lock keys via build_session_lock_keys and
acquiring the mutexes from session_mutexes_for (session_locks) before calling
prepare_group_stanza and keep those locks held across the entire retry path that
performs the prepare/send, ensuring they are released only after the send/retry
completes; apply the same pattern wherever prepare_group_stanza (and similar
group/status send paths) is invoked so the per-sender Signal encrypt/decrypt
operations are serialized.
In `@src/store/signal_adapter.rs`:
- Around line 157-168: The is_trusted_identity stub currently always returns
Ok(true), which bypasses the UntrustedIdentity recovery in process_prekey_bundle
and causes silent acceptance of rotated identities; change is_trusted_identity
to perform a real check against stored device identity data instead of
blanket-allow: look up the stored identity for the given ProtocolAddress/device
(the same source that save_identity writes to), compare it to the provided
IdentityKey, and return Ok(true) only if they match, otherwise return
Err(SignalProtocolError::UntrustedIdentity) or Ok(false) as appropriate so the
caller (e.g., the send path handling UntrustedIdentity) can trigger the
recovery/verification flow. Ensure you reference and reuse the same storage
accessors used by save_identity to avoid races and keep locking minimal while
performing the comparison.
In `@wacore/src/send.rs`:
- Around line 566-679: The current make_session_task/spawn_oneshot path runs
process_prekey_bundle (and elsewhere message_encrypt) directly on the async
runtime, which can starve the executor at high fan-out; change the
implementation so async store/IO work (loading/saving identity/session via
session_store/identity_store) remains on the async task but the CPU-heavy
X3DH/AES/HMAC work is executed inside runtime.spawn_blocking (or a dedicated
crypto thread-pool) — specifically, wrap the call to process_prekey_bundle (and
the retry call after save_identity) inside runtime.spawn_blocking (or offload to
a crypto worker) and await its result from the async task, preserving the same
error handling/log messages; do the same for message_encrypt paths referenced in
the review. Ensure you still use spawn_oneshot for the task envelope but move
the crypto-heavy per-device calls into spawn_blocking so the async executor is
not pinned.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 455a7494-6752-4915-8fc5-84b660774733
📒 Files selected for processing (11)
src/client/device_registry.rssrc/features/signal.rssrc/send.rssrc/store/signal_adapter.rssrc/usync.rsstorages/sqlite-storage/src/sqlite_store.rswacore/Cargo.tomlwacore/benches/send_receive_benchmark.rswacore/src/send.rswacore/src/store/signal_cache.rswacore/src/store/traits.rs
Building the ThreadPool inside run_group_send made every iai-callgrind sample charge the syscalls of thread-pool startup to the measured encrypt path. Move it to GrpSendData so setup_group_send constructs it once and the measured body only runs prepare_group_stanza.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
wacore/benches/send_receive_benchmark.rs (1)
693-695:⚠️ Potential issue | 🟠 Major | ⚡ Quick winStop constructing
BenchRuntimein the measured path.Line 693 builds a new thread pool during the benchmarked send path, so the reported cost includes runtime startup, not just
prepare_group_stanza. We need clean numbers.Minimal fix
struct GrpSendData { alice: User, group_jid: Jid, participants: Vec<Jid>, force_skdm: bool, resolver: MockResolver, + runtime: BenchRuntime, msg: wa::Message, } fn setup_group_send(n: usize) -> GrpSendData { @@ GrpSendData { alice, group_jid, participants, force_skdm: false, resolver: MockResolver(devices), + runtime: BenchRuntime::default(), msg: text_msg(), } } fn run_group_send(d: &mut GrpSendData) { @@ - let runtime = BenchRuntime::default(); let result = futures::executor::block_on(prepare_group_stanza( - &runtime, + &d.runtime, &mut stores, &d.resolver, &mut group_info, &own_jid,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/benches/send_receive_benchmark.rs` around lines 693 - 695, BenchRuntime is being constructed inside the measured path so runtime startup is counted; move creation of the runtime out of the benchmarked call by instantiating BenchRuntime before calling futures::executor::block_on and then pass a reference to that prebuilt runtime into prepare_group_stanza (symbols: BenchRuntime, prepare_group_stanza, futures::executor::block_on). Ensure the benchmark only wraps the block_on(prepare_group_stanza(&runtime, ...)) invocation and not the BenchRuntime::default() construction.wacore/src/send.rs (1)
566-679:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftDon’t run the crypto hot path on the shared async executor.
This is still the X3DH / AES / HMAC-heavy path, and both loops are dispatching it through
Runtime::spawnrather thanRuntime::spawn_blocking. At concurrency 16 that means crypto jobs can compete directly with socket and ACK work on the runtime workers, which is exactly where latency starts getting weird under load. Move the CPU-heavy section behindspawn_blocking(or a dedicated crypto pool) and keep only the async store access on the executor.#!/bin/bash set -euo pipefail sed -n '13,31p' wacore/src/runtime.rs sed -n '566,679p' wacore/src/send.rs sed -n '719,757p' wacore/src/send.rsAs per coding guidelines, "Wrap blocking I/O (e.g.,
ureq) and heavy CPU work intokio::task::spawn_blockingfor async code".Also applies to: 719-757
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/src/send.rs` around lines 566 - 679, The code is spawning CPU-heavy crypto work via Runtime::spawn in make_session_task/spawn_oneshot which lets X3DH/AES/HMAC run on the shared async executor; change it to run the blocking/CPU portion inside a blocking thread pool. Concretely: keep async store/identity I/O (identity_store.save_identity, session_store access if async) on the async task but move calls to process_prekey_bundle (including RNG creation and the retry call) into tokio::task::spawn_blocking (or your Runtime::spawn_blocking wrapper) and await its result; update the spawn_oneshot body to call spawn_blocking for the crypto path and handle returned Result/Err the same way. Apply the same change to the other loop that also calls process_prekey_bundle (the block referenced around 719-757) so all X3DH/AES/HMAC work runs on the blocking pool rather than Runtime::spawn.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@wacore/src/send.rs`:
- Around line 619-672: The UntrustedIdentity match arm is unreachable because
IdentityAdapter::is_trusted_identity always returns Ok(true); remove the entire
Err(SignalProtocolError::UntrustedIdentity(_)) branch including its logs, the
bundle.identity_key() call, identity_store.save_identity(...) await, creation of
rng via rand::make_rng, and the subsequent process_prekey_bundle(...) retry
(UsePQRatchet::No), simplifying the match to omit this recovery path and avoid
dead code in send.rs.
- Around line 435-451: The wasm-specific spawn_oneshot implementation doesn't
enforce Send bounds so it can violate Runtime::spawn's contract; update the
generics on spawn_oneshot so F: Future<Output = T> + Send + 'static and T: Send
+ 'static (i.e., add Send to both F and T bounds) so the async block passed to
rt.spawn meets Runtime::spawn's Send + 'static requirement and the returned
Spawned stays correct.
---
Duplicate comments:
In `@wacore/benches/send_receive_benchmark.rs`:
- Around line 693-695: BenchRuntime is being constructed inside the measured
path so runtime startup is counted; move creation of the runtime out of the
benchmarked call by instantiating BenchRuntime before calling
futures::executor::block_on and then pass a reference to that prebuilt runtime
into prepare_group_stanza (symbols: BenchRuntime, prepare_group_stanza,
futures::executor::block_on). Ensure the benchmark only wraps the
block_on(prepare_group_stanza(&runtime, ...)) invocation and not the
BenchRuntime::default() construction.
In `@wacore/src/send.rs`:
- Around line 566-679: The code is spawning CPU-heavy crypto work via
Runtime::spawn in make_session_task/spawn_oneshot which lets X3DH/AES/HMAC run
on the shared async executor; change it to run the blocking/CPU portion inside a
blocking thread pool. Concretely: keep async store/identity I/O
(identity_store.save_identity, session_store access if async) on the async task
but move calls to process_prekey_bundle (including RNG creation and the retry
call) into tokio::task::spawn_blocking (or your Runtime::spawn_blocking wrapper)
and await its result; update the spawn_oneshot body to call spawn_blocking for
the crypto path and handle returned Result/Err the same way. Apply the same
change to the other loop that also calls process_prekey_bundle (the block
referenced around 719-757) so all X3DH/AES/HMAC work runs on the blocking pool
rather than Runtime::spawn.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: f1421e2e-1d14-4eb0-9b80-34191e3b8e81
📒 Files selected for processing (11)
src/client/device_registry.rssrc/features/signal.rssrc/send.rssrc/store/signal_adapter.rssrc/usync.rsstorages/sqlite-storage/src/sqlite_store.rswacore/Cargo.tomlwacore/benches/send_receive_benchmark.rswacore/src/send.rswacore/src/store/signal_cache.rswacore/src/store/traits.rs
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@wacore/benches/send_receive_benchmark.rs`:
- Line 120: Replace all uses of .unwrap() on Mutex locks in
send_receive_benchmark.rs (e.g., the call self.identities.lock().unwrap() and
the other occurrences noted at lines 134, 189, 192, 195) with explicit failure
handling: at minimum call .expect(...) with a clear diagnostic message
referencing the mutex being locked (for example "failed to lock identities mutex
in send/receive benchmark") or map the PoisonError into a usable error via
.map_err(...) and propagate it. Update the locking sites (the one using
self.identities.lock(), and the other mutex.lock() calls) to use .expect(...)
messages that identify the field (identities, peers, messages, etc.) and the
operation (locking during benchmark) so panics contain actionable context.
- Around line 61-63: The code currently discards the Result of
pool.spawn(future) and returns AbortHandle::noop(), which hides spawn failures
and invalidates benchmark results; update the caller (the function containing
pool.spawn and AbortHandle::noop) to check the Result returned by
pool.spawn(future) and fail fast on error (e.g., propagate a Result or panic)
instead of ignoring it, ensuring any Err from pool.spawn is surfaced with a
clear message referencing the spawn failure so benchmarks never continue
silently on scheduling failures.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 679d04e7-cfca-4c40-a473-d7d9e3bb593b
📒 Files selected for processing (1)
wacore/benches/send_receive_benchmark.rs
All IdentityKeyStore impls in the workspace return Ok(true) from is_trusted_identity (matches WA Web's TOFU model), so the SignalProtocolError::UntrustedIdentity arm in encrypt_for_devices was dead code. process_prekey_bundle's own save_identity persists the new identity on every prekey bundle, so identity rotations are handled transparently without a manual retry path.
A silent spawn failure would skip a device's encrypt task, the oneshot would close, and the fan-out's log+skip path would mask it — measuring fewer encrypts than intended and reporting an inflated speedup. The pool only fails to spawn after shutdown, which never happens in a bench, so .expect surfaces the bug instead of hiding it.
Group sends to large groups (582-member chat in the production logs) showed a CPU spike to ~80% for 3-4s and a ~280ms wall-clock encrypt loop. The root cause was a serial fan-out: 666 X3DH derivations + AES-CBC + HMAC executed back-to-back on a single worker thread. This PR parallelizes that path and fixes the surrounding bottlenecks that would have constrained the win.
What's in here
260ad32is_trusted_identityreturnsOk(true)directly instead of acquiring adevice.read().awaitto delegate to a stub5ecca43signal_cache::get_identityreleases the mutex beforebackend.load_identity().await, matching the existingget_session/has_sessionpattern8133f6e0667958update_device_lists(Vec<DeviceListRecord>)trait method, SQLite override does one transaction;usync.rsnow batches the per-recipient writesa0c4d84HashMap<&Jid, Jid>→Vec<Option<Jid>>indexed by device position in the encrypt fan-outbb8d374process_prekey_bundleandmessage_encryptloops (concurrency 16) viaRuntime::spawn+FuturesUnorderedDesign notes
wacore::Runtimetrait (no new tokio main dep on wacore). Each task is dispatched throughRuntime::spawn+ a oneshot channel for result delivery; concurrency is bounded byFuturesUnordered.<to>participant order. The previous code preserved input order; output now follows completion order. WA Web'sWAWebPhashUtils.phashV2and ourparticipant_list_hashboth sort before computing the participant hash, so server-side validation is order-independent.GroupKeyDistributionMsg.jsis stricter — primary-device failure rejects the whole send — but that's an existing divergence and not part of this PR.Clone + 'staticbounds added on the session/identity store generics. Thesignal_adapter.rstypes already satisfy them (cheap Arc bumps).#[cfg]-gatedspawn_oneshotmirrors the WASM Runtime variant). No CPU parallelism gain on single-threaded runtimes — FuturesUnordered still polls cooperatively, but the crypto stays serialized on the one thread. Slight overhead vs serial; acceptable for now.Empirical setup
Microbenchmarks live at
/tmp/perf-audit/(not committed). Each item has its own bin; numbers above are measured on Oracle ARM64 with a multi-threaded tokio runtime.Test plan
cargo fmt --allcargo clippy --workspace --exclude e2e-tests --tests(clean)cargo test --workspace --exclude e2e-tests(659 wacore + 448 whatsapp-rust + others, 0 failures)