perf(signal): check out sender keys on load instead of deep-cloning - #880
perf(signal): check out sender keys on load instead of deep-cloning#880jlucaso1 wants to merge 1 commit into
Conversation
load_sender_key returned an owned SenderKeyRecord by deep-cloning the cached Arc (the cache kept its copy, so the Arc was never unique). group_decrypt and group encrypt mutate the record and store it back, so a deep clone of the whole SenderKeyState backlog (a VecDeque of up to MAX_MESSAGE_KEYS message keys per state) ran on every group decrypt — the dominant cost in an offline catch-up / out-of-order backlog (measured ~56% of the worst-case group-decrypt bench). Mirror the session checkout that load_session already uses: get_sender_key stays a read-only Arc peek, and a new checkout_sender_key takes the record out of the cache (leaving the slot CheckedOut) so the cached Arc is unique and the load is a move, not a clone. store_sender_key puts it back. Flush and eviction skip CheckedOut entries (deferred, like sessions); a peek of a checked-out key reports absent, which is safe because every current peek is on the local sending key while the load is on a remote sender's key (audited). Steady-state group messaging (small sender-key state) is unaffected; the win is the large-backlog catch-up path. Tracked by bench_group_out_of_order_decrypt_worst_case on CodSpeed. Tests: checkout lifecycle + peek-during-checkout, cold backend load, and flush-defers-checked-out. 1031 wacore + 827 lib tests pass; clippy --all-targets clean.
|
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 CodeRabbitRelease Notes
WalkthroughRefactors the sender-key in-memory cache from ChangesSender-key tri-state checkout pattern
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
🚥 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)
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 |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9e56264cb4
ℹ️ 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".
| if let Some(entry) = state.cache.get_mut(key) { | ||
| if matches!(entry, SenderKeyEntry::Present(_)) { | ||
| let SenderKeyEntry::Present(record) = | ||
| std::mem::replace(entry, SenderKeyEntry::CheckedOut) |
There was a problem hiding this comment.
Restore checked-out sender keys on errors
When the loaded record is replaced with CheckedOut, there is no guard that puts it back if the caller returns before store_sender_key. That happens on normal group decrypt error paths after load_sender_key, for example when group_decrypt sees an unknown rotated chain id and returns NoSenderKeyState before its final store; after that both get_sender_key and later checkouts treat this entry as absent, so one stale/rotated sender-key message can make the valid cached key disappear for subsequent messages from that sender until the cache is cleared or overwritten.
Useful? React with 👍 / 👎.
📦 Binary size report
.text per crate
Top movers (cargo-bloat attribution)
Baseline: |
There was a problem hiding this comment.
1 issue found across 2 files
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="src/store/signal_adapter.rs">
<violation number="1" location="src/store/signal_adapter.rs:294">
P0: `load_sender_key` now uses checkout semantics globally, but some libsignal call sites only read and do not store back. That can strand sender keys in CheckedOut state and cause follow-up group encryption/decryption loads to fail with missing sender key state.</violation>
</file>
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
| self.0 | ||
| .cache | ||
| .get_sender_key(sender_key_name, &*device.backend) | ||
| .checkout_sender_key(sender_key_name, &*device.backend) |
There was a problem hiding this comment.
P0: load_sender_key now uses checkout semantics globally, but some libsignal call sites only read and do not store back. That can strand sender keys in CheckedOut state and cause follow-up group encryption/decryption loads to fail with missing sender key state.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/store/signal_adapter.rs, line 294:
<comment>`load_sender_key` now uses checkout semantics globally, but some libsignal call sites only read and do not store back. That can strand sender keys in CheckedOut state and cause follow-up group encryption/decryption loads to fail with missing sender key state.</comment>
<file context>
@@ -285,13 +285,14 @@ impl wacore::libsignal::protocol::SenderKeyStore for SenderKeyAdapter {
self.0
.cache
- .get_sender_key(sender_key_name, &*device.backend)
+ .checkout_sender_key(sender_key_name, &*device.backend)
.await
- .map(|opt| opt.map(std::sync::Arc::unwrap_or_clone))
</file context>
|
Closing: the checkout/move approach is unsafe for sender keys (group_decrypt/encrypt have no store-on-error or rollback, and outgoing sends aren't serialized, so a CheckedOut slot leaks on the common NoSenderKeyState error and a local-key peek races a concurrent send into a spurious SKDM). Replacing with a structural-sharing approach (persistent message-key backlog) that keeps the trait/concurrency/error semantics intact. |
What
load_sender_keyreturned an ownedSenderKeyRecordby deep-cloning the cachedArc— the cache kept its copy, so theArcwas never unique andArc::unwrap_or_clonealways cloned.group_decrypt(and group encrypt) mutate the record and store it back, so every group decrypt deep-cloned the entireSenderKeyStatebacklog (aVecDequeof up toMAX_MESSAGE_KEYSmessage keys per state).This mirrors the checkout that sessions'
load_sessionalready does:get_sender_keystays a read-onlyArcpeek (refcount bump, used by the send paths that just check existence).checkout_sender_keytakes the record out of the cache, leaving the slotCheckedOut, so the cachedArcis unique and the load is a move, not a clone.store_sender_keyputs it back.CheckedOutentries (deferred until stored back), exactly like the session path.Why it's safe
It's the same
CheckedOutstate machine sessions already use in production (SessionEntry/get_session). A peek of a checked-out key reports absent — which is safe because every current peek is on the local sending key (send.rs,features/signal.rsSKDM checks) while the load is on a remote sender's key, so they never touch the same cache entry; receives of the same key are serialized per-chat bymessage_enqueue_locks.unwrap_or_clonestill falls back to a clone if a peek'sArcis somehow alive, so correctness never depends on uniqueness. No wire/protocol change.Measurement
The win is concentrated in the offline catch-up / out-of-order backlog path, where the sender-key state is large. A prior bench analysis (#858) put this clone at ~56% of
bench_group_out_of_order_decrypt_worst_case. Steady-state group messaging (small sender-key state) is unaffected. CodSpeed'sbench_group_out_of_order_decrypt_worst_caseandbench_group_decrypttrack the delta on this PR.Tests
New unit tests in
signal_cache.rs: checkout lifecycle + peek/second-checkout-during-checkout report absent; cold load from the backend; and flush defers a checked-out key (not persisted while borrowed, persisted once stored back). The existingwarm_sender_key_hit_shares_arc_not_deep_clone(peek still shares theArc) is unchanged.cargo test -p wacore(1031) and-p whatsapp-rust --lib(827) pass;cargo clippy --all-targetsclean.Provenance
Found while profiling the messaging hot path with CodSpeed and confirmed against the code — the group analog of the session Arc-move, which sender keys hadn't received.