perf(client): trim per-message allocations on the send/receive hot path - #1025
Conversation
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
|
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 (4)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe change separates DB-only and L1 recent-message persistence, reuses protocol-address storage during session-cache probing, and changes message-secret pending entries and batch snapshots to use ChangesRecent-message persistence
Session cache probing
Message-secret buffering
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 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 |
|
| Filename | Overview |
|---|---|
| src/client/sender_keys.rs | Avoids constructing an owned chat-message key when recent messages are stored only in the backend. |
| src/client/sessions.rs | Reuses a protocol-address buffer during the synchronous session-cache filter. |
| src/msg_secret_buffer.rs | Uses shared entries for flush snapshots and preserves recaptured entries through pointer-identity cleanup. |
| src/retry.rs | Adds coverage for storing, retrieving, and consuming a recent message without the in-memory cache. |
Reviews (4): Last reviewed commit: "perf(client): trim per-message allocatio..." | Re-trigger Greptile
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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 (2)
src/client/sessions.rs (1)
301-311: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winUse
.retain()instead of.collect()to reuse theVecallocation.Look, we're building infrastructure to connect billions of people, and every byte of memory counts on our servers. You're doing a great job reducing per-message allocations here, but calling
.collect()still allocates a brand newVec. Since you already ownjids, we can just mutate it in-place using.retain(). This reuses the existing buffer and gets us even closer to our performance goals. Let's make it happen so our infra stays blazing fast.⚡ Proposed refactor
- let mut reusable_addr = wacore::types::jid::make_reusable_protocol_address(); - let jids: Vec<Jid> = jids - .into_iter() - .filter(|jid| { - jid.reset_protocol_address(&mut reusable_addr); - !matches!( - self.signal_cache.try_has_session(&reusable_addr), - Some(true) - ) - }) - .collect(); + let mut reusable_addr = wacore::types::jid::make_reusable_protocol_address(); + let mut jids = jids; + jids.retain(|jid| { + jid.reset_protocol_address(&mut reusable_addr); + !matches!( + self.signal_cache.try_has_session(&reusable_addr), + Some(true) + ) + });🤖 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/client/sessions.rs` around lines 301 - 311, Update the JID filtering logic to mutate the existing jids vector with retain instead of consuming it into a new Vec via into_iter and collect. Preserve the reusable_addr reset and signal_cache.try_has_session filtering behavior while reusing the existing allocation.src/msg_secret_buffer.rs (1)
211-218: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winUse
Arc::ptr_eqfor O(1) comparison on the hot path.Listen, we're building WhatsApp to handle billions of messages, we can't afford to waste CPU cycles! You’ve done the hard work of wrapping these entries in
Arc, so drop the byte-by-byte comparison in theunchangedclosure.Since
insert_pendingalways allocates a newArcfor incoming refreshes,Arc::ptr_eqprovides an exact O(1) match. This guarantees we only remove the exact memory snapshot we just wrote to the backend without the deep equality overhead. Let's keep the hot path blazing fast. Move fast and fix things!⚡ Proposed fix to use pointer equality
- let unchanged = |current: &Arc<MsgSecretEntry>| { - current.secret == entry.secret - && current.expires_at == entry.expires_at - && current.message_ts == entry.message_ts - }; - if pending.get(&key).is_some_and(unchanged) { + if pending.get(&key).is_some_and(|current| Arc::ptr_eq(current, entry)) { pending.remove(&key); }🤖 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/msg_secret_buffer.rs` around lines 211 - 218, Update the `unchanged` closure used by `pending.get(&key).is_some_and(...)` to compare the current entry and `entry` with `Arc::ptr_eq` only, removing the field-by-field equality checks. Preserve the existing behavior of removing `pending[key]` only when both Arcs reference the exact same allocation.
🤖 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 `@src/client/sessions.rs`:
- Around line 301-311: Update the JID filtering logic to mutate the existing
jids vector with retain instead of consuming it into a new Vec via into_iter and
collect. Preserve the reusable_addr reset and signal_cache.try_has_session
filtering behavior while reusing the existing allocation.
In `@src/msg_secret_buffer.rs`:
- Around line 211-218: Update the `unchanged` closure used by
`pending.get(&key).is_some_and(...)` to compare the current entry and `entry`
with `Arc::ptr_eq` only, removing the field-by-field equality checks. Preserve
the existing behavior of removing `pending[key]` only when both Arcs reference
the exact same allocation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 0bfa0851-2e47-452e-be0b-f3d1e9d81efa
📒 Files selected for processing (3)
src/client/sender_keys.rssrc/client/sessions.rssrc/msg_secret_buffer.rs
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
No issues found across 3 files
Confidence score: 5/5
- Automated review surfaced no issues in the provided summaries.
- No files require special attention.
You’re at about 98% of the monthly reviewed-line limit. You may want to disable incremental reviews to conserve quota. Reviews will continue until that limit is exceeded. If you need help avoiding interruptions, please contact contact@cubic.dev.
Auto-approved: Behavior-preserving allocation optimizations on three hot paths, each focused and bounded. The diff confirms no functional change, and Cubic found no implementation issues.
Re-trigger cubic
📦 Binary size report
.text per crate
Top movers (cargo-bloat attribution)
Baseline: |
4b2325d to
0be6477
Compare
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)
src/msg_secret_buffer.rs (1)
203-220: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winUse
retainwithArc::ptr_eqto eliminate string allocations.Look, we're building the infrastructure for the world to communicate. We don't have cycles to waste on useless mallocs. I'm looking at this and you're doing three string clones per entry just to construct a lookup key for the map, followed by a deep comparison of the secret bytes. This is allocating on the hot path just to query a map, which completely defeats the purpose of optimizing this.
Since you've smartly wrapped the entries in
Arcs, every new queue or refresh allocates a newArc(even if the metadata was merely merged). You can completely eliminate the string allocations and deep byte comparisons by usingHashMap::retainwith pointer equality (Arc::ptr_eq). For realistic batch sizes, the CPU cache will burn through a linear pointer scan infinitely faster than allocating strings. Get this shipped.💡 Proposed fix
- fn finish_batch(&self, written: &[Arc<MsgSecretEntry>]) { - let mut pending = self.pending.lock().unwrap_or_else(|p| p.into_inner()); - for entry in written { - let key = ( - entry.chat.clone(), - entry.sender.clone(), - entry.msg_id.clone(), - ); - let unchanged = |current: &Arc<MsgSecretEntry>| { - current.secret == entry.secret - && current.expires_at == entry.expires_at - && current.message_ts == entry.message_ts - }; - if pending.get(&key).is_some_and(unchanged) { - pending.remove(&key); - } - } - } + fn finish_batch(&self, written: &[Arc<MsgSecretEntry>]) { + let mut pending = self.pending.lock().unwrap_or_else(|p| p.into_inner()); + pending.retain(|_, current| { + !written.iter().any(|w| Arc::ptr_eq(current, w)) + }); + }🤖 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/msg_secret_buffer.rs` around lines 203 - 220, Update finish_batch to remove matching entries via pending.retain, comparing each retained Arc with the entries in written using Arc::ptr_eq. Eliminate the cloned chat, sender, and msg_id lookup key and the deep secret/metadata comparison, while preserving removal only for entries whose Arc pointer appears in written.
🤖 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/client/sender_keys.rs`:
- Around line 378-382: Update the store_sent_message contract and every
implementation/call site to accept and forward the existing Arc<Vec<u8>>
payload, including the persistence backend and SQLite retry handling, instead of
coercing it to a slice and recreating the buffer with to_vec. Preserve shared
ownership across the DB-only and L1 cache paths without deep-copying message
payloads.
---
Outside diff comments:
In `@src/msg_secret_buffer.rs`:
- Around line 203-220: Update finish_batch to remove matching entries via
pending.retain, comparing each retained Arc with the entries in written using
Arc::ptr_eq. Eliminate the cloned chat, sender, and msg_id lookup key and the
deep secret/metadata comparison, while preserving removal only for entries whose
Arc pointer appears in written.
🪄 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: 2e92328a-19af-4ffd-8f87-6529ee87ec6f
📒 Files selected for processing (3)
src/client/sender_keys.rssrc/client/sessions.rssrc/msg_secret_buffer.rs
0be6477 to
b6e3297
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
src/client/sender_keys.rs (1)
378-382: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftRefactor upstream
store_sent_messageto takeArc<Vec<u8>>and eliminate the payload deep copy.Look, we're out here optimizing string allocations on this DB-only path to keep WhatsApp fast, but we're completely missing a massive bottleneck right below this layer.
You've got an
Arc<Vec<u8>>right here inshared, but you're coercing it down to a&[u8]on both the DB-only and L1 cache paths to satisfy thestore_sent_messagetrait. If you look at our SQLite backend, it immediately does.to_vec()to handle retries! We are needlessly deep-copying the entire message payload on every single message we send across the network.I need things to work right, and burning CPU cycles to clone buffers we already own isn't how we scale to billions of users. We need to update the
store_sent_messagecontract inwacoreto take anArc<Vec<u8>>orbytes::Bytesso we can pass this right through without a copy. Let's get this fixed across the stack.🤖 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/client/sender_keys.rs` around lines 378 - 382, Update the store_sent_message contract in wacore and every backend/call site to accept and propagate Arc<Vec<u8>> (or the project’s established equivalent) instead of a borrowed byte slice. Preserve the shared Arc from sender_keys.rs through the SQLite backend and retry handling, removing the immediate to_vec/deep-copy conversion while retaining existing storage behavior.
🤖 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/client/sessions.rs`:
- Around line 305-308: In the session reuse condition, replace the verbose
!matches!(self.signal_cache.try_has_session(&reusable_addr), Some(true))
expression with a direct Option<bool> comparison against Some(true), preserving
the existing negated condition and behavior.
In `@src/msg_secret_buffer.rs`:
- Around line 205-215: Update finish_batch to build a HashSet of the written Arc
pointer identities before acquiring the pending mutex, then use constant-time
membership checks inside pending.retain instead of scanning written with
iter().any(). Preserve Arc pointer-identity semantics and keep lock acquisition
after the set construction.
---
Duplicate comments:
In `@src/client/sender_keys.rs`:
- Around line 378-382: Update the store_sent_message contract in wacore and
every backend/call site to accept and propagate Arc<Vec<u8>> (or the project’s
established equivalent) instead of a borrowed byte slice. Preserve the shared
Arc from sender_keys.rs through the SQLite backend and retry handling, removing
the immediate to_vec/deep-copy conversion while retaining existing storage
behavior.
🪄 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: 79dff2d2-653c-486d-8e7c-2f95eae75b13
📒 Files selected for processing (4)
src/client/sender_keys.rssrc/client/sessions.rssrc/msg_secret_buffer.rssrc/retry.rs
Three independent, backend-agnostic, behavior-preserving allocation cuts on the pingpong steady-state path: - ensure_sessions_inner: reuse one ProtocolAddress across the warm-cache pre-filter (retain in place) instead of allocating a fresh signal address per jid. A plain local, so the concurrent probe stream keeps its own per-task address. - add_recent_message: in DB-only mode (no L1 recent-message cache) only the chat and id are borrowed, so resolve the chat directly and pass the caller's borrowed id instead of materializing make_chat_message_id's owned ChatMessageId, whose id.to_owned() would just be borrowed away. - msg_secret_buffer: store Arc<MsgSecretEntry> in the pending map so the flush snapshot is a refcount bump instead of a deep clone of every entry (the owned Vec for put_msg_secrets is unwrapped once, was two full-batch deep clones), and key finish_batch removal on Arc pointer identity, dropping the per-entry (chat, sender, id) key rebuild and secret comparison. insert_pending always stores a fresh Arc so a recapture is a distinct pointer; added tests pin that invariant and the edit-recapture survival, plus a DB-only recent-message round-trip. Measured with dhat over a 100k-message pingpong flood: -8.2 allocations per message (229.2 -> 221.0), -0.57 KB/msg. Small in isolation; allocation hygiene on a path that runs once per message.
b6e3297 to
99902fd
Compare
1 participant
Three independent, backend-agnostic, behavior-preserving allocation cuts on the
pingpong steady-state path. Each removes work that runs once per message.
Changes
ensure_sessions_inner(client/sessions.rs): the warm-cache pre-filterallocated a fresh
ProtocolAddressper jid just to use it as a lookup key.Reuse a single one, rewritten in place with
reset_protocol_address(the samepattern already used in
send/encrypt.rs). It is a plain local, so theconcurrent probe stream below keeps its own per-task address. Saves 1 alloc/msg
at fanout 2 and scales as N-1 with device fanout.
add_recent_message(client/sender_keys.rs): in DB-only mode (no L1recent-message cache) only the chat and id are borrowed, yet the code built the
owned
ChatMessageIdfrommake_chat_message_idup front. Resolve the chatdirectly and pass the caller's borrowed
id, so theid.to_owned()that wouldonly be borrowed away is never allocated. The chat resolution
(
resolve_encryption_jid) is preserved, so PN/LID normalization is unchanged.msg_secret_buffer(msg_secret_buffer.rs): the write-behind buffer storedMsgSecretEntryby value, so the flush snapshot (values().cloned()) deep-clonedevery pending entry (3 Strings + a Vec each), the batch was cloned again for the
backend put, and
finish_batchrebuilt the(chat, sender, id)key anddeep-compared the secret for every written entry (3 allocations/msg). Store
Arc<MsgSecretEntry>instead: the snapshot is a refcount bump, the owned Vec forput_msg_secretsis unwrapped once, andfinish_batchkeys removal onArcpointer identity (
retain+ptr_eq), dropping the key rebuild and the secretcomparison entirely.
insert_pendingalways stores a freshArc, so a recaptureduring a flush is a distinct pointer that the in-flight batch cannot evict; the
invariant is documented at both sites and pinned by new tests.
Measurement
dhat, 100k-message pingpong flood, memory backend:
Small in isolation. This is allocation hygiene, not a throughput play; the effect
is below the noise floor of a throughput measurement at this rate.
Tests
New tests, all green:
msg_secret_buffer::recapture_survives_finish_batch_of_real_snapshot-- drivesthe real drain flow (snapshot the actual pending
Arc, recapture, thenfinish_batchthe stale snapshot) and asserts the recapture survives and reachesthe backend.
msg_secret_buffer::identical_recapture_is_a_distinct_arc-- pins theinsert_pending"always a freshArc" invariantfinish_batchrelies on: even abyte-identical recapture is a distinct allocation, so pointer identity (not
content) decides removal. Would fail under the old field comparison.
retry::recent_message_db_only_round_trip-- exercises the capacity-0 (DB-only)add_recent_messagepath end to end, asserting the store + retrieve round-tripsthrough the backend under the resolved chat.
Validation
cargo fmt --all --check,cargo clippy --all-targets -- -D warnings, andcargo test --workspace --exclude e2e-testsall pass (991 lib tests).