perf(store): reserve the prekey map for a batch insert - #1270
Conversation
`store_prekeys_batch` inserted a connect-sized batch one key at a time into a map with no reservation, so the table crossed the load factor eight times on the way to holding it: eight allocations, eight copies of everything already inserted, and eight frees. The batch length is known at the call, and every id in it is new (prekey ids come from the monotonic NEXT_PK_ID counter), so one `reserve` reaches the final size in a single allocation. Measured on the connect batch of 812 (`DEFAULT_WANTED_PRE_KEY_COUNT`): 11 allocations / 84.08 KB become 3 / 42.13 KB, and the in-call transient high-water falls from 63.07 KB to 42.13 KB. Retained is unchanged at 42,072 B — the final table is the same size either way — so this buys allocator traffic, not residency. The observability ledger gains that distinction, since the profiler attributes the retained table to `reserve_rehash` and makes it look like churn. Adds a divan `AllocProfiler` bench for the path and two tests: one that a batch round-trips and stays idempotent per id, one that the reserved table matches an incrementally grown one so the reservation cannot over-grow a map that already holds a window. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A5Vcfvdd5n2PX2Z9M8Ympc
📝 WalkthroughSummary by CodeRabbit
Walkthrough
ChangesPrekey batch storage
Estimated code review effort: 2 (Simple) | ~15 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 |
…t PR Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A5Vcfvdd5n2PX2Z9M8Ympc
|
| Filename | Overview |
|---|---|
| wacore/src/store/in_memory.rs | Adds a conservative reservation before prekey batch insertion and comprehensive capacity and storage regression tests. |
| wacore/benches/prekey_store_benchmark.rs | Adds allocation benchmarks for fresh and already-populated in-memory prekey maps. |
| wacore/Cargo.toml | Registers the new prekey-store benchmark target. |
| agent_docs/observability.md | Clarifies that the optimization reduces transient allocation work rather than retained map memory. |
Reviews (3): Last reviewed commit: "perf(store): reserve for a prekey batch ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| bencher | ||
| .with_inputs(|| (InMemoryBackend::new(), batch(1, count))) | ||
| .bench_refs(|(backend, keys)| { | ||
| block_on(backend.store_prekeys_batch(keys, false)).unwrap(); |
There was a problem hiding this comment.
The new benchmark calls .unwrap() here and again on lines 66 and 71, contrary to the repository-wide prohibition on .unwrap() outside tests. These calls also make benchmark failures panic without contextual error propagation.
Context Used: CLAUDE.md (source)
Prompt To Fix With AI
This is a comment left during a code review.
Path: wacore/benches/prekey_store_benchmark.rs
Line: 54
Comment:
**Benchmark uses prohibited unwraps**
The new benchmark calls `.unwrap()` here and again on lines 66 and 71, contrary to the repository-wide prohibition on `.unwrap()` outside tests. These calls also make benchmark failures panic without contextual error propagation.
**Context Used:** CLAUDE.md ([source](https://github.com/oxidezap/whatsapp-rust/blob/main/CLAUDE.md))
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
![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: fa04c479c8
ℹ️ 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".
| // This does NOT shrink the table that stays resident: the final | ||
| // capacity is the same either way, so it buys allocator traffic and | ||
| // in-call headroom, not retained bytes. | ||
| state.prekeys.reserve(keys.len()); |
There was a problem hiding this comment.
When a caller replays or rewrites an existing batch, HashMap::reserve(keys.len()) assumes every entry will increase len, although the public SignalStore::store_prekeys_batch contract permits overwrites and the new idempotence test exercises one. Replaying an 812-key window therefore grows a map from capacity 896 to 1792 before replacing the same rows, permanently retaining roughly another 42 KiB; reserve only for distinct IDs not already present, or otherwise enforce the all-new invariant at the API boundary.
Useful? React with 👍 / 👎.
📦 Binary size report
.text per crate
Baseline: |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
The reservation took the bare batch length, which assumes every id is new. The trait permits a batch to overwrite ids, and a table grown for rows that were only overwritten never shrinks back — so replaying a stored 812-key window would have taken the map from 896 to 1792 capacity and kept the extra table forever. That is precisely the residency this change claims not to touch. Reserve `keys.len() - prekeys.len()` instead: no batch can overwrite more rows than exist, so that difference is the floor on how many ids must be new. It can only under-reserve and fall back to incremental growth, never inflate the resident table. The connect path is unaffected — the map is empty there, so the whole batch is still reserved: 3 allocations / 42.13 KB, dealloc 1 / 64 B, the same figures as before. A replay now reserves nothing and leaves the table exactly as it found it, which `replaying_a_stored_batch_does_not_ grow_the_table` pins. Reported independently by Greptile and Codex on #1270. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A5Vcfvdd5n2PX2Z9M8Ympc
|
Both review bots landed on the same issue independently. One taken, one declined. Reservation assumes all-new IDs — valid, fixed in The objection is right and it undercut the PR's own thesis: the trait says nothing about ids being new ( Rather than gate on let at_least_new = keys.len().saturating_sub(state.prekeys.len());
state.prekeys.reserve(at_least_new);No batch can overwrite more rows than exist, so Connect path is unchanged, re-measured after the fix: 3 allocations / 42.13 KB, dealloc 1 / 64 B — identical to the numbers in the PR body.
The rule in On the substance: propagating instead of panicking would be wrong here. Also for the record, unrelated to review: the red "Semver Checks (informational)" is pre-existing and not from this PR. It is Generated by Claude Code |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cf70e5e7f7
ℹ️ 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".
| let at_least_new = keys.len().saturating_sub(state.prekeys.len()); | ||
| state.prekeys.reserve(at_least_new); |
There was a problem hiding this comment.
Fresh evidence in the revised code is that keys.len().saturating_sub(state.prekeys.len()) still treats duplicate IDs within the incoming batch as new rows. The trait imposes no uniqueness requirement and its default implementation processes duplicates with last-write-wins semantics, so an empty backend receiving 812 entries for one ID reserves an 812-row table but retains only one row, recreating roughly 42 KiB of permanent over-allocation; reserve for distinct IDs absent from the map instead.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Correct — fixed in 4dfd8b5. Subtracting the stored rows closed the map-side hole and left the batch-side one: keys.len() still counted repeated ids as distinct rows, so 812 entries sharing one id would have reserved a 1024-bucket table to hold one row. Same permanent over-allocation, reached from the other direction.
I did not take the literal suggestion of reserving for distinct absent ids, because the dedup costs more than the optimization saves: a HashSet<u32> of the batch is one allocation plus 812 hashes, against the eight allocations being removed. Instead the reservation is now gated on the batch being strictly ascending, which is a sufficient condition for distinctness:
let ascending = keys.windows(2).all(|pair| pair[0].0 < pair[1].0);
if ascending {
let at_least_new = keys.len().saturating_sub(state.prekeys.len());
state.prekeys.reserve(at_least_new);
}One pass of integer compares, no allocation, no hashing. Both guards are one-sided in the same direction — they can only skip or shrink a reservation and fall back to incremental growth, never inflate the resident table — so the "costs no retained bytes" claim now holds for any caller rather than for well-behaved ones.
The connect path satisfies both conditions (upload_pre_keys_pass emits gen_start + i into an empty map), so the measurements are unchanged: 3 allocations / 42.13 KB, dealloc 1 / 64 B, median 26.0 µs at 812. a_batch_of_repeated_ids_does_not_reserve_for_them pins the case you raised — 812 entries for one id must leave the table sized for one row.
Generated by Claude Code
Subtracting the stored rows still assumed the ids within a batch were distinct from each other. The trait requires no such thing and inserts last-write-wins, so 812 entries sharing one id would have reserved a 1024-bucket table to hold a single row — the same permanent over-allocation the previous commit removed, reached from the other side. Gate the reservation on the batch being strictly ascending, which proves its ids distinct. That is one pass of integer compares and no allocation; deduplicating properly would need a set, whose own allocation and 812 hashes cost more than the eight allocations being saved. Like the subtraction, the gate is one-sided: it can only skip a reservation and fall back to incremental growth, never inflate the resident table. The connect path satisfies it — `upload_pre_keys_pass` emits `gen_start + i` — so the measurements are untouched: 3 allocations / 42.13 KB, dealloc 1 / 64 B, median 26.0 µs at 812. Reported by Codex on #1270. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A5Vcfvdd5n2PX2Z9M8Ympc
Dismissed because a newer commit was pushed; Greptile will re-review the current head.
2 participants
Summary
InMemoryBackend::store_prekeys_batchtakeskeys: &[(u32, Bytes)]and inserts them one at a time intostate.prekeys: HashMap<u32, PreKeyEntry>with no reservation. A connect-sized batch therefore crosses hashbrown's 7/8 load factor eight times on the way in — each crossing allocates a table twice the size, rehashes and moves everything already inserted, and frees the old one. The batch length is known at the call site, so the table can reach its final size in one allocation instead.The number that motivated this is not the number this fixes, and that is the main thing this PR wants on the record — see
## Cost.Changes
wacore/src/store/in_memory.rs: reserve before the insert loop, but only for rows that provably must be added:A table grown for rows that never arrive does not shrink back, so a naive
reserve(keys.len())would cost exactly the retained bytes this PR claims not to touch. Two one-sided guards prevent that, and each closes a different hole:keys.len() - len()is the floor on how many ids must be new. Without it, replaying a stored 812-key window would take the map from 896 to 1792 capacity and keep the extra ~42 KiB forever.Both can only skip or shrink a reservation and fall back to incremental growth — never inflate the resident table — so the guarantee holds for any caller, not just well-behaved ones. The connect path satisfies both: the map is empty and
upload_pre_keys_passemitsgen_start + ifrom the monotonicNEXT_PK_IDcounter, so the whole batch is reserved and gets the full win.wacore/benches/prekey_store_benchmark.rs(new, registered inwacore/Cargo.toml): adivan::AllocProfilerbench over the path at 64 and 812 keys, plus apopulatedrow that writes a second full batch into a map that already holds a window. It lands in the existing CodSpeedcoreshard, which runs both the simulation and memory instruments.wacore/src/store/in_memory.rstests:store_prekeys_batch_stores_every_key,store_prekeys_batch_reserve_does_not_over_grow_the_table,replaying_a_stored_batch_does_not_grow_the_table, anda_batch_of_repeated_ids_does_not_reserve_for_them.agent_docs/observability.md: the prekey-window ledger entry gains the peak-vs-retained clarification below, so the next reader of a dhat profile does not re-derive it.No signature, semantics, map type, or hash strategy changed.
store_prekeys_batchstays idempotent per id.Cost
Batch size, from the source rather than assumed.
upload_pre_keys_passcallsstore_prekeys_batch(&encoded_batch)with exactlyplan.gen_countentries (src/prekeys.rs:572).plan.gen_countcomes fromplan_prekey_upload(first_unupload, next, max_id, wanted), wherewantedisDEFAULT_WANTED_PRE_KEY_COUNT = 812(src/prekeys.rs:25, mirroring WA Web'sUPLOAD_KEYS_COUNTinWAWebUploadPreKeysJob), clamped into[5, 65535].gen_count = wanted - available, andavailablecounts only leftover generated-but-unuploaded window keys. On a fresh connectavailable == 0and the plan isgen_count == 812— pinned by the existing planner tests (src/prekeys.rs:1011,:1021,:1053,:1067). The batch is genuinely hundreds, not dozens. The one path that writes fewer is the retry-receipt allocation, and that calls single-keystore_prekey, not the batch.Allocations and bytes,
divan::AllocProfiler,cargo bench -p wacore --bench prekey_store_benchmark, fresh backend per iteration, 100 samples:The
populatedrow is byte-identical on both sides: a second window crosses the load factor exactly once either way, so reserving on an already-populated map neither helps nor over-grows. The 812 row was re-measured after each of the two guards was added and did not move.Peak against retained — explicitly. Measured with a tracking
GlobalAllocthat records live bytes on every alloc/dealloc, snapshotted after the call returns with the backend still alive, batch of 812:Retained is bit-identical. It has to be: 812 rows need 1024 buckets on both paths, and 1024 × (
size_of::<(u32, PreKeyEntry)>()= 40 + 1 control byte) = 41,984 B is the table that stays. So the 41.0 KiB/session that dhat attributed tohashbrown::RawTable::reserve_rehashunder this function is the final resident table, not rehash churn — 41,984 B is that figure to three digits. dhat names the frame that allocated the block, and the frame that allocates the last table isreserve_rehash; the intermediate tables are all freed before the block that survives to the global peak. This change moves that allocation fromreserve_rehashtoreserveand does not shrink it by one byte. A profile diff taken after this lands will still show 41.0 KiB there.What the change actually removes, per connect, per session: 8 allocations, 8 frees, ~42 KB of transient allocator traffic, ~890 element moves with a rehash each, and 20.9 KB of in-call transient high-water. Whether that 20.9 KB shows up in process RSS depends on whether sessions open concurrently; the external profile it came from reported exactly 41.0 KiB/session with no intermediate contribution, which is what sequential opening looks like — so the honest expectation is that RSS at the peak barely moves and the win is allocator work, not memory. Reported as such rather than as retention.
Time is secondary and reported with its methodology: divan wall-clock,
--release, single unloaded container, 100 samples, median 48.4 µs → 26.0 µs at 812 (fastest 44.0 → 25.3 µs). Consistent with dropping ~890 rehash moves, but a single-box wall-clock number on a shared runner; the instruction-count instrument in CodSpeed is the one to trust.Checked and not changed
storages/sqlite-storage/src/sqlite_store.rshas the analog, and it is worse.store_prekeys_batchthere is not a transaction per key — the 812 inserts do share oneconn.transaction— but inside it every row is its owninsert_into(...).on_conflict(...).execute(conn), i.e. 812 round trips through Diesel where a multi-rowVALUESinsert (chunked under SQLite's variable limit) would do. On top of that the batchVecis copied twice before any SQL runs:keys.to_vec()once per call, thenkeys_clone = keys.clone()once per retry attempt, which at 812 ×size_of::<(u32, Bytes)>()= 40 B is 32.5 KB per copy — more than the whole in-memory map's table. Measured: one 812-key call against a shared-cache in-memory SQLite takes 8.0 ms release / 26.2 ms debug. Left alone per this batch's scope; worth a separate batch, which now has a number to beat.store_prekeys_batchis the only bulk populator ofstate.prekeys. The map has exactly two writers inin_memory.rs: this one and single-keystore_prekey.store_prekey's only non-test callers are the retry-receipt path (src/prekeys.rs:403, one key) and the libsignal adapter (src/store/signal_adapter.rs:432, one key) — neither loops. The defaultstore_prekeys_batchinstore/traits.rsdoes fan out tostore_prekey, butInMemoryBackendoverrides it, so that fallback is not a path into this map.RandomStatewould change these numbers too and is a different argument with a different risk; not attempted here.Validation
cargo fmt --all --checkcargo clippy -p wacore --all-targets -- -D warnings— clean locally. The fullcargo clippy --workspace --all-targetscould not run in the dev container:alsa-sys's build script fails for want oflibasound/pkg-config, pulled in byexamples/voip-cli(cpal). Confirmed environmental by reproducing it on a stashed, clean tree — and CI's Build & Lint (all features) covers it and passes.cargo test(workspace default members, includes doctests) — 55 test binaries, 0 failures. All four new tests pass.cargo bench -p wacore --bench prekey_store_benchmarkbefore and after, numbers above.4dfd8b5: Build & Test, Clippy, Format, Test Stable (MSRV), Feature Matrix, Build & Lint (all features), Rustdoc, E2E, all three Miri jobs, wasm32, Cargo Deny, Binary Size, and all four CodSpeed jobs.continue-on-error: trueby design ("Advisory only. The workspace is pre-1.0 and intentionally breaks API between minors"), and every failure it lists is inmex_operations.rs, the removedsimdfeature, orBinaryError::UnexpectedFormatByte. This PR adds no public API surface. Theunstablemergeable state is that check alone.