perf: cut a dozen per-message allocations from the DM round trip - #1122
Conversation
Every one of these is a buffer or string the hot path built and then either threw away unused or copied straight into its real home. Message-secret rows. `MsgSecretEntry::new` is now the single place the three `Arc<str>` identifiers are derived, replacing the same three lines copied into the inbound capture, the outbound persist and the history-sync collector. It renders each JID straight into its `Arc<str>` via a new `Jid::to_non_ad_arc_str` (the old `to_non_ad_string().into()` paid for an intermediate `String` it immediately copied out of), and it shares one allocation between `chat` and `sender` when both name the same user, which is every direct message. A capture with no bot alias now takes the buffer's existing single-entry path instead of a two-slot Vec that only ever held one. Enc classification. The three per-kind payload buckets were each reserved to the stanza's enc count up front; a stanza's enc nodes are overwhelmingly all one kind, so two of the three buffers were allocated and stayed empty. They are now reserved on first push, which leaves a mixed stanza at exactly the allocation count it had before. Encrypt fan-out. `SessionPlan` carried an all-`None` override slot per device plus a prekey index list, both allocated on every send and both empty on a warm one. Both are materialized only when something is actually recorded; the device count moves onto the plan so the "built for a different device list" assertion stays exact. Stanza building. The DM content vector reserved one slot and then reallocated for each of `<device-identity>`, `<reporting>` and the caller's extra nodes, memcpying the whole `Node` values each time. The reporting token's version attribute went through `i32::to_string()`, where handing `NodeValue` the integer formats it into an inline `CompactString`. And `edit` is parsed from the borrowed attribute, since `EditAttribute::from(String)` re-borrows it anyway and drops the String for every known variant.
|
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 (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR centralizes ChangesMessage-secret canonicalization and persistence
Lazy payload and encryption storage
Wire and stanza allocation paths
Estimated code review effort: 4 (Complex) | ~45 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/message/msg_secret.rs | Avoids a one-element batch allocation while preserving the atomic two-alias persistence path. |
| src/message/receive.rs | Allocates encryption payload buckets only when their first payload is classified. |
| wacore/binary/src/jid.rs | Adds a stack-rendered Arc form that remains byte-identical to the existing canonical non-AD representation. |
| wacore/src/store/traits.rs | Centralizes message-secret entry construction and safely shares identifiers for equivalent chat identities. |
| wacore/src/send/encrypt.rs | Lazily materializes encryption override and prekey-index vectors while retaining device-count tracking. |
| wacore/src/send/dm.rs | Pre-sizes the message-content node vector for all optional and caller-provided children. |
| wacore/src/messages.rs | Parses edit attributes from borrowed text without changing known or unknown value handling. |
| wacore/src/reporting_token.rs | Uses the integer NodeValue conversion directly while preserving the reporting-token wire value. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart LR
JID[JID inputs] --> Canon[Canonical non-AD identifiers]
Canon --> Secret[Message-secret entry]
Secret --> Buffer[Write-behind buffer or batch]
Enc[Inbound enc nodes] --> Buckets[Lazy per-kind buckets]
Devices[Recipient devices] --> Plan[Lazy session override plan]
Plan --> Encrypt[Signal encryption fan-out]
Encrypt --> Stanza[Pre-sized message stanza]
Stanza --> Wire[Encoded wire message]
Reviews (2): Last reviewed commit: "test(jid): pin the non-AD render of a no..." | Re-trigger Greptile
📦 Binary size report
.text per crate
Baseline: |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@wacore/binary/src/jid.rs`:
- Around line 1993-2034: Add a nonzero-agent Bot/Interop JID case such as
“user.5:10@bot” to the cases exercised by
to_non_ad_arc_str_matches_to_non_ad_string, preserving the existing parity
assertion and making the intentional agent-information drop explicit.
🪄 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: 096696f8-c67d-4a16-9217-3d7cff11cf0f
📒 Files selected for processing (13)
src/history_sync.rssrc/message/msg_secret.rssrc/message/receive.rssrc/msg_secret_buffer.rssrc/send/mod.rswacore/binary/src/jid.rswacore/binary/tests/jid_non_ad_arc_alloc.rswacore/src/messages.rswacore/src/reporting_token.rswacore/src/send/dm.rswacore/src/send/encrypt.rswacore/src/store/traits.rswacore/src/types/message.rs
Every parity case used agent 0, so none covered the one shape where the non-AD forms deliberately drop information that is_same_chat_as treats as identity-significant. Preserving the agent in to_non_ad_arc_str now fails this case and no other.
A DM round trip made ~167 allocator calls. About a dozen of them were building the same string twice, sizing buffers that stay empty, or heap-allocating a value that fits inline.
What was removed
Message-secret rows, ~5 per round trip.
msg_secret.rs,send/mod.rsandhistory_sync.rseach builtMsgSecretEntrywith the same three lines, andchat.to_non_ad_string().into()is two allocations per identifier: aString, then theArc<str>its bytes are copied into.Jid::to_non_ad_arc_str()renders through the file's existingJidStackWriterand builds theArc<str>in one allocation, with a heap fallback when the 64-byte stack buffer overflows.MsgSecretEntry::new/sender_id_forare now the single chokepoint, and they share one allocation betweenchatandsenderwhen both name the same user, which is every DM. The history-sync collector already did that aliasing by hand; the reasoning now lives in one place.Enc classification, 2 per inbound message. All three per-kind payload buckets were reserved to the stanza's enc count, but a stanza's enc nodes are essentially always one kind, so two buffers were allocated and never written. Reserving on first push means a mixed stanza still pays one allocation per non-empty bucket.
Encrypt fan-out, 2 per send.
ensure_sessions_for_devicesbuiltvec![None; devices.len()]plus a secondVecof the same length on every call; on a warm send neither is ever written. Both are lazy now.SessionPlangained adevice_countfield so the existing "built for a different device list"debug_assertstays exact rather than being weakened by a now-possibly-empty vector.Stanza building, ~2 per send.
vec![participants]reserves exactly one slot, so appending<device-identity>,<reporting>and caller nodes reallocated and memcpy'd wholeNodevalues.result.version.to_string()heap-allocated a one-byteStringwhereNodeValue's integer conversion formats viaitoainto an inlineCompactString.EditAttribute::from(s.to_string())was pure waste: theFrom<String>impl immediately re-borrows and drops theString.Measurements
Harness
pingpong, 120k messages at 12k/s,MODE=rss, four interleaved ABBA/BAAB pairs againstmain, pre-built binaries with the sha256 recorded per run. Every run hadlost=0andack=120000.That matches the static count of removed call sites on this exact path (1 recipient device, no companions), so the two agree independently.
CPU is deliberately not quoted here. The host was under competing load during this run (pong latency on the baseline side swung between 1.2 ms and 30 ms), which makes any timing number from it worthless. The allocation figures survive that because they are deterministic counts rather than timings. A clean CPU A/B can be run on request; the honest expectation is a small effect, since malloc+cfree together are ~2.8% of the profile and this removes ~6% of the calls.
What was investigated and deliberately left alone
foldhash(0.92% of the profile, andfoldhashis already in the tree via hashbrown, so it would cost no new dependency). Not done here: foldhash's own README states it does not claim HashDoS resistance against interactive attackers, and the keys in these maps (JIDs, Signal protocol addresses) come from the network. The maps are bounded, which caps a flood at O(N) probes, and the user part is a server-assigned numeric identifier rather than attacker-chosen bytes, so the risk is low rather than zero. It is a security-relevant global change worth its own reviewed PR, not a bundle rider. Calibration: libsignal keysHashMap<ProtocolAddress, _>with std's SipHash and whatsmeow stringifies the address on every session lookup, so neither reference implementation treats this as hot.__vdso_clock_gettime(1.25%). Our own call count is already budgeted and tested (dm_send_stays_within_its_clock_budgetpins 1 wall + 2 monotonic reads per send;received_stanza_handling_reads_no_clockpins 0 on receive), and the hot coordination caches are built without TTL soPortableCache::entry_timereturns a sentinel instead of reading the clock. What remains is Tokio's scheduler and timer park; Tokio does not cache time and has no pluggable clock (tokio#3918).EncryptResult::encrypted_devicesis built and then discarded by the DM path. Removing it is ~1 allocation per message but changes apubfield on apubstruct; not worth the API churn.Box::pin(ready(...))in the Signal store adapters, ~3 boxed ready futures per device per crypto op. The hand-desugaring already there removed the state-machine size but not the allocation; removing the allocation requires the trait to stop returning a boxed future, which is a wide change across vendored libsignal.push_name: .map(|s| s.to_string())looked like a freeCow::into_ownedwin and is not:optional_stringon a string attribute always yieldsCow::Borrowed, so both spellings allocate identically.SmallVecfor the override map and the enc buckets was considered and rejected: all of them live across an.await, where inline storage inflates the future and can cost more in memcpy than it saves (andsend_message_with_optionshas a future-size test that this would have risked tripping).Tests
Nine new tests, each mutation-checked (implementation broken deliberately, failure confirmed, then restored):
to_non_ad_arc_str_matches_to_non_ad_string(multibyte, 80-char user overflowing the stack buffer, empty JID)self.agent, self.deviceinstead of0, 0to_non_ad_arc_str_costs_one_allocationdirect_message_shares_one_identifier_allocation+ 2 siblingssender_id_forentry_construction_allocates_once_per_distinct_identifierchat_idtoArc::from(to_non_ad_string())a_one_enc_stanza_sizes_one_bucket_exactly+the_first_push_reserves_the_whole_stanzareserve_exactan_empty_map_answers_every_index_without_allocatingassume_readyeager againrecording_materializes_the_whole_map_onceindex + 1instead ofdevice_counttest_build_reporting_node(extended).attr("v", "2")edit_attribute_parses_identically_from_borrowed_and_ownedFrom<String>always yieldUnknownBad-path coverage: empty message id, empty/default JID, zero enc count,
usize::MAXindex into an empty override map, unknowneditwire value, unicode/multibyte user parts, and a user part that overflows the stack buffer.Security
No hasher, data structure or lock discipline changed.
sender_id_foraliases two identifiers only whenis_same_chat_asholds, which is stricter than string equality of the rendered non-AD form (it also comparesintegrator, andagentwhere the server renders it), so it can only ever fail to alias, never merge two distinct lookup keys. A test pins that the LID and PN forms of the same user part stay separate.to_non_ad_arc_stris byte-identical toto_non_ad_stringincluding the overflow fallback, so no persisted or looked-up key changes shape.Verification
cargo fmt --all,cargo clippy --workspace --all-targetswith zero warnings, and all test binaries acrosswacore-binary,wacoreandwhatsapp-rustgreen (1216 and 1239 in the two lib suites).