Skip to content

perf: cut a dozen per-message allocations from the DM round trip - #1122

Merged
jlucaso1 merged 2 commits into
mainfrom
perf/allocs
Jul 26, 2026
Merged

perf: cut a dozen per-message allocations from the DM round trip#1122
jlucaso1 merged 2 commits into
mainfrom
perf/allocs

Conversation

@jlucaso1

Copy link
Copy Markdown
Collaborator

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.rs and history_sync.rs each built MsgSecretEntry with the same three lines, and chat.to_non_ad_string().into() is two allocations per identifier: a String, then the Arc<str> its bytes are copied into. Jid::to_non_ad_arc_str() renders through the file's existing JidStackWriter and builds the Arc<str> in one allocation, with a heap fallback when the 64-byte stack buffer overflows. MsgSecretEntry::new / sender_id_for are now the single chokepoint, and they share one allocation between chat and sender when 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_devices built vec![None; devices.len()] plus a second Vec of the same length on every call; on a warm send neither is ever written. Both are lazy now. SessionPlan gained a device_count field so the existing "built for a different device list" debug_assert stays 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 whole Node values. result.version.to_string() heap-allocated a one-byte String where NodeValue's integer conversion formats via itoa into an inline CompactString. EditAttribute::from(s.to_string()) was pure waste: the From<String> impl immediately re-borrows and drops the String.

Measurements

Harness pingpong, 120k messages at 12k/s, MODE=rss, four interleaved ABBA/BAAB pairs against main, pre-built binaries with the sha256 recorded per run. Every run had lost=0 and ack=120000.

metric main branch delta Welch t
allocator calls per message 167.04 (sd 0.39) 156.33 (sd 0.43) -6.41% -37.04
bytes requested per message 29 564 (sd 19) 29 187 (sd 24) -1.28% -24.48

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

  • Swapping SipHash for foldhash (0.92% of the profile, and foldhash is 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 keys HashMap<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_budget pins 1 wall + 2 monotonic reads per send; received_stanza_handling_reads_no_clock pins 0 on receive), and the hot coordination caches are built without TTL so PortableCache::entry_time returns 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_devices is built and then discarded by the DM path. Removing it is ~1 allocation per message but changes a pub field on a pub struct; 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.
  • The node decoder, ~5-8 allocations per inbound stanza for attribute and child boxed slices. Real, but it needs an arena or node pool tied to the yoke.
  • push_name: .map(|s| s.to_string()) looked like a free Cow::into_owned win and is not: optional_string on a string attribute always yields Cow::Borrowed, so both spellings allocate identically.
  • SmallVec for 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 (and send_message_with_options has a future-size test that this would have risked tripping).

Tests

Nine new tests, each mutation-checked (implementation broken deliberately, failure confirmed, then restored):

Test Mutation Caught
to_non_ad_arc_str_matches_to_non_ad_string (multibyte, 80-char user overflowing the stack buffer, empty JID) pass self.agent, self.device instead of 0, 0 yes
to_non_ad_arc_str_costs_one_allocation delete the stack-buffer fast path yes, 2 vs 1
direct_message_shares_one_identifier_allocation + 2 siblings drop the aliasing in sender_id_for yes, 2 of 3 (the third is the negative case and correctly still passes)
entry_construction_allocates_once_per_distinct_identifier revert chat_id to Arc::from(to_non_ad_string()) yes
a_one_enc_stanza_sizes_one_bucket_exactly + the_first_push_reserves_the_whole_stanza remove the reserve_exact yes, capacity 4 vs 1
an_empty_map_answers_every_index_without_allocating make assume_ready eager again yes
recording_materializes_the_whole_map_once resize to index + 1 instead of device_count yes
test_build_reporting_node (extended) hardcode .attr("v", "2") yes
edit_attribute_parses_identically_from_borrowed_and_owned make From<String> always yield Unknown yes

Bad-path coverage: empty message id, empty/default JID, zero enc count, usize::MAX index into an empty override map, unknown edit wire 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_for aliases two identifiers only when is_same_chat_as holds, which is stricter than string equality of the rendered non-AD form (it also compares integrator, and agent where 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_str is byte-identical to to_non_ad_string including the overflow fallback, so no persisted or looked-up key changes shape.

Verification

cargo fmt --all, cargo clippy --workspace --all-targets with zero warnings, and all test binaries across wacore-binary, wacore and whatsapp-rust green (1216 and 1239 in the two lib suites).

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

coderabbitai Bot commented Jul 26, 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 (base), Organization UI (inherited)

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: ca06342e-6a75-46f6-a8d4-7a08f4e2d22a

📥 Commits

Reviewing files that changed from the base of the PR and between e56cc52 and d2dbc73.

📒 Files selected for processing (1)
  • wacore/binary/src/jid.rs

📝 Walkthrough

Summary by CodeRabbit

  • Bug Fixes
    • Improved consistency when capturing and persisting message secrets, including correct sender identifier handling and ensuring alias pairs are stored together.
    • Strengthened edit and reporting-token parsing to preserve original wire values.
  • Performance
    • Reduced allocations across message classification, DM stanza construction, encryption fan-out, and attribute parsing.
    • Added a more allocation-efficient JID non-AD string conversion path for smoother operation under load.

Walkthrough

This PR centralizes MsgSecretEntry construction and sender canonicalization, adds direct Arc<str> JID conversion, and reduces eager allocations in message reception, encryption fan-out, parsing, reporting nodes, and DM stanza assembly. Tests cover aliasing, allocation counts, capacity behavior, and wire compatibility.

Changes

Message-secret canonicalization and persistence

Layer / File(s) Summary
Canonical entry construction
wacore/src/store/traits.rs, wacore/binary/src/jid.rs
MsgSecretEntry::new and sender_id_for centralize identifier conversion and same-chat aliasing; Jid::to_non_ad_arc_str provides direct Arc<str> conversion with fallback handling.
Message-secret call sites
src/message/msg_secret.rs, src/history_sync.rs, src/send/mod.rs
Inbound, history-sync, and outbound persistence use the centralized constructor; inbound aliases are persisted together when both exist.
Canonicalization and allocation tests
wacore/src/store/traits.rs, wacore/binary/src/jid.rs, wacore/binary/tests/*, src/msg_secret_buffer.rs
Tests cover identifier aliasing, namespace separation, JID output, and allocation counts.

Lazy payload and encryption storage

Layer / File(s) Summary
Incoming encryption buckets
src/message/receive.rs
Encoding buckets reserve capacity only when first used, with tests for untouched buckets, filling behavior, and zero-count stanzas.
Per-device encryption overrides
wacore/src/send/encrypt.rs
Session plans store device counts separately, materialize override vectors on demand, and use safe indexed lookup across session preparation and encryption fan-out.
Override validation tests
wacore/src/send/encrypt.rs
Tests cover empty-map allocation, first override materialization, indexed updates, and single-device lookup.

Wire and stanza allocation paths

Layer / File(s) Summary
Wire parsing and serialization
wacore/src/messages.rs, wacore/src/types/message.rs, wacore/src/reporting_token.rs
Edit attributes are parsed from borrowed values, reporting-token versions flow as integers, and tests verify unchanged parsing and rendered values.
DM stanza assembly
wacore/src/send/dm.rs
Message-content nodes pre-allocate for participants, optional nodes, and extra stanza nodes before appending them.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

Suggested labels: api-design, performance

Suggested reviewers: cubic-dev-ai

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly and concisely summarizes the main change: reducing per-message allocations in the DM round trip.
Description check ✅ Passed The description is directly related to the PR and accurately describes the allocation and correctness changes in the patch.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch perf/allocs

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.

@greptile-apps

greptile-apps Bot commented Jul 26, 2026

Copy link
Copy Markdown

Greptile Summary

The PR reduces per-message allocations without changing wire or persistence semantics.

  • Centralizes message-secret identifier construction and shares canonical direct-message identifiers.
  • Lazily allocates receive buckets and per-device encryption override storage.
  • Pre-sizes stanza children and removes temporary strings during attribute construction and parsing.
  • Adds focused tests for canonicalization, allocation behavior, aliasing, lazy storage, and wire values.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

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]
Loading

Reviews (2): Last reviewed commit: "test(jid): pin the non-AD render of a no..." | Re-trigger Greptile

@github-actions

github-actions Bot commented Jul 26, 2026

Copy link
Copy Markdown

📦 Binary size report

Metric main PR Δ
bin size (stripped) 9.98 MiB 9.98 MiB 0
bin .text 8.03 MiB 8.03 MiB -192 B (-0.00%) 🔽
bin allocated (text+data+bss) 9.98 MiB 9.98 MiB -80 B (-0.00%) 🔽
llvm-lines wacore 492,078 492,046 -32 (-0.01%) 🔽
llvm-lines wacore copies 16,335 16,327 -8 (-0.05%) 🔽
llvm-lines whatsapp-rust lib 720,635 720,617 -18 (-0.00%) 🔽
llvm-lines whatsapp-rust lib copies 22,772 22,770 -2 (-0.01%) 🔽
deps crates (Cargo.lock) 471 471 0
.text per crate
Crate main PR Δ
.text whatsapp_rust 1.82 MiB 1.82 MiB -187 B (-0.01%) 🔽
.text wacore 646.24 KiB 645.75 KiB -505 B (-0.08%) 🔽
.text wacore_binary 89.34 KiB 89.70 KiB +375 B (+0.41%) 🔺
.text wacore_libsignal 161.84 KiB 161.84 KiB 0
.text wacore_appstate 22.36 KiB 22.36 KiB 0
.text wacore_noise 21.60 KiB 21.60 KiB 0
.text waproto 1.74 MiB 1.74 MiB 0
.text whatsapp_rust_sqlite_storage 514.91 KiB 514.91 KiB 0
.text whatsapp_rust_tokio_transport 39.91 KiB 39.91 KiB 0
.text whatsapp_rust_ureq_http_client 10.40 KiB 10.40 KiB 0
.text std 1.07 MiB 1.07 MiB +385 B (+0.03%) 🔺
.text other deps 1.89 MiB 1.89 MiB -273 B (-0.01%) 🔽

Baseline: 7dad9099b (latest main run) · Head: 0d703a1ed · Graphs

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

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

📥 Commits

Reviewing files that changed from the base of the PR and between 4d781fe and e56cc52.

📒 Files selected for processing (13)
  • src/history_sync.rs
  • src/message/msg_secret.rs
  • src/message/receive.rs
  • src/msg_secret_buffer.rs
  • src/send/mod.rs
  • wacore/binary/src/jid.rs
  • wacore/binary/tests/jid_non_ad_arc_alloc.rs
  • wacore/src/messages.rs
  • wacore/src/reporting_token.rs
  • wacore/src/send/dm.rs
  • wacore/src/send/encrypt.rs
  • wacore/src/store/traits.rs
  • wacore/src/types/message.rs

Comment thread wacore/binary/src/jid.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.
@jlucaso1
jlucaso1 merged commit 718def6 into main Jul 26, 2026
19 of 20 checks passed
@jlucaso1
jlucaso1 deleted the perf/allocs branch July 26, 2026 17:59
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