From fa04c479c8bc4f161878aab94fe8edcef0fc2b59 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 21:23:36 +0000 Subject: [PATCH 1/4] perf(store): reserve the prekey map for a batch insert MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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 Claude-Session: https://claude.ai/code/session_01A5Vcfvdd5n2PX2Z9M8Ympc --- agent_docs/observability.md | 12 ++++ wacore/Cargo.toml | 4 ++ wacore/benches/prekey_store_benchmark.rs | 74 ++++++++++++++++++++++ wacore/src/store/in_memory.rs | 81 ++++++++++++++++++++++++ 4 files changed, 171 insertions(+) create mode 100644 wacore/benches/prekey_store_benchmark.rs diff --git a/agent_docs/observability.md b/agent_docs/observability.md index 5060660b7..91864c646 100644 --- a/agent_docs/observability.md +++ b/agent_docs/observability.md @@ -191,6 +191,18 @@ staying alive because the `Bytes` slices handed to `store_prekeys_batch` *are* what the backend stores (one allocation instead of 812), and 41 KiB is that map's `RawTable` at 1024 buckets. Nothing to optimise; do not re-derive it. +That 41 KiB deserves one clarification, because a heap profiler hands it to you +under a name that invites the wrong fix. dhat attributes the final table to +`hashbrown::RawTable::reserve_rehash`, the frame that happened to allocate it, +so a per-session diff reads "41.0 KiB in reserve_rehash" and looks like rehash +churn. It is not: 1024 buckets × (`size_of::<(u32, PreKeyEntry)>()` + 1 control +byte) = 41,984 B is the table that *stays*, and the intermediate tables are all +freed before the process peak. `store_prekeys_batch` does reserve for the batch +length (#1266), which cuts the call from 11 allocations / 84.1 KB to 3 / 42.1 KB +and its in-call transient high-water from 63.1 KB to 42.1 KB — but retained is +bit-identical at 42,072 B either way, because the final table is the same size. +Reserving is worth it for the allocator traffic; it will never move the 41 KiB. + **The rustls session cache is 5 KiB, not 44.** A whole retained `default_tls_connector()` measures 14.0 KiB; disabling resumption entirely takes it to 9.0 KiB, and sizing the store for the one host a factory dials takes it to diff --git a/wacore/Cargo.toml b/wacore/Cargo.toml index 8029c99e4..3164cb2f2 100644 --- a/wacore/Cargo.toml +++ b/wacore/Cargo.toml @@ -142,6 +142,10 @@ harness = false name = "sender_key_derivation_benchmark" harness = false +[[bench]] +name = "prekey_store_benchmark" +harness = false + [[bench]] name = "voip_benchmark" harness = false diff --git a/wacore/benches/prekey_store_benchmark.rs b/wacore/benches/prekey_store_benchmark.rs new file mode 100644 index 000000000..a7d1d3c01 --- /dev/null +++ b/wacore/benches/prekey_store_benchmark.rs @@ -0,0 +1,74 @@ +//! Allocation accounting for the prekey batch write on the connect path. +//! +//! `upload_pre_keys_pass` generates `DEFAULT_WANTED_PRE_KEY_COUNT` (812) one-time +//! prekeys and hands them to `store_prekeys_batch` as one call. `divan::AllocProfiler` +//! is wired as the global allocator so each row reports allocation count and bytes +//! next to wall time -- the count is the signal here, since the map's growth is the +//! only thing this path allocates: the records themselves are `Bytes` slices of one +//! shared buffer the caller already owns, so they cost refcount bumps, not copies. +//! +//! The `populated` row is the second upload pass: the same batch size arriving at a +//! map that already holds a window. It exists so a reservation that over-grows an +//! already-populated table would show up as bytes here rather than silently. + +use bytes::Bytes; +use divan::{Bencher, black_box}; +use futures::executor::block_on; +use wacore::store::in_memory::InMemoryBackend; +use wacore::store::traits::SignalStore; + +#[global_allocator] +static ALLOC: divan::AllocProfiler = divan::AllocProfiler::system(); + +fn main() { + divan::main(); +} + +/// `DEFAULT_WANTED_PRE_KEY_COUNT` in `src/prekeys.rs`, mirroring WA Web's +/// UPLOAD_KEYS_COUNT. The small row is there to show the growth is what scales. +const CONNECT_BATCH: usize = 812; + +/// Upper bound on one encoded `PreKeyRecordStructure`, the same figure +/// `upload_pre_keys_pass` sizes its shared buffer with. +const RECORD_LEN: usize = 74; + +/// The batch as the upload path actually hands it over: every record is a slice of +/// one contiguous buffer, so building it allocates once regardless of `count`. +fn batch(first_id: u32, count: usize) -> Vec<(u32, Bytes)> { + let shared = Bytes::from(vec![7u8; count * RECORD_LEN]); + (0..count) + .map(|i| { + ( + first_id + i as u32, + shared.slice(i * RECORD_LEN..(i + 1) * RECORD_LEN), + ) + }) + .collect() +} + +#[divan::bench(args = [64, CONNECT_BATCH])] +fn store_prekeys_batch(bencher: Bencher, count: usize) { + bencher + .with_inputs(|| (InMemoryBackend::new(), batch(1, count))) + .bench_refs(|(backend, keys)| { + block_on(backend.store_prekeys_batch(keys, false)).unwrap(); + black_box(&backend); + }); +} + +/// Prekey ids are minted from the monotonic NEXT_PK_ID counter, so a second pass +/// carries ids past the first window's -- every row is new, none overwrites. +#[divan::bench(name = "store_prekeys_batch/populated")] +fn store_prekeys_batch_populated(bencher: Bencher) { + bencher + .with_inputs(|| { + let backend = InMemoryBackend::new(); + block_on(backend.store_prekeys_batch(&batch(1, CONNECT_BATCH), true)).unwrap(); + let next = batch(CONNECT_BATCH as u32 + 1, CONNECT_BATCH); + (backend, next) + }) + .bench_refs(|(backend, keys)| { + block_on(backend.store_prekeys_batch(keys, false)).unwrap(); + black_box(&backend); + }); +} diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index 86f033954..11382c8a9 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -320,6 +320,16 @@ impl SignalStore for InMemoryBackend { async fn store_prekeys_batch(&self, keys: &[(u32, Bytes)], _uploaded: bool) -> Result<()> { let mut state = self.state.lock().await; + // The batch length is known and every id in it is new: prekey ids are + // minted from the monotonic NEXT_PK_ID counter, so a batch never + // overwrites a stored row (unlike `put_msg_secrets`, where reserving a + // mostly-overwrite batch would grow the table for rows it never adds). + // Growing incrementally instead allocates and copies a whole table per + // rehash — a connect-sized batch crosses the load factor eight times. + // 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()); for (id, record) in keys { state.prekeys.insert( *id, @@ -1243,6 +1253,77 @@ mod tests { ); } + /// A connect-sized batch: every id must be readable back, and a later batch + /// repeating an id must overwrite it rather than duplicate or drop it. This + /// is what pins `store_prekeys_batch` idempotent per id across the reserve. + #[tokio::test] + async fn store_prekeys_batch_stores_every_key() { + const COUNT: u32 = 812; + let backend = InMemoryBackend::new(); + + let batch: Vec<(u32, Bytes)> = (1..=COUNT) + .map(|id| (id, Bytes::from(format!("record-{id}")))) + .collect(); + backend.store_prekeys_batch(&batch, false).await.unwrap(); + + for id in 1..=COUNT { + assert_eq!( + backend.load_prekey(id).await.unwrap(), + Some(Bytes::from(format!("record-{id}"))), + "prekey {id} must survive the batch write" + ); + } + assert_eq!(backend.get_max_prekey_id().await.unwrap(), COUNT); + + backend + .store_prekeys_batch(&[(7, Bytes::from_static(b"rewritten"))], true) + .await + .unwrap(); + assert_eq!( + backend.load_prekey(7).await.unwrap(), + Some(Bytes::from_static(b"rewritten")) + ); + assert_eq!( + backend.state.lock().await.prekeys.len(), + COUNT as usize, + "re-storing an existing id must not add a row" + ); + } + + /// Reserving for the batch length must leave the table exactly the size the + /// row count alone demands — the point of the reserve is to reach that size + /// in one allocation, not to reach a bigger one. The control map is grown + /// one insert at a time, which is the un-reserved shape; hashbrown sizes a + /// table from the element count alone, so the two must agree. The second + /// pass covers the reserve landing on an already-populated map, where + /// reserving the full batch length on top of the existing rows would be + /// visible as a doubled table. + #[tokio::test] + async fn store_prekeys_batch_reserve_does_not_over_grow_the_table() { + const COUNT: u32 = 812; + let backend = InMemoryBackend::new(); + let mut control: HashMap = HashMap::new(); + + for pass in 0..2u32 { + let first = pass * COUNT + 1; + let batch: Vec<(u32, Bytes)> = (first..first + COUNT) + .map(|id| (id, Bytes::from_static(b"record"))) + .collect(); + backend.store_prekeys_batch(&batch, false).await.unwrap(); + for id in first..first + COUNT { + control.insert(id, ()); + } + + let state = backend.state.lock().await; + assert_eq!(state.prekeys.len(), control.len()); + assert_eq!( + state.prekeys.capacity(), + control.capacity(), + "pass {pass}: the reserved table must match an incrementally grown one" + ); + } + } + #[tokio::test] async fn group_metadata_round_trip() { use crate::store::traits::ProtocolStore; From 10d40579187712c963e6ceb61f26cb76ad2d3e00 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 21:25:40 +0000 Subject: [PATCH 2/4] docs(observability): point the prekey-reserve ledger note at the right PR Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01A5Vcfvdd5n2PX2Z9M8Ympc --- agent_docs/observability.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agent_docs/observability.md b/agent_docs/observability.md index 91864c646..585d4a9ad 100644 --- a/agent_docs/observability.md +++ b/agent_docs/observability.md @@ -198,7 +198,7 @@ so a per-session diff reads "41.0 KiB in reserve_rehash" and looks like rehash churn. It is not: 1024 buckets × (`size_of::<(u32, PreKeyEntry)>()` + 1 control byte) = 41,984 B is the table that *stays*, and the intermediate tables are all freed before the process peak. `store_prekeys_batch` does reserve for the batch -length (#1266), which cuts the call from 11 allocations / 84.1 KB to 3 / 42.1 KB +length (#1270), which cuts the call from 11 allocations / 84.1 KB to 3 / 42.1 KB and its in-call transient high-water from 63.1 KB to 42.1 KB — but retained is bit-identical at 42,072 B either way, because the final table is the same size. Reserving is worth it for the allocator traffic; it will never move the 41 KiB. From cf70e5e7f715ecb2e8f6ef3dd821eb9ebc5f08df Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 21:40:49 +0000 Subject: [PATCH 3/4] perf(store): reserve only the prekey rows a batch must add MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_01A5Vcfvdd5n2PX2Z9M8Ympc --- wacore/src/store/in_memory.rs | 54 ++++++++++++++++++++++++++++++----- 1 file changed, 47 insertions(+), 7 deletions(-) diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index 11382c8a9..aca69b401 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -320,16 +320,27 @@ impl SignalStore for InMemoryBackend { async fn store_prekeys_batch(&self, keys: &[(u32, Bytes)], _uploaded: bool) -> Result<()> { let mut state = self.state.lock().await; - // The batch length is known and every id in it is new: prekey ids are - // minted from the monotonic NEXT_PK_ID counter, so a batch never - // overwrites a stored row (unlike `put_msg_secrets`, where reserving a - // mostly-overwrite batch would grow the table for rows it never adds). - // Growing incrementally instead allocates and copies a whole table per - // rehash — a connect-sized batch crosses the load factor eight times. + // Growing one insert at a time allocates and copies a whole table per + // rehash, and a connect-sized batch arriving at an empty map crosses + // the load factor eight times. The batch length is known, so the table + // can reach its final size in one allocation instead. + // + // Reserve the batch length MINUS the rows already stored, not the batch + // length: a batch may legally overwrite ids (the trait permits it), and + // a table grown for rows that were only overwritten never shrinks back. + // `keys.len() - len()` is the floor on how many ids must be new, since + // no batch can overwrite more rows than exist — so it can only + // under-reserve (falling back to incremental growth), never inflate the + // resident table. The connect path, where the map is empty and every id + // is freshly minted from the monotonic NEXT_PK_ID counter, reserves the + // whole batch and gets the full win; a replay of a stored window + // reserves nothing and leaves the table exactly as it found it. + // // 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()); + let at_least_new = keys.len().saturating_sub(state.prekeys.len()); + state.prekeys.reserve(at_least_new); for (id, record) in keys { state.prekeys.insert( *id, @@ -1324,6 +1335,35 @@ mod tests { } } + /// Replaying a stored window must not grow the table by one bucket. The + /// trait permits a batch to overwrite ids, and a table grown for rows that + /// were only overwritten never shrinks back — so a reservation taken on the + /// bare batch length would retain an extra table forever, which is exactly + /// the residency this change claims not to touch. + #[tokio::test] + async fn replaying_a_stored_batch_does_not_grow_the_table() { + const COUNT: u32 = 812; + let backend = InMemoryBackend::new(); + let batch: Vec<(u32, Bytes)> = (1..=COUNT) + .map(|id| (id, Bytes::from_static(b"record"))) + .collect(); + + backend.store_prekeys_batch(&batch, false).await.unwrap(); + let settled = backend.state.lock().await.prekeys.capacity(); + + // Same ids twice more: every row is an overwrite, so nothing is added. + backend.store_prekeys_batch(&batch, true).await.unwrap(); + backend.store_prekeys_batch(&batch, true).await.unwrap(); + + let state = backend.state.lock().await; + assert_eq!(state.prekeys.len(), COUNT as usize, "no rows were added"); + assert_eq!( + state.prekeys.capacity(), + settled, + "an all-overwrite batch must not enlarge the table" + ); + } + #[tokio::test] async fn group_metadata_round_trip() { use crate::store::traits::ProtocolStore; From 4dfd8b54978d1be55a6ff186e2304b233998548e Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 21:53:25 +0000 Subject: [PATCH 4/4] perf(store): reserve for a prekey batch only when its ids are distinct MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_01A5Vcfvdd5n2PX2Z9M8Ympc --- wacore/src/store/in_memory.rs | 68 ++++++++++++++++++++++++++++------- 1 file changed, 56 insertions(+), 12 deletions(-) diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index aca69b401..3286eed9e 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -325,22 +325,33 @@ impl SignalStore for InMemoryBackend { // the load factor eight times. The batch length is known, so the table // can reach its final size in one allocation instead. // - // Reserve the batch length MINUS the rows already stored, not the batch - // length: a batch may legally overwrite ids (the trait permits it), and - // a table grown for rows that were only overwritten never shrinks back. - // `keys.len() - len()` is the floor on how many ids must be new, since - // no batch can overwrite more rows than exist — so it can only - // under-reserve (falling back to incremental growth), never inflate the - // resident table. The connect path, where the map is empty and every id - // is freshly minted from the monotonic NEXT_PK_ID counter, reserves the - // whole batch and gets the full win; a replay of a stored window - // reserves nothing and leaves the table exactly as it found it. + // Two things stop that reservation from over-growing a table, because a + // table grown for rows that were never added does not shrink back and + // this is meant to cost no retained bytes: + // + // 1. Subtract the rows already stored. A batch may legally overwrite + // ids, and no batch can overwrite more rows than exist, so + // `keys.len() - len()` is the floor on how many ids must be new. + // 2. Only reserve at all when the batch is strictly ascending, which + // proves its ids are distinct. Without that, 812 entries sharing one + // id would reserve a 1024-bucket table to hold a single row. Testing + // the order costs 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. + // + // Both are one-sided: they can only under-reserve and fall back to + // incremental growth, never inflate the resident table. The connect path + // satisfies both — the map is empty and `upload_pre_keys_pass` emits + // `gen_start + i`, so the whole batch is reserved and gets the full win. // // 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. - let at_least_new = keys.len().saturating_sub(state.prekeys.len()); - state.prekeys.reserve(at_least_new); + 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); + } for (id, record) in keys { state.prekeys.insert( *id, @@ -1364,6 +1375,39 @@ mod tests { ); } + /// A batch whose ids repeat stores one row per distinct id, so sizing the + /// table from the batch length would leave it holding a table for rows that + /// never existed. The reservation is skipped unless the batch is strictly + /// ascending, which is what makes its ids provably distinct. + #[tokio::test] + async fn a_batch_of_repeated_ids_does_not_reserve_for_them() { + const COUNT: usize = 812; + let backend = InMemoryBackend::new(); + let batch: Vec<(u32, Bytes)> = (0..COUNT) + .map(|i| (7, Bytes::from(format!("record-{i}")))) + .collect(); + + backend.store_prekeys_batch(&batch, false).await.unwrap(); + + // One row survives — the last write for id 7 — so the table must be + // sized for one row, not for the 812 entries that were handed over. + let mut control: HashMap = HashMap::new(); + control.insert(7, ()); + + let state = backend.state.lock().await; + assert_eq!(state.prekeys.len(), 1, "last write wins per id"); + assert_eq!( + state.prekeys.capacity(), + control.capacity(), + "a repeated-id batch must not size the table by its length" + ); + drop(state); + assert_eq!( + backend.load_prekey(7).await.unwrap(), + Some(Bytes::from(format!("record-{}", COUNT - 1))) + ); + } + #[tokio::test] async fn group_metadata_round_trip() { use crate::store::traits::ProtocolStore;