perf(send): establish sessions before taking the sender-key chain lock - #807
Merged
Conversation
prepare_group_stanza held the per-(group, sender) chain lock across the whole SKDM path: device resolution and the prekey fetch + X3DH inside encrypt_for_devices — network round-trips — ran inside the critical section, so concurrent sends to the same group serialized behind an RTT (or its retries) whenever any target device lacked a session. Split encrypt_for_devices into its two halves: ensure_sessions_for_devices (network: LID-first lookup, prekey fetch, parallel X3DH; touches only session/identity state) and encrypt_for_devices_with_sessions (CPU: the pairwise fan-out). The combined function remains as a composition for the DM path. The group path now resolves devices and ensures sessions before the lock; the lock covers only SKDM creation + pairwise encrypt + skmsg — the chain-consistency invariant it exists for. Failure semantics preserved: session-setup errors log and continue without distribution (WA Web GroupSkmsgJob), and the sender-key record is still created under the lock so the skmsg always encrypts. The regression test probes the actual chain lock from inside the mock resolver's fetch; it fails against the previous code. Test mem stores now share state across clones (Arc), matching production store semantics so spawned-task session writes are visible.
📝 WalkthroughWalkthroughThis PR refactors group message encryption to serialize per-group session setup (prekey fetch and X3DH) separately from the sender-key chain lock. The changes split encryption into session-preparation and encryption phases, move session establishment outside the chain critical section, and add comprehensive lock-ordering tests to validate the concurrency improvements. ChangesSKDM Session Setup Concurrency
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: fad388adea
ℹ️ 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".
Benchmark Results67 unchanged benchmark(s)
|
Review follow-up: hoisting ensure_sessions_for_devices out of the chain lock let two concurrent cold sends to the same group race prekey fetch + X3DH writes to the same per-device sessions (previously serialized as a side effect of the chain lock). Add SenderKeyStore::session_setup_lock — a per-group lock held only across session setup, never the chain critical section. Cold same-group sends serialize their setup again; warm sends never take it, so the chain lock stays network-free. Default impl is uncontended (mirrors sender_key_lock); the signal cache shares the chain-lock map under a disjoint '::setup'-suffixed key. The regression test now also asserts the fetch runs UNDER the setup lock, alongside the existing not-under-chain-lock assertion.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
3 tasks
Merged
This was referenced Jul 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
prepare_group_stanzatakes the per-(group, sender) chain lock at the top and holds it through the entire SKDM path. That puts two network phases inside the critical section:resolver.resolve_devices), andencrypt_for_devices(fetch_prekeys_for_identity_check).Any send to a group with at least one session-less target device holds the chain lock across a server round-trip, so concurrent sends to the same group serialize behind that RTT (and its failure/retry time). The invariant the lock exists for — the SKDM and the skmsg must describe the same chain iteration — only requires it to cover SKDM creation and the skmsg encrypt, both CPU-bound.
Change
encrypt_for_devicesis split into its two existing halves:ensure_sessions_for_devices— the network half: LID-first session lookup, batch prekey fetch, parallel X3DH. Touches only session/identity state, never a sender-key chain. Returns aSessionPlan(the per-device LID encryption overrides + the 406 flag).encrypt_for_devices_with_sessions— the CPU half: the bounded pairwise encrypt fan-out, consuming theSessionPlan. Safe to run under locks that must not span I/O.encrypt_for_devicesremains as the composition of the two, so the DM path is unchanged. The group path now resolves devices and ensures sessions before taking the chain lock; the lock covers only SKDM creation + the pairwise fan-out + the skmsg encrypt. This also matches WA Web, whereensureE2ESessionsis a separate step beforeGroupSkmsgJob's encrypt.Session-setup serialization (review follow-up): hoisting setup out of the chain lock would have let two concurrent cold sends to the same group race prekey fetch + X3DH writes to the same per-device sessions — serialization the chain lock previously provided as a side effect. A new
SenderKeyStore::session_setup_lock(per-group, default-uncontended likesender_key_lock; the signal cache shares the chain-lock map under a disjoint::setupkey) is held only acrossensure_sessions_for_devices. Same-group cold sends serialize their setup exactly as onmain; warm sends never take it, so the chain lock stays network-free. Cross-group/-path races over a shared device are unchanged frommain(the chain lock never covered those).Failure semantics are preserved: a session-setup error logs and continues without SKDM distribution (WA Web's try/catch-without-rethrow rule),
is_device_unregistered_errorstill flips the stale-device flag, and the sender-key record is still created under the chain lock even when distribution is skipped — so the skmsg always encrypts, exactly as before (previously the record creation happened before the failing encrypt call; skipping it would have turned a distribution failure into a hardNoSenderKeyStatesend failure).Verification
prekey_fetch_runs_outside_chain_lock: the mock resolver probes the store's actual lock instances from insidefetch_prekeys_for_identity_checkand asserts both directions — the fetch runs not under the chain lock and under the per-group setup lock — plus end-to-end that the pre-established session produces B's pairwise SKDM under the chain lock. The not-under-chain-lock assertion fails against the previousgroup.rs.Arc<Mutex<…>>), matching production stores ("the shared cache provides interior mutability") — without this, sessions established inside the spawned X3DH tasks were silently lost in tests, which made the end-to-end half impossible to express.cargo test -p wacore --lib: 947 passed.cargo test -p whatsapp-rust --lib: 744 passed.cargo clippy -p wacore -p wacore-libsignal -p whatsapp-rust --testsclean; fmt clean.Breaking
None for callers:
encrypt_for_devicesandprepare_group_stanzakeep their signatures.ensure_sessions_for_devices/encrypt_for_devices_with_sessions/SessionPlanand the defaulted trait methodSenderKeyStore::session_setup_lockare new pub items. Tracing topology shifts slightly: the DM path now emitswa.send.ensure_sessions+wa.send.encrypt_fanoutinstead of one span covering both.