perf: cut a dozen per-message allocations from the DM round trip

src/history_sync.rs

@@ -164,11 +164,7 @@ impl HistorySecretSeedCollector {
             Err(_) => return,
         };
         for sender in senders.into_iter().flatten() {
-            let sender_id = if sender.is_same_chat_as(chat) {
-                Arc::clone(&chat_id)
-            } else {
-                Arc::from(sender.to_non_ad_string())
-            };
+            let sender_id = MsgSecretEntry::sender_id_for(chat, &chat_id, &sender);
             self.entries.push(MsgSecretEntry {
                 chat: Arc::clone(&chat_id),
                 sender: sender_id,

src/message/msg_secret.rs

@@ -31,34 +31,43 @@ impl Client {
         let class = wacore::msg_secret::classify(msg, chat_is_bot);
         let message_ts = u64::try_from(info.timestamp.timestamp()).ok();
 
-        // Build both aliases (primary, plus the bot-DM LID key) and write them
-        // in one batch so a partial write can't leave only one stored.
-        let mut entries = Vec::with_capacity(2);
-        if let Some(entry) = self.build_msg_secret_entry(
+        let primary = self.build_msg_secret_entry(
             &info.source.chat,
             &info.source.sender,
             &info.id,
             secret_bytes,
             class,
             message_ts,
-        ) {
-            entries.push(entry);
-        }
+        );
+        // The bot-DM LID key is the only alias a capture ever adds, so a plain
+        // chat writes exactly one row. Deciding that before touching a Vec keeps
+        // the common capture off the batch allocation entirely.
+        let mut bot_alias = None;
         if chat_is_bot
             && let Some(sender) = self.dm_sender_identity_for(&info.source.chat).await
             && sender.to_non_ad() != info.source.sender.to_non_ad()
-            && let Some(entry) = self.build_msg_secret_entry(
+        {
+            bot_alias = self.build_msg_secret_entry(
                 &info.source.chat,
                 &sender,
                 &info.id,
                 secret_bytes,
                 class,
                 message_ts,
-            )
-        {
-            entries.push(entry);
+            );
+        }
+
+        match (primary, bot_alias) {
+            // Both aliases go out in one batch so a partial write can't leave
+            // only one of them stored.
+            (Some(primary), Some(alias)) => {
+                self.persist_msg_secret_entries(vec![primary, alias]).await
+            }
+            (Some(entry), None) | (None, Some(entry)) => {
+                self.msg_secret_buffer.queue_one(entry).await
+            }
+            (None, None) => {}
         }
-        self.persist_msg_secret_entries(entries).await;
     }
 
     /// Build one retention entry, applying the policy gates and computing the
@@ -92,14 +101,14 @@ impl Client {
             message_ts,
             wacore::time::now_secs(),
         );
-        Some(wacore::store::traits::MsgSecretEntry {
-            chat: chat.to_non_ad_string().into(),
-            sender: sender.to_non_ad_string().into(),
-            msg_id: msg_id.into(),
-            secret: *secret,
+        Some(wacore::store::traits::MsgSecretEntry::new(
+            chat,
+            sender,
+            msg_id,
+            *secret,
             expires_at,
-            message_ts: message_ts.and_then(|t| i64::try_from(t).ok()).unwrap_or(0),
-        })
+            message_ts.and_then(|t| i64::try_from(t).ok()).unwrap_or(0),
+        ))
     }
 
     /// Queue a batch of secret aliases on the write-behind buffer: immediately

src/message/receive.rs

@@ -22,6 +22,21 @@ impl ParsedSessionMessage {
     }
 }
 
+/// Append to one of the three per-kind enc buckets, allocating that bucket only
+/// once it actually receives a payload.
+///
+/// A stanza's enc nodes are overwhelmingly all one kind, so reserving all three
+/// buckets up front spent two allocations per message on buffers that stayed
+/// empty for their whole lifetime. Reserving the stanza's full enc count on the
+/// first push keeps a mixed stanza at one allocation per non-empty bucket,
+/// exactly as before.
+fn push_enc_payload(bucket: &mut Vec<EncPayload>, stanza_enc_count: usize, payload: EncPayload) {
+    if bucket.capacity() == 0 {
+        bucket.reserve_exact(stanza_enc_count);
+    }
+    bucket.push(payload);
+}
+
 async fn decrypt_session_message(
     message: &mut ParsedSessionMessage,
     signal_address: &wacore::libsignal::protocol::ProtocolAddress,
@@ -222,9 +237,9 @@ impl Client {
             return None;
         }
 
-        let mut session_payloads = Vec::with_capacity(all_enc_nodes.len());
-        let mut group_payloads = Vec::with_capacity(all_enc_nodes.len());
-        let mut bot_payloads = Vec::with_capacity(all_enc_nodes.len());
+        let mut session_payloads = Vec::new();
+        let mut group_payloads = Vec::new();
+        let mut bot_payloads = Vec::new();
         let mut max_sender_retry_count = 0;
         let mut has_hide_fail = false;
         let mut had_unknown_enc = false;
@@ -301,13 +316,14 @@ impl Client {
                 }
             };
 
-            if payload.enc_type.is_bot_secret() {
-                bot_payloads.push(payload);
+            let bucket = if payload.enc_type.is_bot_secret() {
+                &mut bot_payloads
             } else if payload.enc_type.is_session() {
-                session_payloads.push(payload);
+                &mut session_payloads
             } else {
-                group_payloads.push(payload);
-            }
+                &mut group_payloads
+            };
+            push_enc_payload(bucket, all_enc_nodes.len(), payload);
         }
 
         // WA Web diagnostic: validate skmsg is not first in multi-enc messages.
@@ -1750,6 +1766,76 @@ impl Client {
     }
 }
 
+#[cfg(test)]
+mod enc_bucket_tests {
+    use super::push_enc_payload;
+    use crate::message::EncPayload;
+    use wacore::message_processing::EncType;
+
+    fn payload(enc_type: EncType) -> EncPayload {
+        EncPayload {
+            ciphertext: bytes::Bytes::from_static(b"ct"),
+            enc_type,
+            padding_version: 2,
+        }
+    }
+
+    /// The DM shape: one enc node, one bucket used, two buckets untouched. The
+    /// used bucket must be sized to the stanza exactly (not to `Vec`'s default
+    /// growth step) and the unused ones must own no buffer at all.
+    #[test]
+    fn a_one_enc_stanza_sizes_one_bucket_exactly_and_leaves_the_rest_empty() {
+        let mut used: Vec<EncPayload> = Vec::new();
+        let unused: Vec<EncPayload> = Vec::new();
+
+        push_enc_payload(&mut used, 1, payload(EncType::Message));
+
+        assert_eq!(used.len(), 1);
+        assert_eq!(
+            used.capacity(),
+            1,
+            "the bucket must be reserved to the stanza's enc count, not grown"
+        );
+        assert_eq!(unused.capacity(), 0, "an empty bucket must own no buffer");
+        assert!(unused.is_empty());
+    }
+
+    /// The first push must reserve the stanza's whole enc count, so a multi-enc
+    /// stanza still pays exactly one allocation for the bucket it fills. Eight
+    /// is past `Vec`'s own first growth step, so plain pushes would reallocate.
+    #[test]
+    fn the_first_push_reserves_the_whole_stanza() {
+        const ENC_COUNT: usize = 8;
+        let mut bucket: Vec<EncPayload> = Vec::new();
+        push_enc_payload(&mut bucket, ENC_COUNT, payload(EncType::Message));
+        let reserved = bucket.as_ptr();
+        assert_eq!(bucket.capacity(), ENC_COUNT);
+
+        for _ in 1..ENC_COUNT {
+            push_enc_payload(&mut bucket, ENC_COUNT, payload(EncType::PreKeyMessage));
+        }
+        assert_eq!(bucket.len(), ENC_COUNT);
+        assert_eq!(
+            bucket.as_ptr(),
+            reserved,
+            "filling up to the stanza's enc count must not reallocate"
+        );
+        assert_eq!(bucket[0].enc_type, EncType::Message);
+        assert_eq!(bucket[ENC_COUNT - 1].enc_type, EncType::PreKeyMessage);
+    }
+
+    /// A degenerate count must not make the helper skip its reservation and
+    /// leave the bucket re-reserving on every later push.
+    #[test]
+    fn a_zero_count_still_stores_the_payload() {
+        let mut bucket: Vec<EncPayload> = Vec::new();
+        push_enc_payload(&mut bucket, 0, payload(EncType::SenderKey));
+        assert_eq!(bucket.len(), 1);
+        push_enc_payload(&mut bucket, 0, payload(EncType::SenderKey));
+        assert_eq!(bucket.len(), 2);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::test_utils::create_test_client_with_failing_http;

src/msg_secret_buffer.rs

@@ -616,6 +616,31 @@ mod tests {
         buf.wait_flushed().await;
     }
 
+    /// Building a row costs one allocation per *distinct* identifier and
+    /// nothing else. The naive spelling (`to_non_ad_string().into()` per JID)
+    /// cost five for a direct message; this is the per-message saving on both
+    /// the inbound capture and the outbound persist.
+    #[test]
+    fn entry_construction_allocates_once_per_distinct_identifier() {
+        let chat: wacore_binary::Jid = "5511987650001@s.whatsapp.net".parse().unwrap();
+        let peer_device: wacore_binary::Jid = "5511987650001:33@s.whatsapp.net".parse().unwrap();
+        let me: wacore_binary::Jid = "5511987650002@s.whatsapp.net".parse().unwrap();
+        let secret = [0u8; wacore::reporting_token::MESSAGE_SECRET_SIZE];
+
+        // A DM: chat and sender are the same user, so the row holds two
+        // allocations (the shared identifier and the message id).
+        let dm = crate::test_alloc::min_allocs(2, || {
+            MsgSecretEntry::new(&chat, &peer_device, "3EB0AABBCCDDEEFF0011", secret, 0, 0)
+        });
+        assert_eq!(dm, 2, "a direct-message row: shared identifier + msg id");
+
+        // Outbound: the sender is us, a different user, so it needs its own.
+        let outbound = crate::test_alloc::min_allocs(3, || {
+            MsgSecretEntry::new(&chat, &me, "3EB0AABBCCDDEEFF0011", secret, 0, 0)
+        });
+        assert_eq!(outbound, 3, "distinct users cost one identifier each");
+    }
+
     /// Cloning a buffered entry must stay allocation-free: identifiers share
     /// their Arc allocations and the protocol-sized secret lives inline.
     #[test]

src/send/mod.rs

@@ -2432,14 +2432,9 @@ impl Client {
             u64::try_from(now).ok(),
             now,
         );
-        let entry = wacore::store::traits::MsgSecretEntry {
-            chat: chat.to_non_ad_string().into(),
-            sender: sender.to_non_ad_string().into(),
-            msg_id: msg_id.into(),
-            secret: *secret,
-            expires_at,
-            message_ts: now,
-        };
+        let entry = wacore::store::traits::MsgSecretEntry::new(
+            chat, sender, msg_id, *secret, expires_at, now,
+        );
         // Same write-behind buffer as inbound captures: visible immediately,
         // flushed off the send path (msmsg replies read buffer-first).
         self.msg_secret_buffer.queue_one(entry).await;

wacore/binary/src/jid.rs

@@ -659,6 +659,20 @@ impl Jid {
         buf
     }
 
+    /// [`Self::to_non_ad_string`] as a shareable `Arc<str>`, in exactly one
+    /// allocation. Going through the `String` first costs two — the buffer, then
+    /// the `Arc<str>` its bytes are copied into — and the message-secret rows
+    /// build two of these per message.
+    pub fn to_non_ad_arc_str(&self) -> std::sync::Arc<str> {
+        let mut writer = JidStackWriter::new();
+        if write_jid_fallible(&mut writer, &self.user, self.server, 0, 0).is_ok() {
+            return std::sync::Arc::from(writer.as_str());
+        }
+        // A user part too long for the stack buffer (never seen on the wire)
+        // still renders, just back through the heap.
+        std::sync::Arc::from(self.to_non_ad_string())
+    }
+
     /// Check if this JID matches the user or their LID.
     /// Useful for checking if a participant is "us" in group messages.
     #[inline]
@@ -1976,6 +1990,53 @@ mod tests {
         }
     }
 
+    /// The stack-buffered `Arc<str>` form must render exactly what the `String`
+    /// form does, including for inputs that overflow the stack buffer and fall
+    /// back to the heap, and for multibyte user parts (the buffer is bounded in
+    /// bytes, and a split fragment would be invalid UTF-8).
+    #[test]
+    fn to_non_ad_arc_str_matches_to_non_ad_string() {
+        let long_user = "9".repeat(80);
+        let multibyte_user = "ẞünïcodé-ñ".repeat(3);
+        let owned = [
+            format!("{long_user}:12@s.whatsapp.net"),
+            format!("{multibyte_user}@g.us"),
+            format!("{multibyte_user}@s.whatsapp.net"),
+        ];
+        let cases = [
+            "1234567890:33@s.whatsapp.net",
+            "1234567890@s.whatsapp.net",
+            "100000012345678:25@lid",
+            "867051314767696:0@bot",
+            // Nonzero agent on a bot JID: both forms drop the agent (matching
+            // whatsmeow's ToNonAD), while `is_same_chat_as` still treats it as
+            // identity-significant. Pinning it here keeps that asymmetry
+            // deliberate rather than something a later edit can erase quietly.
+            "867051314767696.5:10@bot",
+            "120363021033254949@g.us",
+            "status@broadcast",
+        ]
+        .into_iter()
+        .chain(owned.iter().map(String::as_str));
+
+        for s in cases {
+            let jid: Jid = s.parse().unwrap_or_else(|e| panic!("parse {s}: {e}"));
+            assert_eq!(
+                &*jid.to_non_ad_arc_str(),
+                jid.to_non_ad_string().as_str(),
+                "mismatch for {s}"
+            );
+        }
+
+        // A default (empty-user) JID has no wire form to render but must still
+        // agree with the String path rather than panic in the stack writer.
+        let empty = Jid::default();
+        assert_eq!(
+            &*empty.to_non_ad_arc_str(),
+            empty.to_non_ad_string().as_str()
+        );
+    }
+
     #[test]
     fn test_into_non_ad_matches_to_non_ad() {
         // into_non_ad (consuming) must produce a JID identical to to_non_ad (cloning).