perf(send): trim group-send warm-path CPU and cold fan-out allocations - #936
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI (base), Organization UI (inherited) Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds a warm memo for repeated empty SKDM resolution, switches multi-device encryption to chunked fan-out, and moves group stanza locking so SKDM creation and SKMSG encryption occur under the sender-key chain lock. ChangesSend path updates
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested labels: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 84891260e4
ℹ️ 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".
📦 Binary size report
.text per crate
Top movers (cargo-bloat attribution)
Baseline: |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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/src/send/group.rs (1)
384-430: 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy liftDon’t hold the sender-key chain lock during SKDM fan-out.
The guard is acquired before
encrypt_for_devices_with_sessions(...).await, so a cold large-group send holds the per-chain mutex while doing pairwise device encryption. That work is chain-independent aftercreate_sender_key_distribution_message_for_groupreturns, and it serializes same-group sends longer than the “ratchet steps” comment claims.🔒 Suggested shape for narrowing the critical section
- let chain_lock = stores - .sender_key_store - .sender_key_lock(&sender_key_name) - .await; - let chain_guard = chain_lock.lock().await; + let chain_lock = stores + .sender_key_store + .sender_key_lock(&sender_key_name) + .await; + let mut axolotl_skdm_bytes = None; + let skmsg = { + let chain_guard = chain_lock.lock().await; + if distribution_list.is_some() { + axolotl_skdm_bytes = Some( + create_sender_key_distribution_message_for_group( + stores.sender_key_store, + &sender_key_name, + ) + .await?, + ); + } + let skmsg = encrypt_group_message( + stores.sender_key_store, + &sender_key_name, + &plaintext, + &mut rand::make_rng::<rand::rngs::StdRng>(), + ) + .await?; + drop(chain_guard); + skmsg + }; + if let Some(ref distribution_list) = distribution_list { - let axolotl_skdm_bytes = create_sender_key_distribution_message_for_group( - stores.sender_key_store, - &sender_key_name, - ) - .await?; + let Some(axolotl_skdm_bytes) = axolotl_skdm_bytes else { + unreachable!("SKDM bytes are created whenever distribution_list is present"); + }; if let Some(plan) = session_plan { let skdm_wrapper_msg = wa::Message {Also applies to: 474-485
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@wacore/src/send/group.rs` around lines 384 - 430, The sender-key chain lock in send_group is held too long because it remains active through encrypt_for_devices_with_sessions, which serializes cold large-group sends unnecessarily. Narrow the critical section in send/group.rs by keeping the lock only around create_sender_key_distribution_message_for_group and any chain-state updates, then release it before SKDM fan-out and device encryption; use the existing sender_key_lock and chain_guard flow as the anchor for the refactor. Apply the same lock-scope reduction to the matching later block referenced in the diff so the SKDM and main message paths both avoid holding the per-chain mutex across chain-independent work.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@wacore/src/send/encrypt.rs`:
- Around line 730-736: The fan-out partitioning in encrypt.rs does not reliably
preserve ENCRYPT_FANOUT_CONCURRENCY because using div_ceil with step_by can
yield too few chunks for totals just above the limit. Update the chunking logic
in the device fan-out loop around total, num_chunks, and in_flight to partition
by chunk index instead of stepping by chunk_size, so the work is split into
exactly num_chunks slices and the intended parallelism is maintained.
---
Outside diff comments:
In `@wacore/src/send/group.rs`:
- Around line 384-430: The sender-key chain lock in send_group is held too long
because it remains active through encrypt_for_devices_with_sessions, which
serializes cold large-group sends unnecessarily. Narrow the critical section in
send/group.rs by keeping the lock only around
create_sender_key_distribution_message_for_group and any chain-state updates,
then release it before SKDM fan-out and device encryption; use the existing
sender_key_lock and chain_guard flow as the anchor for the refactor. Apply the
same lock-scope reduction to the matching later block referenced in the diff so
the SKDM and main message paths both avoid holding the per-chain mutex across
chain-independent work.
🪄 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 (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 69e31442-95d6-4a28-985f-5ab443f202b4
📒 Files selected for processing (5)
src/client.rssrc/client/lifecycle.rssrc/send/mod.rswacore/src/send/encrypt.rswacore/src/send/group.rs
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
3 issues found across 5 files
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
013adce to
5f4c159
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 (1)
wacore/src/send/tests.rs (1)
3482-3491: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winThis doc comment is now attached to the wrong test — that's sloppy, let's fix it.
The comment at Lines 3482-3485 describes reporting a replaced identity via the resolver, but you inserted the new empty-device-set test right underneath it. In Rust doc comments bind to the next item, so this now documents
encrypt_for_devices_with_sessions_raw_handles_empty_device_set— which has nothing to do with replaced identities. Meanwhileencrypt_for_devices_reports_replaced_identityat Line 3524 lost its description entirely. Move it so each comment sits above the test it actually explains. Things need to be right.🧹 Proposed fix: move the comment back to its test
- /// The send path must report a replaced identity via the resolver when - /// establishing a session whose bundle carries a new identity key for an - /// address we already knew (peer reinstall). Mirrors WA Web saveIdentity - /// -> handleNewIdentity firing during outbound session setup. /// Regression (PR `#936` review): the chunked encrypt fan-out must return an /// empty result — not divide by zero — when the device set is empty. This is /// reachable on the cold force-SKDM path once the sender/hosted devices are /// filtered out, and the previous per-device fan-out returned empty here. #[tokio::test] async fn encrypt_for_devices_with_sessions_raw_handles_empty_device_set() {Then re-add the replaced-identity comment above
encrypt_for_devices_reports_replaced_identityat Line 3524:+ /// The send path must report a replaced identity via the resolver when + /// establishing a session whose bundle carries a new identity key for an + /// address we already knew (peer reinstall). Mirrors WA Web saveIdentity + /// -> handleNewIdentity firing during outbound session setup. #[tokio::test] async fn encrypt_for_devices_reports_replaced_identity() {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@wacore/src/send/tests.rs` around lines 3482 - 3491, The doc comment describing the replaced-identity resolver behavior is attached to the wrong test, so move that comment from encrypt_for_devices_with_sessions_raw_handles_empty_device_set to the test it actually documents, encrypt_for_devices_reports_replaced_identity. Keep the empty-device-set test with only the regression comment that matches its behavior, and restore the replaced-identity description immediately above encrypt_for_devices_reports_replaced_identity so Rust doc binding is correct.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@wacore/src/send/tests.rs`:
- Around line 3482-3491: The doc comment describing the replaced-identity
resolver behavior is attached to the wrong test, so move that comment from
encrypt_for_devices_with_sessions_raw_handles_empty_device_set to the test it
actually documents, encrypt_for_devices_reports_replaced_identity. Keep the
empty-device-set test with only the regression comment that matches its
behavior, and restore the replaced-identity description immediately above
encrypt_for_devices_reports_replaced_identity so Rust doc binding is correct.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: ddcc4d59-94e8-4cf0-a5fd-53d0a202447f
📒 Files selected for processing (2)
wacore/src/send/encrypt.rswacore/src/send/tests.rs
5f4c159 to
44a1f97
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 (1)
wacore/src/send/group.rs (1)
383-428: 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy liftDon’t hold the sender-key chain lock during pairwise SKDM fan-out.
chain_guardcurrently spansencrypt_for_devices_with_sessions, so a large cold group send keeps same-group sends blocked while pairwise device encryption runs. Keep only SKDM creation +encrypt_group_messageunder the lock, then drop before fan-out/stanza assembly.Also applies to: 480-482
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@wacore/src/send/group.rs` around lines 383 - 428, The sender-key chain lock is held too long in group send flow, blocking other same-group sends during pairwise SKDM fan-out. Keep `chain_guard` only around sender-key record creation and `encrypt_group_message`/SKDM preparation in `send/group.rs`, then release it before calling `encrypt_for_devices_with_sessions` and stanza assembly. Adjust the scope near `sender_key_lock`, `create_sender_key_distribution_message_for_group`, and the matching SKDM/send path referenced in the later block so fan-out runs without the chain lock held.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@wacore/src/send/group.rs`:
- Around line 383-428: The sender-key chain lock is held too long in group send
flow, blocking other same-group sends during pairwise SKDM fan-out. Keep
`chain_guard` only around sender-key record creation and
`encrypt_group_message`/SKDM preparation in `send/group.rs`, then release it
before calling `encrypt_for_devices_with_sessions` and stanza assembly. Adjust
the scope near `sender_key_lock`,
`create_sender_key_distribution_message_for_group`, and the matching SKDM/send
path referenced in the later block so fan-out runs without the chain lock held.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: ecc65f00-3ba1-48a5-968a-09dcc142bad2
📒 Files selected for processing (6)
src/client.rssrc/client/lifecycle.rssrc/send/mod.rswacore/src/send/encrypt.rswacore/src/send/group.rswacore/src/send/tests.rs
9dfbd27 to
c8179d3
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c8179d3b78
ℹ️ 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 per-(group,sender) sender-key lock must span SKDM creation, the pairwise SKDM fan-out and the skmsg encrypt: the first and last touch the sender-key chain atomically, and the fan-out mutates shared per-device Signal sessions that the group path (unlike the DM path) does not lock. Plaintext padding is chain-independent though, so compute it before taking the lock, and drop the guard before the stanza build and marshalling instead of holding it for the rest of the function.
filter_skdm_targets is O(devices) per send (three hash lookups per member), which dominates warm group sends to large groups. Memoize the last (resolved-devices, sender-key-device-map) Arc pair that produced an empty needs_skdm; a repeat send observing the same two Arcs is still fully warm and skips the scan. Weak pointers keep the comparison ABA-safe, and either Arc swaps the moment warm state or membership changes, so a stale skip is impossible. Gated on the device memo (otherwise the device Arc is rebuilt each call and the check can never hit).
c8179d3 to
b2ea28b
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@wacore/src/send/group.rs`:
- Around line 383-388: The sender-key chain lock in the group send path only
serializes by sender-key name and does not cover cross-group Signal session
fan-out, so concurrent cold sends can still race the per-device ratchet. Update
the group send flow around the SKDM fan-out in the send/group.rs logic to
acquire the same per-device session lock mechanism used by the DM path, either
by reusing session_locks directly or by threading a lock provider into this
layer, and keep the SKDM creation, fan-out, and encrypt steps under that
serialization.
🪄 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 (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 031bbdfd-4c86-43fe-b6a2-27a8389a2d4a
📒 Files selected for processing (6)
src/client.rssrc/client/lifecycle.rssrc/send/mod.rswacore/src/send/encrypt.rswacore/src/send/group.rswacore/src/send/tests.rs
The fan-out spawned one task + oneshot channel + two store clone_box()es per recipient device, so an 800-member group allocated on the order of 5x800 times per cold send. Encrypt in index-partitioned chunks instead (bounded by ENCRYPT_FANOUT_CONCURRENCY): one task per chunk clones the Arc-backed stores once and encrypts its slice sequentially, so ratchet advances still persist, spawns and clones drop to the concurrency limit, and parallelism is unchanged. Also fixes a divide-by-zero on an empty device set. Covered by new tests for the full fan-out (multi-chunk), a sessionless device being skipped, and the empty device set. Wire order is irrelevant (phash sorts before hashing on both ends).
b2ea28b to
a27e43d
Compare
1 participant
Three self-contained optimizations to the group send path, found by profiling an 802-member LID group send. Behavior is unchanged:
cargo test -p wacore --lib send::is green (95 tests) and a full group-send scenario keepsno sender keyat 0.1. Trim the chain lock to the crypto sections
prepare_group_stanzaheld the per-(group, sender)sender-key lock across the whole function. The lock must span SKDM creation, the pairwise SKDM fan-out and the skmsg encrypt: the first and last touch the sender-key chain atomically, and the fan-out mutates shared per-device Signal sessions that the group path (unlike the DM path) does not lock. Plaintext padding is chain-independent though, so it now runs before the lock, and the guard drops before the stanza build and marshalling instead of being held for the rest of the function.2. Skip the warm SKDM device scan when nothing changed
filter_skdm_targetsis O(devices) per send (three hash lookups per member; ~24M lookups over a 10k-send run to an 802-member group). Memoize the last(resolved-devices, sender-key-device-map)Arcpair that produced an emptyneeds_skdm; a repeat send observing the same twoArcs is still fully warm and skips the scan.Weakpointers keep the comparison ABA-safe (same contract asGroupDevicesMemo), and eitherArcswaps the moment warm state or membership changes, so a stale skip is impossible. Gated on the device memo.3. Chunk the cold SKDM fan-out
The fan-out spawned one task + oneshot channel + two store
clone_box()es per recipient device, so an 800-member group allocated on the order of 5x800 times per cold send (dhat: 123 MB / 45.7k blocks). Encrypt in index-partitioned chunks bounded byENCRYPT_FANOUT_CONCURRENCYinstead: one task per chunk clones the Arc-backed stores once and encrypts its slice sequentially, so ratchet advances still persist while spawns and clones drop to the concurrency limit and parallelism is unchanged. Also fixes a divide-by-zero on an empty device set. Covered by new tests: full multi-chunk fan-out, a sessionless device being skipped, and the empty device set.Measured (802-member cold send): encrypt fan-out allocation 123 MB to 28.8 MB, blocks 45.7k to 4.1k; dhat allocation peak 107 MB to 49 MB.