Skip to content

perf(signal): check out sender keys on load instead of deep-cloning - #880

Closed
jlucaso1 wants to merge 1 commit into
mainfrom
perf/sender-key-checkout
Closed

perf(signal): check out sender keys on load instead of deep-cloning#880
jlucaso1 wants to merge 1 commit into
mainfrom
perf/sender-key-checkout

Conversation

@jlucaso1

@jlucaso1 jlucaso1 commented Jun 15, 2026

Copy link
Copy Markdown
Collaborator

What

load_sender_key returned an owned SenderKeyRecord by deep-cloning the cached Arc — the cache kept its copy, so the Arc was never unique and Arc::unwrap_or_clone always cloned. group_decrypt (and group encrypt) mutate the record and store it back, so every group decrypt deep-cloned the entire SenderKeyState backlog (a VecDeque of up to MAX_MESSAGE_KEYS message keys per state).

This mirrors the checkout that sessions' load_session already does:

  • get_sender_key stays a read-only Arc peek (refcount bump, used by the send paths that just check existence).
  • 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 until stored back), exactly like the session path.

Why it's safe

It's the same CheckedOut state 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.rs SKDM 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 by message_enqueue_locks. unwrap_or_clone still falls back to a clone if a peek's Arc is 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's bench_group_out_of_order_decrypt_worst_case and bench_group_decrypt track 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 existing warm_sender_key_hit_shares_arc_not_deep_clone (peek still shares the Arc) is unchanged. cargo test -p wacore (1031) and -p whatsapp-rust --lib (827) pass; cargo clippy --all-targets clean.

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.

Review in cubic

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.
@coderabbitai

coderabbitai Bot commented Jun 15, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 14c97d01-08b4-4b6d-bca7-710ca7fc98f0

📥 Commits

Reviewing files that changed from the base of the PR and between a03672e and 9e56264.

📒 Files selected for processing (2)
  • src/store/signal_adapter.rs
  • wacore/src/store/signal_cache.rs

📝 Walkthrough

Summary by CodeRabbit

Release Notes

  • Refactor

    • Optimized internal sender-key cache management to improve handling of concurrent operations and resource efficiency.
  • Tests

    • Added tests to validate sender-key caching behavior and concurrent access patterns.

Walkthrough

Refactors the sender-key in-memory cache from Option<Arc<SenderKeyRecord>> to a tri-state SenderKeyEntry (Present, Absent, CheckedOut). A new checkout_sender_key method returns an owned record and marks the entry as CheckedOut; eviction skips checked-out entries; flush defers persistence for CheckedOut entries. load_sender_key in the signal adapter is updated to use checkout_sender_key.

Changes

Sender-key tri-state checkout pattern

Layer / File(s) Summary
SenderKeyEntry tri-state type and store state
wacore/src/store/signal_cache.rs
Introduces SenderKeyEntry enum (Present, Absent, CheckedOut) and updates SenderKeyStoreState to store these variants; rewrites put and delete to write the correct variant.
Eviction, get_sender_key, and checkout_sender_key
wacore/src/store/signal_cache.rs
Rewrites evict_if_needed to skip CheckedOut entries; rewrites get_sender_key to do backend I/O outside the mutex and treat CheckedOut as absent; adds the new checkout_sender_key method that transitions PresentCheckedOut and returns an owned SenderKeyRecord.
Flush persistence with CheckedOut deferral
wacore/src/store/signal_cache.rs
Updates flush to serialize Present, delete Absent from backend, and skip CheckedOut entirely; CheckedOut keys remain dirty until stored back.
Signal adapter wired to checkout_sender_key
src/store/signal_adapter.rs
Updates load_sender_key to call cache.checkout_sender_key(...) instead of cache.get_sender_key(...), with revised comments reflecting the checkout/store-back ownership pattern.
Tests for checkout semantics and flush deferral
wacore/src/store/signal_cache.rs
Adds tokio tests covering checkout take-out/restore, cold backend loading, and flush deferral while a sender key is CheckedOut.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

  • oxidezap/whatsapp-rust#713: Directly modifies sender-key cache/adapter ownership semantics in the same two files, changing how load_sender_key returns an owned record — the immediate predecessor to this checkout pattern.
  • oxidezap/whatsapp-rust#503: Modifies SenderKeyStore::load_sender_key in src/store/signal_adapter.rs — the exact same method this PR changes to use checkout_sender_key.
  • oxidezap/whatsapp-rust#326: Modifies signal_cache.rs sender-key caching internals — the same cache code this PR refactors to the tri-state model.

Suggested labels

performance

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately reflects the main optimization: replacing deep-cloning with a checkout mechanism for sender keys to improve performance.
Description check ✅ Passed The description comprehensively covers the what, why, safety rationale, performance measurement, tests, and provenance—all directly related to the changeset.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch perf/sender-key-checkout

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

@github-actions

Copy link
Copy Markdown

📦 Binary size report

Metric main PR Δ
bin size (stripped) 10.61 MiB 10.61 MiB +3.66 KiB (+0.03%) 🔺
bin .text 8.74 MiB 8.74 MiB +3.69 KiB (+0.04%) 🔺
bin allocated (text+data+bss) 10.61 MiB 10.61 MiB +3.99 KiB (+0.04%) 🔺
llvm-lines wacore 645,616 645,617 +1 (+0.00%) 🔺
llvm-lines wacore copies 17,669 17,668 -1 (-0.01%) 🔽
llvm-lines whatsapp-rust lib 657,082 657,739 +657 (+0.10%) 🔺
llvm-lines whatsapp-rust lib copies 20,006 20,016 +10 (+0.05%) 🔺
deps crates (Cargo.lock) 354 354 0
.text per crate
Crate main PR Δ
.text whatsapp_rust 1.48 MiB 1.49 MiB +7.93 KiB (+0.52%) 🔺
.text wacore 544.63 KiB 540.78 KiB -3.85 KiB (-0.71%) 🔽
.text wacore_binary 158.83 KiB 158.83 KiB 0
.text wacore_libsignal 170.47 KiB 170.47 KiB 0
.text wacore_appstate 35.26 KiB 35.26 KiB 0
.text wacore_noise 30.68 KiB 30.68 KiB 0
.text waproto 895.34 KiB 895.34 KiB 0
.text whatsapp_rust_sqlite_storage 206.21 KiB 206.21 KiB 0
.text whatsapp_rust_tokio_transport 33.09 KiB 33.09 KiB 0
.text whatsapp_rust_ureq_http_client 6.19 KiB 6.19 KiB 0
.text std 1.13 MiB 1.14 MiB +228 B (+0.02%) 🔺
.text other deps 4.02 MiB 4.02 MiB -618 B (-0.01%) 🔽
Top movers (cargo-bloat attribution)
Crate main PR Δ
whatsapp_rust 1.48 MiB 1.49 MiB +7.93 KiB (+0.52%)
wacore 544.63 KiB 540.78 KiB -3.85 KiB (-0.71%)
prost 466.39 KiB 464.26 KiB -2.13 KiB (-0.46%)
rustix 131 B 1.41 KiB +1.28 KiB (+1002.29%)

Baseline: a03672ee5 (latest main run) · Head: 22534de47 · Graphs

@jlucaso1
jlucaso1 marked this pull request as draft June 15, 2026 23:56

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>

@codspeed-hq

codspeed-hq Bot commented Jun 15, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 174 untouched benchmarks


Comparing perf/sender-key-checkout (9e56264) with main (a03672e)

Open in CodSpeed

@jlucaso1

Copy link
Copy Markdown
Collaborator Author

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.

@jlucaso1 jlucaso1 closed this Jun 16, 2026
@jlucaso1
jlucaso1 deleted the perf/sender-key-checkout branch July 3, 2026 02:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant