perf(store): cache the device snapshot as Arc<Device> - #808
Merged
Conversation
get_device_snapshot() cloned the whole Device on every call — Jids, push_name, props_hash, edge_routing_info and nct_salt heap allocations — and it is called on every inbound message plus ~70 other sites. The snapshot is now an Arc<Device> rebuilt under the device write guard in modify_device (the single mutation funnel), so reads become a refcount bump with no lock against writers and no clone, and can never observe a stale snapshot relative to a committed mutation. Fallout, in the same spirit of avoiding needless work: - get_device_snapshot() and the get_pn/get_lid/get_push_name/require_pn accessors are now sync (no async lock dance); call sites updated. - Read-only get_device_arc() users (retry, sessions, prekeys, signal feature, send force-skdm checks, receive parse_message_info) switch to the snapshot — each was paying an RwLock read just to reach .backend or clone pn/lid. - Send/receive paths borrow pn/lid/account from the held snapshot instead of cloning them; clones remain only where ownership leaves the scope. - Tests that mutated Device through the raw write lock now go through modify_device — direct mutation would bypass the cached snapshot, so the docs now state that invariant explicitly (AGENTS.md, agent_docs). Breaking (pre-1.0): get_device_snapshot returns Arc<Device> and is no longer async; the three accessors are no longer async; get_device_arc is documented as store-adapter-only.
|
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 (7)
📝 WalkthroughSummary by CodeRabbit
WalkthroughConvert PersistenceManager to expose a cached Arc via a synchronous get_device_snapshot(), make client identity getters synchronous, refactor runtime/send/session/prekey/retry flows to use borrowed snapshot fields, and update tests and docs to the new mutation/read conventions. ChangesSynchronous Device Snapshot Access Refactoring
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 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)
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 |
Benchmark Results67 unchanged benchmark(s)
|
![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: 23ef08c178
ℹ️ 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 borrowed pn/lid bindings made the explicit &x at call sites a double reference (clippy needless_borrow under -D warnings). Review follow-up: set the dirty flag immediately after the modifier, before the snapshot clone — a shutdown flush racing that window checked dirty, saw clean, and could exit without persisting the committed mutation. The flush's device read lock still serializes against the held write guard, so it always saves post-mutation state.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/message/tests.rs (1)
1753-1765: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winExtract a shared identity-seeding helper for these tests.
You repeat the same
pm.modify_devicesetup four times. Pulling this into one helper will keep these tests easier to evolve and reduces drift risk when device identity setup changes again.Also applies to: 2452-2464, 2550-2562, 2644-2656
🤖 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 `@src/message/tests.rs` around lines 1753 - 1765, Extract the repeated pm.modify_device block into a single async test helper (e.g., seed_test_device_identity or setup_test_device_identity) that takes &pm (or &mut PM test handle) and performs the device.pn and device.lid assignments currently done inside pm.modify_device; replace each duplicated inline pm.modify_device({...}).await with a call to that helper from the four test sites (the occurrences around lines 1753, 2452, 2550, 2644) so tests call await seed_test_device_identity(&pm). Ensure the helper is async and lives in the same test module (or a common tests::helpers mod) so it can be reused by all tests.src/features/events.rs (1)
133-141: 🧹 Nitpick | 🔵 Trivial | 💤 Low valueConsider removing
asyncfromresolve_responder_jid.Look, this method doesn't await anything anymore since
get_lid()is now synchronous. Theasynckeyword here is just dead weight. Not a blocker, but you should clean this up - we don't ship unnecessary complexity at Meta.♻️ Suggested cleanup
- async fn resolve_responder_jid(&self, event_creator_jid: &Jid, own_pn: &Jid) -> Jid { + fn resolve_responder_jid(&self, event_creator_jid: &Jid, own_pn: &Jid) -> Jid {And update the call site at line 86-88:
let responder = self .resolve_responder_jid(event_creator_jid, &my_base) - .await; + ;🤖 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 `@src/features/events.rs` around lines 133 - 141, The function resolve_responder_jid no longer needs to be async because it does not await anything; change its signature to remove async (e.g., fn resolve_responder_jid(&self, event_creator_jid: &Jid, own_pn: &Jid) -> Jid) and keep the same body (calling self.client.get_lid()). Then update all call sites that currently await this method to call it synchronously (remove .await and handle the returned Jid directly); also update any trait signatures or impls if they declared the method as async so the types remain consistent.
🤖 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 `@src/message/tests.rs`:
- Line 1040: Tests call client.persistence_manager.get_device_snapshot()
multiple times within the same JID resolution fallback chain (e.g., when
resolving lid/pn), which can produce flaky behavior; fix by calling
get_device_snapshot() once at the start of the resolution path, store the
returned Arc<Device> in a local variable (e.g., snapshot) and then use
snapshot.lid.clone() / snapshot.pn.clone() (or equivalent fields) for all
fallback .or(...) branches instead of re-invoking get_device_snapshot(); update
all occurrences in the test (the chains currently calling get_device_snapshot()
repeatedly) to reference that single snapshot variable so the resolution uses a
consistent point-in-time Device.
In `@src/request.rs`:
- Around line 91-95: The public method generate_message_id in impl Request is
unnecessarily async because get_device_snapshot is now synchronous and there are
no await points; change generate_message_id to a synchronous fn returning String
(remove async) and update its signature and any internal .await usage, keeping
the same body calling self.persistence_manager.get_device_snapshot() and
self.get_request_utils().generate_message_id(...); then update all callers to
remove `.await` when calling Request::generate_message_id so call sites compile
cleanly.
---
Outside diff comments:
In `@src/features/events.rs`:
- Around line 133-141: The function resolve_responder_jid no longer needs to be
async because it does not await anything; change its signature to remove async
(e.g., fn resolve_responder_jid(&self, event_creator_jid: &Jid, own_pn: &Jid) ->
Jid) and keep the same body (calling self.client.get_lid()). Then update all
call sites that currently await this method to call it synchronously (remove
.await and handle the returned Jid directly); also update any trait signatures
or impls if they declared the method as async so the types remain consistent.
In `@src/message/tests.rs`:
- Around line 1753-1765: Extract the repeated pm.modify_device block into a
single async test helper (e.g., seed_test_device_identity or
setup_test_device_identity) that takes &pm (or &mut PM test handle) and performs
the device.pn and device.lid assignments currently done inside pm.modify_device;
replace each duplicated inline pm.modify_device({...}).await with a call to that
helper from the four test sites (the occurrences around lines 1753, 2452, 2550,
2644) so tests call await seed_test_device_identity(&pm). Ensure the helper is
async and lives in the same test module (or a common tests::helpers mod) so it
can be reused by all tests.
🪄 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: 33cff748-7b2a-467a-afc7-670ecd56a162
📒 Files selected for processing (52)
AGENTS.mdagent_docs/feature_implementation.mdsrc/bot.rssrc/client/accessors.rssrc/client/app_state.rssrc/client/device_registry.rssrc/client/iq_ops.rssrc/client/lifecycle.rssrc/client/messaging.rssrc/client/node_io.rssrc/client/sender_keys.rssrc/client/sessions.rssrc/features/contacts.rssrc/features/events.rssrc/features/media_reupload.rssrc/features/polls.rssrc/features/presence.rssrc/features/signal.rssrc/handlers/call.rssrc/handlers/notification/device.rssrc/handshake.rssrc/history_sync.rssrc/message/msg_secret.rssrc/message/receive.rssrc/message/special.rssrc/message/tests.rssrc/pair.rssrc/pair_code.rssrc/pdo.rssrc/prekeys.rssrc/receipt.rssrc/request.rssrc/retry.rssrc/send.rssrc/store/persistence_manager.rssrc/usync.rssrc/version.rstests/e2e/src/lib.rstests/e2e/tests/app_state.rstests/e2e/tests/chat_actions.rstests/e2e/tests/community.rstests/e2e/tests/digest_key.rstests/e2e/tests/groups.rstests/e2e/tests/lid_sessions.rstests/e2e/tests/media.rstests/e2e/tests/memory_soak.rstests/e2e/tests/prekey_sessions.rstests/e2e/tests/privacy_tokens.rstests/e2e/tests/profile.rstests/e2e/tests/profile_picture.rstests/e2e/tests/session_reuse.rstests/handshake_integration.rs
💤 Files with no reviewable changes (6)
- tests/e2e/tests/community.rs
- tests/e2e/tests/prekey_sessions.rs
- tests/e2e/tests/profile_picture.rs
- tests/e2e/tests/groups.rs
- tests/e2e/tests/chat_actions.rs
- tests/e2e/tests/media.rs
…test chain Review follow-ups: generate_message_id lost its only await point with the sync snapshot, so drop the async (call sites updated). The test JID fallback chains re-read the snapshot mid-chain — besides defeating the point-in-time contract, .or(...) evaluated the second snapshot + clone eagerly; they now take one snapshot and borrow.
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
get_device_snapshot()deep-cloned the wholeDeviceon every call (persistence_manager.rs:66-68): twoJids,push_name,props_hash,edge_routing_info(up to ~KB),nct_salt— ~5-8 heap allocations per call. It has ~72 call sites, including once per inbound message (parse_message_inforeads pn/lid) and throughout the send/receipt paths. Separately, many read-only paths paid anRwLockread on the device arc just to reach.backendor clone two fields, and theget_pn/get_lid/get_push_nameaccessors did an async lock dance per call.Change
PersistenceManagerkeeps adevice_snapshot: std::sync::RwLock<Arc<Device>>, rebuilt under the device write guard insidemodify_device— the single mutation funnel (process_commandand every directmodify_devicecaller go through it; the signal-store write locks only delegate to the backend and never mutate Device fields, verified across allimpl *Store for Device). Reads can therefore never observe a stale snapshot relative to a committed mutation, and the rebuild cost is one Device clone per mutation (rare: pairing, push-name sync, prekey counter) instead of one per read.get_device_snapshot()is now sync and returnsArc<Device>— a refcount bump, no lock against writers, no clone.get_pn/get_lid/get_push_name/require_pnare now sync via the snapshot; the owned return is the only clone left.get_device_arc()sites (retry base-key/session checks,ensure_sessions_inner, prekey upload paths, signal feature checks, send force-SKDM checks,parse_message_info) switch to the snapshot.parse_message_infonow borrow pn/lid/account from the held snapshot instead of cloning (the old code cloned the whole Device precisely so it could.take()two fields); clones remain only where ownership escapes the scope (e.g. PDO'sMessageInfo.sender).get_device_arc()stays, documented as store-adapter-only (&mut Devicetrait access).Invariant made explicit
Direct Device mutation through the raw write lock now bypasses the cached snapshot — it was already against the documented convention, but four tests in
message/tests.rsdid it; they now usemodify_device.AGENTS.md(State convention) andagent_docs/feature_implementation.mdupdated: mutations viaDeviceCommand/modify_deviceonly (tests included), reads via the cheap snapshot, borrow over clone.Verification
Full workspace
cargo test: 2087 passed, 0 failed (includes e2e compile viacargo check -p e2e-tests --tests).cargo clippy --all --testsclean, fmt clean. Diff: 52 files, +254/−365.Breaking (pre-1.0)
get_device_snapshot():async fn → Devicebecomesfn → Arc<Device>.get_push_name()/get_pn()/get_lid(): no longerasync.Call-site migration is mechanical: drop
.await, and either deref fields (snapshot.pn.as_ref()) or clone where ownership is needed.Generated by Claude Code