perf: four more cuts to the per-message allocation count

src/client/messaging.rs

@@ -39,10 +39,16 @@ impl Client {
     ///
     /// Always drains `frames`, including when no socket is installed, while
     /// retaining its outer allocation for the persistent workers to reuse.
+    ///
+    /// Results land in `results`, which the caller owns and reuses too. A
+    /// returned `Vec` would allocate once per burst, and the common burst is a
+    /// single frame, so that allocation was the dominant cost of sending one.
     pub(crate) async fn send_raw_bytes_burst(
         &self,
         frames: &mut Vec<Vec<u8>>,
-    ) -> Result<Vec<crate::socket::error::EncryptSendResult>, ClientError> {
+        results: &mut Vec<crate::socket::error::EncryptSendResult>,
+    ) -> Result<(), ClientError> {
+        results.clear();
         let noise_socket = match self.get_noise_socket().await {
             Ok(socket) => socket,
             Err(error) => {
@@ -52,16 +58,18 @@ impl Client {
         };
         if frames.len() == 1 {
             let plaintext = frames.pop().expect("length checked");
-            return Ok(vec![
+            results.push(
                 noise_socket
                     .encrypt_and_send(bytes::Bytes::from(plaintext))
                     .await,
-            ]);
+            );
+            return Ok(());
         }
         let sends = frames
             .drain(..)
             .map(|plaintext| noise_socket.encrypt_and_send(bytes::Bytes::from(plaintext)));
-        Ok(futures::future::join_all(sends).await)
+        results.extend(futures::future::join_all(sends).await);
+        Ok(())
     }
 
     #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.node", level = "debug", skip_all, fields(tag = %node.tag), err(Debug)))]

src/client/node_io.rs

@@ -697,6 +697,7 @@ impl Client {
             let mut batch = Vec::with_capacity(Self::MAX_ACK_BURST);
             let mut frames = Vec::with_capacity(Self::MAX_ACK_BURST);
             let mut guards = Vec::with_capacity(Self::MAX_ACK_BURST);
+            let mut results = Vec::with_capacity(Self::MAX_ACK_BURST);
             while let Ok(first) = rx.recv().await {
                 let Some(client) = client.upgrade() else {
                     break;
@@ -752,9 +753,9 @@ impl Client {
                 // EnteredSpan is not Send and cannot cross the await.
                 let frame_count = frames.len();
                 let send_and_report = async {
-                    match client.send_raw_bytes_burst(&mut frames).await {
-                        Ok(results) => {
-                            for result in results {
+                    match client.send_raw_bytes_burst(&mut frames, &mut results).await {
+                        Ok(()) => {
+                            for result in results.drain(..) {
                                 if let Err(e) = result
                                     && !e.is_transport_unavailable()
                                 {

src/client/tests.rs

@@ -3409,22 +3409,43 @@ async fn raw_bytes_burst_drains_and_reuses_input_on_happy_paths() {
     let mut frames = Vec::with_capacity(4);
     let retained_capacity = frames.capacity();
     frames.push(vec![0x11; 32]);
-    let single = client
-        .send_raw_bytes_burst(&mut frames)
+    // Sized for the largest burst below, so a reallocation here would mean the
+    // callee replaced the buffer rather than filling it.
+    let mut results = Vec::with_capacity(4);
+    // Captured before the first call, not between the two: taken after it, a
+    // replacement made on the single-frame path would already be the buffer
+    // this compares against and would go unnoticed.
+    let results_ptr = results.as_ptr();
+    client
+        .send_raw_bytes_burst(&mut frames, &mut results)
         .await
         .expect("installed socket");
-    assert_eq!(single.len(), 1);
-    assert!(single.into_iter().all(|result| result.is_ok()));
+    assert_eq!(results.len(), 1);
+    assert!(results.iter().all(|result| result.is_ok()));
+    assert_eq!(
+        results.as_ptr(),
+        results_ptr,
+        "the single-frame path must fill the caller's buffer, not replace it"
+    );
     assert!(frames.is_empty(), "the single-frame fast path must drain");
     assert_eq!(frames.capacity(), retained_capacity);
 
     frames.extend((0..4).map(|index| vec![index; 32]));
-    let burst = client
-        .send_raw_bytes_burst(&mut frames)
+    client
+        .send_raw_bytes_burst(&mut frames, &mut results)
         .await
         .expect("installed socket");
-    assert_eq!(burst.len(), 4);
-    assert!(burst.into_iter().all(|result| result.is_ok()));
+    assert_eq!(results.len(), 4);
+    assert!(results.iter().all(|result| result.is_ok()));
+    // Identity, not capacity: a fresh Vec of the same capacity would satisfy a
+    // capacity check while defeating the whole point of the out-parameter. The
+    // buffer is preallocated above so the second burst cannot legitimately
+    // reallocate it.
+    assert_eq!(
+        results.as_ptr(),
+        results_ptr,
+        "the caller's results buffer must be the same allocation, not an equal one"
+    );
     assert!(frames.is_empty(), "the joined path must drain");
     assert_eq!(frames.capacity(), retained_capacity);
     assert_eq!(transport.sent_count(), 5, "every frame must reach the wire");
@@ -3442,7 +3463,8 @@ async fn raw_bytes_burst_drains_input_when_disconnected() {
     let retained_capacity = frames.capacity();
     frames.extend([vec![0x21; 32], vec![0x22; 32]]);
 
-    let result = client.send_raw_bytes_burst(&mut frames).await;
+    let mut results = Vec::new();
+    let result = client.send_raw_bytes_burst(&mut frames, &mut results).await;
     assert!(
         matches!(result, Err(ClientError::NotConnected)),
         "a missing socket must remain an outer NotConnected error: {result:?}"
@@ -3469,11 +3491,12 @@ async fn raw_bytes_burst_surfaces_transport_then_poisoned_per_frame() {
     let mut frames = Vec::with_capacity(4);
     let retained_capacity = frames.capacity();
     frames.push(vec![0x31; 32]);
-    let mut failed = client
-        .send_raw_bytes_burst(&mut frames)
+    let mut results = Vec::new();
+    client
+        .send_raw_bytes_burst(&mut frames, &mut results)
         .await
         .expect("the socket lookup itself succeeds");
-    let transport_error = failed
+    let transport_error = results
         .pop()
         .expect("one result")
         .expect_err("the transport is configured to fail");
@@ -3486,11 +3509,11 @@ async fn raw_bytes_burst_surfaces_transport_then_poisoned_per_frame() {
     assert_eq!(frames.capacity(), retained_capacity);
 
     frames.push(vec![0x32; 32]);
-    let mut poisoned = client
-        .send_raw_bytes_burst(&mut frames)
+    client
+        .send_raw_bytes_burst(&mut frames, &mut results)
         .await
         .expect("the installed socket remains reachable");
-    let poisoned_error = poisoned
+    let poisoned_error = results
         .pop()
         .expect("one result")
         .expect_err("the sender must reject work after an ambiguous write");
@@ -3524,8 +3547,9 @@ async fn raw_bytes_burst_surfaces_a_closed_sender_per_frame() {
     let mut frames = Vec::with_capacity(4);
     let retained_capacity = frames.capacity();
     frames.push(vec![0x41; 32]);
-    let mut results = client
-        .send_raw_bytes_burst(&mut frames)
+    let mut results = Vec::new();
+    client
+        .send_raw_bytes_burst(&mut frames, &mut results)
         .await
         .expect("the installed socket remains reachable");
     let error = results

src/message/dispatch.rs

@@ -158,6 +158,7 @@ impl Client {
                 let mut batch = Vec::with_capacity(Self::MAX_RECEIPT_BURST);
                 let mut frames = Vec::with_capacity(Self::MAX_RECEIPT_BURST);
                 let mut guards = Vec::with_capacity(Self::MAX_RECEIPT_BURST);
+                let mut results = Vec::with_capacity(Self::MAX_RECEIPT_BURST);
                 while let Ok(first) = rx.recv().await {
                     let Some(client) = client.upgrade() else {
                         break;
@@ -198,9 +199,9 @@ impl Client {
                     // single-receipt path, which this one does not use.
                     let frame_count = frames.len();
                     let send_and_report = async {
-                        match client.send_raw_bytes_burst(&mut frames).await {
-                            Ok(results) => {
-                                for result in results {
+                        match client.send_raw_bytes_burst(&mut frames, &mut results).await {
+                            Ok(()) => {
+                                for result in results.drain(..) {
                                     if let Some(error) = delivery_receipt_burst_warning(&result) {
                                         log::warn!(target: "Client/Receipt", "Failed to send delivery receipt: {error:?}");
                                     }

wacore/libsignal/src/protocol/state/session.rs

@@ -97,6 +97,28 @@ pub(crate) enum ReceiverChainState {
     Closed { next_index: u32 },
 }
 
+/// Writes a chain key's material into the protobuf field, reusing the buffer
+/// already there when it is ours to reuse.
+///
+/// `Bytes` is immutable, so assigning a fresh `copy_from_slice` allocates on
+/// every ratchet advance, and the ratchet advances three times per message
+/// round trip. After a checkout the record is uniquely owned (the cache takes
+/// it out of its `Arc` with `try_unwrap`), so the old buffer is almost always
+/// reusable: take it, overwrite in place, and freeze it back. A buffer that is
+/// still shared, or one of an unexpected length, falls back to allocating,
+/// which is exactly the previous behaviour.
+fn write_chain_key(field: &mut Option<bytes::Bytes>, key: &[u8]) {
+    if let Some(existing) = field.take()
+        && existing.len() == key.len()
+        && let Ok(mut owned) = existing.try_into_mut()
+    {
+        owned.copy_from_slice(key);
+        *field = Some(owned.freeze());
+        return;
+    }
+    *field = Some(bytes::Bytes::copy_from_slice(key));
+}
+
 impl SessionState {
     pub fn from_session_structure(session: SessionStructure) -> Self {
         Self { session }
@@ -456,7 +478,7 @@ impl SessionState {
         match chain.chain_key.as_option_mut() {
             Some(existing) => {
                 existing.index = Some(next_chain_key.index());
-                existing.key = Some(Bytes::copy_from_slice(next_chain_key.key()));
+                write_chain_key(&mut existing.key, next_chain_key.key());
             }
             None => {
                 chain.chain_key = MessageField::some(session_structure::chain::ChainKey {
@@ -600,7 +622,7 @@ impl SessionState {
         match target.as_option_mut() {
             Some(existing) => {
                 existing.index = Some(chain_key.index());
-                existing.key = Some(Bytes::copy_from_slice(chain_key.key()));
+                write_chain_key(&mut existing.key, chain_key.key());
             }
             None => {
                 *target = MessageField::some(session_structure::chain::ChainKey {
@@ -2023,3 +2045,75 @@ mod tests {
         );
     }
 }
+
+#[cfg(test)]
+mod chain_key_buffer_tests {
+    use super::*;
+    use bytes::Bytes;
+
+    /// The ratchet advances three times per message round trip, so the buffer
+    /// this writes into is the difference between one allocation per advance
+    /// and none. Reuse is only valid when the buffer is ours alone.
+    #[test]
+    fn a_uniquely_owned_buffer_is_written_in_place() {
+        let mut field = Some(Bytes::copy_from_slice(&[0u8; 32]));
+        let before = field.as_ref().expect("seeded").as_ptr();
+
+        write_chain_key(&mut field, &[7u8; 32]);
+
+        let after = field.as_ref().expect("written");
+        assert_eq!(
+            after.as_ref(),
+            &[7u8; 32],
+            "the new key must be what is read back"
+        );
+        assert_eq!(
+            after.as_ptr(),
+            before,
+            "a uniquely owned buffer must be reused, not replaced"
+        );
+    }
+
+    /// Bad path: a buffer someone else still holds cannot be overwritten, or
+    /// that holder would observe a key it never asked for. Falling back to a
+    /// fresh allocation is the whole point of the guard.
+    #[test]
+    fn a_shared_buffer_is_never_overwritten() {
+        let shared = Bytes::copy_from_slice(&[1u8; 32]);
+        let observer = shared.clone();
+        let mut field = Some(shared);
+
+        write_chain_key(&mut field, &[9u8; 32]);
+
+        assert_eq!(
+            observer.as_ref(),
+            &[1u8; 32],
+            "the other holder must still see what it had"
+        );
+        assert_eq!(field.expect("written").as_ref(), &[9u8; 32]);
+    }
+
+    /// Bad path: a stored key of the wrong length (a record from an older
+    /// format) must not be partially overwritten, leaving a key that is half
+    /// old and half new.
+    #[test]
+    fn a_buffer_of_the_wrong_length_is_replaced_whole() {
+        let mut field = Some(Bytes::copy_from_slice(&[3u8; 16]));
+
+        write_chain_key(&mut field, &[4u8; 32]);
+
+        let written = field.expect("written");
+        assert_eq!(written.len(), 32, "the new key's length wins");
+        assert_eq!(written.as_ref(), &[4u8; 32]);
+    }
+
+    /// An empty field is the None arm: nothing to reuse, so it allocates.
+    #[test]
+    fn an_absent_buffer_is_created() {
+        let mut field: Option<Bytes> = None;
+
+        write_chain_key(&mut field, &[5u8; 32]);
+
+        assert_eq!(field.expect("written").as_ref(), &[5u8; 32]);
+    }
+}