From 3478da79d161c22613500e06eb1dca4a3f355adf Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 18:48:30 +0000 Subject: [PATCH 1/2] perf(prekeys): stream prekey generation to cut the connect-time peak MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit upload_pre_keys_pass built the whole batch of PreKeyRecordStructures in a Vec before encoding them, so all 812 records — each owning two heap Vecs for its public/private key bytes — were resident at once alongside the public keys carried to the upload. That batch was the dominant controllable peak on the connect/registration path. Encode each record straight into the shared buffer and drop it immediately, keeping only its public key. The records are no longer held collectively, and the buffer is pre-sized by the 73-byte max record length to stay a single allocation. Wire output is unchanged. --- src/prekeys.rs | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/src/prekeys.rs b/src/prekeys.rs index e0a41bdda..f2114d424 100644 --- a/src/prekeys.rs +++ b/src/prekeys.rs @@ -473,24 +473,25 @@ impl Client { // Seed one CSPRNG and advance it per key, rather than reseeding from // entropy on every iteration. let mut rng = rand::make_rng::(); - let mut records = Vec::with_capacity(gen_count); + // Encode each record into the shared buffer and drop it immediately, so + // the full batch of PreKeyRecordStructures (each owning two heap Vecs for + // its key bytes) is never resident at once — that batch was the dominant + // controllable peak on the connect path. A record is at most 73 wire bytes + // (id field <=5B + two 34B key fields), so reserving by that bound keeps + // the buffer at a single allocation. + const MAX_RECORD_LEN: usize = 73; let mut pubkeys = Vec::with_capacity(gen_count); + let mut offsets = Vec::with_capacity(gen_count); + let mut buf = Vec::with_capacity(gen_count * MAX_RECORD_LEN); for i in 0..gen_count { let pre_key_id = gen_start + i as u32; let key_pair = KeyPair::generate(&mut rng); - pubkeys.push((pre_key_id, key_pair.public_key)); - records.push((pre_key_id, new_pre_key_record(pre_key_id, &key_pair))); - } - - let total_len: usize = records.iter().map(|(_, r)| r.encoded_len()).sum(); - let mut buf = Vec::with_capacity(total_len); - let mut offsets = Vec::with_capacity(records.len()); - for (id, record) in &records { let start = buf.len(); - record + new_pre_key_record(pre_key_id, &key_pair) .encode(&mut buf) .expect("prost encode into pre-sized Vec"); - offsets.push((*id, start..buf.len())); + offsets.push((pre_key_id, start..buf.len())); + pubkeys.push((pre_key_id, key_pair.public_key)); } let shared = bytes::Bytes::from(buf); let encoded_batch: Vec<(u32, bytes::Bytes)> = offsets From c6017e96e0f945c4eae73876a496151ed497b4b2 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 19:12:30 +0000 Subject: [PATCH 2/2] perf(prekeys): use the type-level record-size bound for the buffer hint MAX_RECORD_LEN pre-sizes the encode buffer to a single allocation; being only a capacity hint, derive it from the u32 id type (1-byte tag + 5-byte varint + two 34 B key fields = 74) rather than the tighter 24-bit id cap, so it is self-evidently correct without depending on the id-range invariant. No behavior change. --- src/prekeys.rs | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/prekeys.rs b/src/prekeys.rs index f2114d424..152982118 100644 --- a/src/prekeys.rs +++ b/src/prekeys.rs @@ -473,13 +473,14 @@ impl Client { // Seed one CSPRNG and advance it per key, rather than reseeding from // entropy on every iteration. let mut rng = rand::make_rng::(); - // Encode each record into the shared buffer and drop it immediately, so - // the full batch of PreKeyRecordStructures (each owning two heap Vecs for - // its key bytes) is never resident at once — that batch was the dominant - // controllable peak on the connect path. A record is at most 73 wire bytes - // (id field <=5B + two 34B key fields), so reserving by that bound keeps - // the buffer at a single allocation. - const MAX_RECORD_LEN: usize = 73; + // Encode each record into the shared buffer and drop it immediately so the + // whole batch of PreKeyRecordStructures (each owns two heap Vecs of key bytes) + // is never resident at once — that batch was the dominant controllable peak on + // the connect path. MAX_RECORD_LEN keeps the buffer to a single allocation; + // being just a capacity hint, it uses the type-level upper bound (1-byte id + // tag + <=5-byte u32 varint + two 34 B key fields = 74) rather than depending + // on the tighter 24-bit id cap. + const MAX_RECORD_LEN: usize = 74; let mut pubkeys = Vec::with_capacity(gen_count); let mut offsets = Vec::with_capacity(gen_count); let mut buf = Vec::with_capacity(gen_count * MAX_RECORD_LEN);