perf(message): reduce hot-path allocations

src/client/device_registry.rs

@@ -60,11 +60,14 @@ enum UserLookupKeys {
 
 impl UserLookupKeys {
     /// Returns all keys to try for lookups, in preference order.
-    fn all_keys(&self) -> Vec<&str> {
-        match self {
-            Self::LidWithPn { lid, pn } | Self::PnWithLid { lid, pn } => vec![lid, pn],
-            Self::Unknown { user } => vec![user],
-        }
+    fn all_keys(&self) -> impl Iterator<Item = &str> {
+        let (first, second) = match self {
+            Self::LidWithPn { lid, pn } | Self::PnWithLid { lid, pn } => {
+                (lid.as_str(), Some(pn.as_str()))
+            }
+            Self::Unknown { user } => (user.as_str(), None),
+        };
+        std::iter::once(first).chain(second)
     }
 
     /// Returns the canonical (preferred) key for storage.
@@ -311,7 +314,6 @@ impl Client {
         self.resolve_lookup_keys(user)
             .await
             .all_keys()
-            .into_iter()
             .map(String::from)
             .collect()
     }
@@ -329,19 +331,17 @@ impl Client {
             return true;
         }
 
-        // Borrowed `&str` keys (like get_devices_from_registry), bound once so both
-        // loops share one Vec<&str>: avoids the per-message get_lookup_keys churn.
+        // Borrowed keys avoid allocating the owned lookup variants on this hot path.
         let lookup = self.resolve_lookup_keys(user).await;
-        let keys = lookup.all_keys();
 
-        for &key in &keys {
+        for key in lookup.all_keys() {
             if let Some(record) = self.device_registry_cache.get(key).await {
                 return record.devices.iter().any(|d| d.device_id == device_id);
             }
         }
 
         let backend = self.persistence_manager.backend();
-        for &key in &keys {
+        for key in lookup.all_keys() {
             match backend.get_devices(key).await {
                 Ok(Some(record)) => {
                     let has_device = record.devices.iter().any(|d| d.device_id == device_id);
@@ -397,7 +397,6 @@ impl Client {
                 Arc::new(record_for_cache),
                 lookup
                     .all_keys()
-                    .into_iter()
                     .chain(std::iter::once(original_user.as_str())),
             )
             .await;
@@ -464,7 +463,6 @@ impl Client {
                     Arc::new(record_for_cache),
                     lookup
                         .all_keys()
-                        .into_iter()
                         .chain(std::iter::once(original_user.as_str())),
                 )
                 .await;

src/client/node_io.rs

@@ -159,13 +159,29 @@ impl Client {
                                         // - Critical nodes (success/failure/stream:error): inline, required for state
                                         // - Message nodes: inline, preserves arrival order for per-chat queues
                                         //   (MessageHandler just enqueues + ACKs, heavy crypto runs in workers)
+                                        // - Acks/receipts: inline when unobserved; retry work detaches itself
                                         // - ib (in-band): inline, ensures offline sync tracking (expected count)
                                         //   is set up before offline messages are processed
                                         // - Everything else: spawned concurrently for parallelism
-                                        let process_inline = matches!(
-                                            node.tag(),
-                                            "success" | "failure" | "stream:error" | "message" | "ib"
-                                        );
+                                        let process_inline = match node.tag() {
+                                            "success" | "failure" | "stream:error" | "message" | "ib" => true,
+                                            // Preserve concurrent callback behavior when an app
+                                            // observes these events or every raw node.
+                                            "receipt" => {
+                                                !self.synchronous_ack
+                                                    && !self.raw_node_forwarding.load(Ordering::Relaxed)
+                                                    && !self.core.event_bus.has_handler_for(
+                                                        wacore::types::events::EventKind::Receipt,
+                                                    )
+                                            }
+                                            "ack" => {
+                                                !self.raw_node_forwarding.load(Ordering::Relaxed)
+                                                    && !self.core.event_bus.has_handler_for(
+                                                        wacore::types::events::EventKind::ServerAck,
+                                                    )
+                                            }
+                                            _ => false,
+                                        };
 
                                         if process_inline {
                                             self.process_decrypted_node(node).await;
@@ -284,6 +300,19 @@ impl Client {
         self: &Arc<Self>,
         node: wacore_binary::OwnedNodeRef,
     ) {
+        // ACKs need shared ownership only for opt-in raw/node observers. The
+        // usual response-waiter path borrows the node and can skip the Arc.
+        if node.tag() == "ack"
+            && !self.raw_node_forwarding.load(Ordering::Relaxed)
+            && self.node_waiter_count.load(Ordering::Acquire) == 0
+            && !self.offline_sync_metrics.active.load(Ordering::Acquire)
+        {
+            use wacore::xml::DisplayableNodeRef;
+            debug!(target: "Client/Recv", "{}", DisplayableNodeRef(node.get()));
+            self.handle_ack_response_inline(node.get());
+            return;
+        }
+
         // Wrap in Arc once - all handlers will share this same allocation
         let node_arc = Arc::new(node);
         self.process_node(node_arc).await;
@@ -438,13 +467,32 @@ impl Client {
             return;
         }
 
-        // Dispatch to appropriate handler using the router
-        // Clone Arc (cheap - just reference count) not the Node itself
-        if !self
-            .stanza_router
-            .dispatch(self.clone(), Arc::clone(&node), &mut cancelled)
-            .await
-        {
+        // Bypass async_trait's boxed future for the hot built-in handlers while
+        // retaining router registration for direct router callers.
+        let handled = match nr.tag.as_ref() {
+            "ack" => {
+                self.handle_ack_response_inline(nr);
+                true
+            }
+            "receipt" => {
+                self.handle_receipt_inline(Arc::clone(&node));
+                true
+            }
+            "message" => {
+                crate::handlers::message::MessageHandler::handle_inline(
+                    self.clone(),
+                    Arc::clone(&node),
+                    &mut cancelled,
+                )
+                .await
+            }
+            _ => {
+                self.stanza_router
+                    .dispatch(self.clone(), Arc::clone(&node), &mut cancelled)
+                    .await
+            }
+        };
+        if !handled {
             warn!(
                 "Received unknown top-level node: {}",
                 DisplayableNodeRef(nr)
@@ -1096,11 +1144,15 @@ impl Client {
     ///
     /// If an ack with an ID that matches a pending task in `response_waiters`,
     /// the task is resolved and the function returns `true`. Otherwise, returns `false`.
+    pub(crate) async fn handle_ack_response(&self, node: &wacore_binary::NodeRef<'_>) -> bool {
+        self.handle_ack_response_inline(node)
+    }
+
     #[cfg_attr(
         feature = "tracing",
         tracing::instrument(name = "wa.conn.ack_response", level = "debug", skip_all)
     )]
-    pub(crate) async fn handle_ack_response(&self, node: &wacore_binary::NodeRef<'_>) -> bool {
+    pub(crate) fn handle_ack_response_inline(&self, node: &wacore_binary::NodeRef<'_>) -> bool {
         let ack_id = node.get_attr("id");
         let ack_error = node.get_attr("error");
 

src/handlers/message.rs

@@ -15,19 +15,12 @@ const MAX_MESSAGE_DELAY_MS: u64 = 20_000;
 #[derive(Default)]
 pub struct MessageHandler;
 
-#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
-#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
-impl StanzaHandler for MessageHandler {
-    fn tag(&self) -> &'static str {
-        "message"
-    }
-
+impl MessageHandler {
     #[cfg_attr(
         feature = "tracing",
         tracing::instrument(name = "wa.recv.message_enqueue", level = "debug", skip_all)
     )]
-    async fn handle(
-        &self,
+    pub(crate) async fn handle_inline(
         client: Arc<Client>,
         node: Arc<wacore_binary::OwnedNodeRef>,
         cancelled: &mut bool,
@@ -62,6 +55,23 @@ impl StanzaHandler for MessageHandler {
     }
 }
 
+#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
+impl StanzaHandler for MessageHandler {
+    fn tag(&self) -> &'static str {
+        "message"
+    }
+
+    async fn handle(
+        &self,
+        client: Arc<Client>,
+        node: Arc<wacore_binary::OwnedNodeRef>,
+        cancelled: &mut bool,
+    ) -> bool {
+        Self::handle_inline(client, node, cancelled).await
+    }
+}
+
 /// Construct a ChatLane with a spawned worker task. Extracted to keep the
 /// init closure passed to `get_with_by_ref` small.
 fn create_chat_lane(client: &Arc<Client>) -> ChatLane {

src/msg_secret_buffer.rs

@@ -72,28 +72,41 @@ impl MsgSecretWriteBuffer {
             return;
         }
         {
-            use wacore::store::traits::{merge_msg_secret_expiry, merge_msg_secret_message_ts};
             let mut pending = self.pending.lock().unwrap_or_else(|p| p.into_inner());
-            for mut entry in entries {
-                let key = (
-                    entry.chat.clone(),
-                    entry.sender.clone(),
-                    entry.msg_id.clone(),
-                );
-                // Two captures coalescing in the same window must merge the
-                // retention metadata exactly like the backend upsert would
-                // have for two sequential writes (never-expire wins, windows
-                // never shrink, a known parent time is never clobbered).
-                if let Some(existing) = pending.get(&key) {
-                    entry.expires_at =
-                        merge_msg_secret_expiry(existing.expires_at, entry.expires_at);
-                    entry.message_ts =
-                        merge_msg_secret_message_ts(existing.message_ts, entry.message_ts);
-                }
-                pending.insert(key, entry);
+            for entry in entries {
+                Self::insert_pending(&mut pending, entry);
             }
         }
-        // The mutex above orders this load against seal(): an insert that the
+        self.schedule_or_flush().await;
+    }
+
+    /// Single-entry fast path for live sends, avoiding a temporary one-element Vec.
+    pub(crate) async fn queue_one(self: &Arc<Self>, entry: MsgSecretEntry) {
+        {
+            let mut pending = self.pending.lock().unwrap_or_else(|p| p.into_inner());
+            Self::insert_pending(&mut pending, entry);
+        }
+        self.schedule_or_flush().await;
+    }
+
+    fn insert_pending(pending: &mut HashMap<Key, MsgSecretEntry>, mut entry: MsgSecretEntry) {
+        use wacore::store::traits::{merge_msg_secret_expiry, merge_msg_secret_message_ts};
+
+        let key = (
+            entry.chat.clone(),
+            entry.sender.clone(),
+            entry.msg_id.clone(),
+        );
+        // Coalesced captures merge retention metadata like sequential backend writes.
+        if let Some(existing) = pending.get(&key) {
+            entry.expires_at = merge_msg_secret_expiry(existing.expires_at, entry.expires_at);
+            entry.message_ts = merge_msg_secret_message_ts(existing.message_ts, entry.message_ts);
+        }
+        pending.insert(key, entry);
+    }
+
+    async fn schedule_or_flush(self: &Arc<Self>) {
+        // The insertion mutex orders this load against seal(): an insert that the
         // shutdown flush's snapshot missed observes sealed and writes inline.
         if self.sealed.load(Ordering::Acquire) {
             self.flush().await;

src/receipt.rs

@@ -349,11 +349,15 @@ impl Client {
             || info.source.is_self_fanout()
     }
 
+    pub(crate) async fn handle_receipt(self: &Arc<Self>, node: Arc<OwnedNodeRef>) {
+        self.handle_receipt_inline(node);
+    }
+
     #[cfg_attr(
         feature = "tracing",
         tracing::instrument(name = "wa.receipt.handle", level = "debug", skip_all)
     )]
-    pub(crate) async fn handle_receipt(self: &Arc<Self>, node: Arc<OwnedNodeRef>) {
+    pub(crate) fn handle_receipt_inline(self: &Arc<Self>, node: Arc<OwnedNodeRef>) {
         let nr = node.get();
         let mut attrs = nr.attrs();
         let from = attrs.jid("from");

src/send/mod.rs

@@ -2020,7 +2020,7 @@ impl Client {
         };
         // Same write-behind buffer as inbound captures: visible immediately,
         // flushed off the send path (msmsg replies read buffer-first).
-        self.msg_secret_buffer.queue(vec![entry]).await;
+        self.msg_secret_buffer.queue_one(entry).await;
     }
 
     /// Decide the identity (LID vs PN) under which an outbound DM's

wacore/src/store/in_memory.rs

@@ -5,6 +5,7 @@
 //! when the struct is dropped.
 
 use std::collections::HashMap;
+use std::sync::Arc;
 use std::sync::atomic::{AtomicI32, Ordering};
 
 use crate::appstate::hash::HashState;
@@ -146,6 +147,19 @@ impl SignalStore for InMemoryBackend {
         Ok(())
     }
 
+    async fn put_sessions_batch(&self, sessions: &[(Arc<str>, Bytes)]) -> Result<()> {
+        let mut state = self.state.lock().await;
+        state.sessions.reserve(sessions.len());
+        for (address, session) in sessions {
+            if let Some(stored) = state.sessions.get_mut(address.as_ref()) {
+                *stored = session.clone();
+            } else {
+                state.sessions.insert(address.to_string(), session.clone());
+            }
+        }
+        Ok(())
+    }
+
     async fn has_session(&self, address: &str) -> Result<bool> {
         Ok(self.state.lock().await.sessions.contains_key(address))
     }
@@ -845,6 +859,34 @@ mod tests {
         is_backend::<InMemoryBackend>();
     }
 
+    #[tokio::test]
+    async fn put_sessions_batch_inserts_and_updates() {
+        let backend = InMemoryBackend::new();
+        let first: Arc<str> = "15550000001:1@s.whatsapp.net".into();
+        let second: Arc<str> = "15550000002:2@s.whatsapp.net".into();
+
+        backend
+            .put_sessions_batch(&[
+                (first.clone(), Bytes::from_static(b"first")),
+                (second.clone(), Bytes::from_static(b"second")),
+            ])
+            .await
+            .unwrap();
+        backend
+            .put_sessions_batch(&[(first.clone(), Bytes::from_static(b"updated"))])
+            .await
+            .unwrap();
+
+        assert_eq!(
+            backend.get_session(&first).await.unwrap().unwrap(),
+            Bytes::from_static(b"updated")
+        );
+        assert_eq!(
+            backend.get_session(&second).await.unwrap().unwrap(),
+            Bytes::from_static(b"second")
+        );
+    }
+
     #[tokio::test]
     async fn group_metadata_round_trip() {
         use crate::store::traits::ProtocolStore;