perf(prekeys): stream prekey generation to cut the connect-time peak

[jlucaso1]

What

Follow-up to #900. upload_pre_keys_pass generates the one-time pre-key batch (812 keys by default, WA Web fidelity) in a blocking offload, then carries their public keys to the upload IQ. That offload built all the records up front — a Vec<(u32, PreKeyRecordStructure)> where every record owns two heap Vecs (public + private key bytes) — so the whole batch sat resident at once, alongside the public keys and the encoded buffer, just to be encoded and dropped immediately after.

Encode each record straight into the shared batch buffer and drop it in the same iteration, keeping only its public key:

• No collective records Vec — at most one PreKeyRecordStructure is live at a time.
• The buffer is pre-sized by the 73-byte max record length (id field ≤5 B + two 34 B key fields) so it stays a single allocation, instead of relying on an up-front encoded_len() sum over the full materialized batch.

Why

#900 cut allocation count/volume on the connect path by carrying generated public keys to the upload instead of reloading + decoding them, but it left the records batch fully materialized during generation — the dominant controllable contributor to connect_to_ready's peak memory (the irreducible part being the 812 X25519 keygens themselves). Removing that batch's residency targets peak directly while preserving #900's allocation win.

Verification

• cargo fmt --all / cargo clippy -p whatsapp-rust --lib --tests — clean
• cargo test -p whatsapp-rust --lib prekeys — 16 tests pass, including the window_tests that drive upload_pre_keys_pass against a mock backend (upload-window, retry-single-key, collapse/regenerate)
• Wire output is unchanged: the encoded records and the upload IQ are byte-identical; only the order of construction (and what stays resident) changes.
• CodSpeed will quantify the connect_to_ready peak-memory delta on this PR.

Scope

Generation-phase only. The upload itself stays a single batched IQ (WA Web WAWebUploadPreKeysJob), and the public keys + final encoded frame are still resident at send time — that floor is inherent to the protocol and small.

---

Generated by Claude Code

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: c20e0cd4-f77b-487f-875c-e33f945146c8

📥 Commits

Reviewing files that changed from the base of the PR and between 3478da7 and c6017e9.

📒 Files selected for processing (1)

• src/prekeys.rs

---

📝 Walkthrough

Summary by CodeRabbit

• Chores
  • Improved the efficiency of pre-key batch encoding during upload, reducing extra in-memory work and streamlining the preparation of the upload payload.

Walkthrough

Inside Client::upload_pre_keys_pass, the pre-key batch encoding is rewritten for correctness and efficiency: the records vector and total_len precomputation are removed. Each generated pre-key record is now encoded directly into a single shared Vec<u8> using a constant MAX_RECORD_LEN = 74 for capacity reservation, with byte offsets tracked to slice per-key Bytes for the upload payload. This eliminates unnecessary intermediate allocations and ensures the buffer is sized correctly from the start.

Changes

Pre-key batch encoding refactor

Layer / File(s)

Summary

Shared buffer encoding with offset slicing src/prekeys.rs

Removes the intermediate records vector and total_len precomputation to eliminate allocations that don't work reliably. Each pre-key record is encoded directly into one shared Vec<u8> buffer reserved using MAX_RECORD_LEN = 74. Byte offsets (id, start..end) are collected and used to slice the buffer into per-key Bytes for the Vec<(u32, Bytes)> upload payload, ensuring the encoding path works correctly without intermediate steps.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• oxidezap/whatsapp-rust#695: Also modifies the pre-key upload encoding path in src/prekeys.rs, refactoring how batches are constructed as contiguous encoded buffers — shares the same buffer construction concerns.
• oxidezap/whatsapp-rust#833: Modifies Client::upload_pre_keys_pass at the same function level, refactoring watermark handling and batch planning logic that directly surrounds this encoding change.
• oxidezap/whatsapp-rust#324: Addresses the broader pre-key upload and validation flow in src/prekeys.rs — the encoding refactor here feeds directly into that path.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the core optimization: streaming prekey generation to reduce peak memory during connection, which directly matches the refactoring described in the changeset.
Description check	✅ Passed	The description is comprehensive and directly related to the changeset, explaining the motivation, implementation details, verification steps, and scope of the optimization.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch perf/prekey-gen-stream-peak

---

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

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

📦 Binary size report

Metric	main	PR	Δ
bin size (stripped)	10.59 MiB	10.59 MiB	-960 B (-0.01%) 🔽
bin .text	8.70 MiB	8.70 MiB	-896 B (-0.01%) 🔽
bin allocated (text+data+bss)	10.59 MiB	10.59 MiB	+8 B (+0.00%) 🔺
llvm-lines wacore	639,594	639,594	0
llvm-lines wacore copies	17,666	17,666	0
llvm-lines whatsapp-rust lib	655,265	654,931	-334 (-0.05%) 🔽
llvm-lines whatsapp-rust lib copies	20,171	20,159	-12 (-0.06%) 🔽
deps crates (Cargo.lock)	347	347	0

.text per crate

Crate	main	PR	Δ
.text whatsapp_rust	1.44 MiB	1.44 MiB	-99 B (-0.01%) 🔽
.text wacore	518.09 KiB	517.29 KiB	-821 B (-0.15%) 🔽
.text wacore_binary	156.73 KiB	156.73 KiB	0
.text wacore_libsignal	166.63 KiB	166.63 KiB	0
.text wacore_appstate	36.98 KiB	36.98 KiB	0
.text wacore_noise	27.71 KiB	27.71 KiB	0
.text waproto	876.21 KiB	876.21 KiB	0
.text whatsapp_rust_sqlite_storage	207.48 KiB	207.48 KiB	0
.text whatsapp_rust_tokio_transport	32.49 KiB	32.49 KiB	0
.text whatsapp_rust_ureq_http_client	5.93 KiB	5.93 KiB	0
.text std	1.13 MiB	1.13 MiB	0
.text other deps	4.07 MiB	4.07 MiB	0

Top movers (cargo-bloat attribution)

Crate	main	PR	Δ
prost	471.26 KiB	472.66 KiB	+1.39 KiB (+0.30%)

Baseline: 9d7e6de88 (latest main run) · Head: cca7ed2e1 · Graphs

[coderabbitai]

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 `@src/prekeys.rs`:
- Around line 476-482: The constant MAX_RECORD_LEN is set to 73 bytes but the
actual maximum protobuf-encoded record size is 74 bytes. The calculation is: id
field requires 6 bytes maximum (1-byte tag plus 5-byte varint), publicKey field
requires 34 bytes (1-byte tag, 1-byte length, 32 bytes), and privateKey field
requires 34 bytes (same breakdown), totaling 74 bytes. Change the MAX_RECORD_LEN
constant from 73 to 74 to match the correct maximum record length and update the
associated comment to reflect the accurate byte breakdown for each field,
particularly correcting the id field from 5 bytes to 6 bytes to include the tag
byte.

🪄 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: 994f1ecb-fef5-40f7-984f-4c281d5fa2a1

📥 Commits

Reviewing files that changed from the base of the PR and between 9d7e6de and 3478da7.

📒 Files selected for processing (1)

• src/prekeys.rs

[cubic-dev-ai]

No issues found across 1 file

_{Re-trigger cubic}

[codspeed-hq]

Merging this PR will not alter performance

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 1 improved benchmark
❌ 1 (👁 1) regressed benchmark
✅ 186 untouched benchmarks

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
⚡	Simulation	send_message[20]	5.6 ms	5 ms	+10.84%
👁	Memory	send_and_receive[20]	65.7 KB	79 KB	-16.89%

Tip

Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing perf/prekey-gen-stream-peak (c6017e9) with main (931a5d6)}