perf(node): cut ack re-encode, exact stanza sizing, warm session pre-filter

src/client/messaging.rs

@@ -28,7 +28,9 @@ impl Client {
             self.resolve_sent_node_waiters(&Arc::new(node.clone()));
         }
 
-        let plaintext_buf = wacore_binary::marshal::marshal_auto(&node).map_err(|e| {
+        // Exact two-pass sizing: typical stanzas are a few hundred bytes, so
+        // the 1 KiB default reserve of the one-pass path mostly over-allocates.
+        let plaintext_buf = wacore_binary::marshal::marshal_exact(&node).map_err(|e| {
             error!("Failed to marshal node: {e:?}");
             SocketError::Marshal(e)
         })?;

src/client/node_io.rs

@@ -309,7 +309,7 @@ impl Client {
         {
             use wacore::xml::DisplayableNodeRef;
             debug!(target: "Client/Recv", "{}", DisplayableNodeRef(node.get()));
-            self.handle_ack_response_inline(node.get());
+            self.handle_ack_response_owned(node);
             return;
         }
 
@@ -471,7 +471,7 @@ impl Client {
         // retaining router registration for direct router callers.
         let handled = match nr.tag.as_ref() {
             "ack" => {
-                self.handle_ack_response_inline(nr);
+                self.handle_ack_response_arc(&node);
                 true
             }
             "receipt" => {
@@ -1140,19 +1140,49 @@ impl Client {
         })).detach();
     }
 
-    /// Handles incoming `<ack/>` stanzas by resolving pending response waiters.
-    ///
-    /// 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)
+    /// Ack entry point for callers that already share the node: the waiter
+    /// receives an `Arc` clone instead of a ~1 KB re-encode + re-parse.
+    pub(crate) fn handle_ack_response_arc(&self, node: &Arc<wacore_binary::OwnedNodeRef>) -> bool {
+        let Some(waiter) = self.take_ack_waiter(node.get()) else {
+            return false;
+        };
+        if let Err(rejected) = waiter.send(Arc::clone(node)) {
+            Self::warn_ack_waiter_dropped(&rejected);
+        }
+        true
+    }
+
+    /// Ack entry point for the read-loop fast path, which owns the node: the
+    /// `Arc` is built from the existing allocation, and only when a waiter is
+    /// actually waiting.
+    pub(crate) fn handle_ack_response_owned(&self, node: wacore_binary::OwnedNodeRef) -> bool {
+        let Some(waiter) = self.take_ack_waiter(node.get()) else {
+            return false;
+        };
+        if let Err(rejected) = waiter.send(Arc::new(node)) {
+            Self::warn_ack_waiter_dropped(&rejected);
+        }
+        true
+    }
+
+    fn warn_ack_waiter_dropped(rejected: &Arc<wacore_binary::OwnedNodeRef>) {
+        warn!(
+            target: "Client/Ack",
+            "Failed to send ACK response to waiter for ID {:?}. Receiver was likely dropped.",
+            rejected.get().get_attr("id")
+        );
     }
 
+    /// Shared ack prologue: log nack codes, dispatch `ServerAck` when
+    /// observed, and pull the matching response waiter out of the map.
     #[cfg_attr(
         feature = "tracing",
         tracing::instrument(name = "wa.conn.ack_response", level = "debug", skip_all)
     )]
-    pub(crate) fn handle_ack_response_inline(&self, node: &wacore_binary::NodeRef<'_>) -> bool {
+    fn take_ack_waiter(
+        &self,
+        node: &wacore_binary::NodeRef<'_>,
+    ) -> Option<futures::channel::oneshot::Sender<Arc<wacore_binary::OwnedNodeRef>>> {
         let ack_id = node.get_attr("id");
         let ack_error = node.get_attr("error");
 
@@ -1215,28 +1245,8 @@ impl Client {
                 .dispatch(wacore::types::events::Event::ServerAck(ack));
         }
 
-        if let Some(id) = ack_id.map(|v| v.as_str())
-            && let Some(waiter) = self.response_waiters_guard().remove(id.as_ref())
-        {
-            // ACK responses are infrequent; re-encode into OwnedNodeRef for the channel.
-            // marshal_ref prepends a leading 0x00 format byte; OwnedNodeRef::new expects raw
-            // protocol bytes without it, matching what unpack() produces from the network.
-            // slice(1..) drops that byte as a zero-copy view instead of re-allocating.
-            match wacore_binary::marshal::marshal_ref(node).and_then(|buf| {
-                wacore_binary::OwnedNodeRef::new(bytes::Bytes::from(buf).slice(1..))
-            }) {
-                Ok(onr) => {
-                    if waiter.send(Arc::new(onr)).is_err() {
-                        warn!(target: "Client/Ack", "Failed to send ACK response to waiter for ID {id}. Receiver was likely dropped.");
-                    }
-                }
-                Err(e) => {
-                    warn!(target: "Client/Ack", "Failed to re-encode ACK node for waiter: {e}");
-                }
-            }
-            return true;
-        }
-        false
+        let id = ack_id.map(|v| v.as_str())?;
+        self.response_waiters_guard().remove(id.as_ref())
     }
 
     #[cfg_attr(

src/client/sessions.rs

@@ -292,6 +292,23 @@ impl Client {
     async fn ensure_sessions_inner(&self, jids: Vec<Jid>) -> Result<()> {
         use wacore::types::jid::JidExt;
 
+        // Warm-cache pre-filter: a cached session answers synchronously, so
+        // the common live-send case skips the probe-stream machinery below
+        // entirely. Contended or unknown entries fall through to the probe.
+        let jids: Vec<Jid> = jids
+            .into_iter()
+            .filter(|jid| {
+                !matches!(
+                    self.signal_cache
+                        .try_has_session(&jid.to_protocol_address()),
+                    Some(true)
+                )
+            })
+            .collect();
+        if jids.is_empty() {
+            return Ok(());
+        }
+
         let device_snapshot = self.persistence_manager.get_device_snapshot();
 
         // Probe sessions concurrently: a cold-cache multi-recipient ensure would

src/client/tests.rs

@@ -143,7 +143,7 @@ async fn test_ack_waiter_resolves() {
         .build();
 
     // 3. Handle the ack
-    let handled = client.handle_ack_response(&ack_node.as_node_ref()).await;
+    let handled = client.handle_ack_response_arc(&Arc::new(to_owned_node(&ack_node)));
     assert!(
         handled,
         "handle_ack_response should return true when waiter exists"
@@ -197,7 +197,7 @@ async fn test_ack_without_matching_waiter() {
         .build();
 
     // Should return false since there's no waiter
-    let handled = client.handle_ack_response(&ack_node.as_node_ref()).await;
+    let handled = client.handle_ack_response_arc(&Arc::new(to_owned_node(&ack_node)));
     assert!(
         !handled,
         "handle_ack_response should return false when no waiter exists"
@@ -208,6 +208,80 @@ async fn test_ack_without_matching_waiter() {
     );
 }
 
+/// Round-trip a built `Node` into the raw-bytes shape `unpack()` produces
+/// from the network (marshal_ref prepends a 0x00 format byte that
+/// `OwnedNodeRef::new` does not expect).
+fn to_owned_node(node: &Node) -> wacore_binary::OwnedNodeRef {
+    wacore_binary::marshal::marshal_ref(&node.as_node_ref())
+        .and_then(|buf| wacore_binary::OwnedNodeRef::new(bytes::Bytes::from(buf).slice(1..)))
+        .expect("valid node")
+}
+
+fn owned_ack_node(id: &str) -> wacore_binary::OwnedNodeRef {
+    to_owned_node(
+        &NodeBuilder::new("ack")
+            .attr("id", id)
+            .attr("from", SERVER_JID)
+            .build(),
+    )
+}
+
+/// The Arc entry point must hand the waiter the SAME allocation it was given
+/// (no re-encode + re-parse round trip).
+#[tokio::test]
+async fn ack_arc_delivery_shares_allocation() {
+    let client = crate::test_utils::create_test_client().await;
+
+    let test_id = "ack-arc-456";
+    let (tx, rx) = oneshot::channel();
+    client
+        .response_waiters_guard()
+        .insert(test_id.to_string(), tx);
+
+    let node = Arc::new(owned_ack_node(test_id));
+    assert!(client.handle_ack_response_arc(&node));
+
+    let received = tokio::time::timeout(Duration::from_secs(1), rx)
+        .await
+        .expect("waiter should resolve")
+        .expect("sender must not drop");
+    assert!(
+        Arc::ptr_eq(&received, &node),
+        "waiter must receive the original allocation, not a re-encoded copy"
+    );
+
+    // No waiter: must report unhandled without consuming anything.
+    assert!(!client.handle_ack_response_arc(&Arc::new(owned_ack_node("ack-arc-none"))));
+}
+
+/// The owned entry point (read-loop fast path) resolves the waiter from the
+/// node it already owns.
+#[tokio::test]
+async fn ack_owned_delivery_resolves_waiter() {
+    let client = crate::test_utils::create_test_client().await;
+
+    let test_id = "ack-owned-789";
+    let (tx, rx) = oneshot::channel();
+    client
+        .response_waiters_guard()
+        .insert(test_id.to_string(), tx);
+
+    assert!(client.handle_ack_response_owned(owned_ack_node(test_id)));
+    let received = tokio::time::timeout(Duration::from_secs(1), rx)
+        .await
+        .expect("waiter should resolve")
+        .expect("sender must not drop");
+    assert!(
+        received
+            .get()
+            .get_attr("id")
+            .is_some_and(|v| v.as_str() == test_id),
+        "delivered node must carry the ack id"
+    );
+
+    assert!(!client.handle_ack_response_owned(owned_ack_node("ack-owned-none")));
+}
+
 /// Every server `<ack>` with an id dispatches an observe-only
 /// `Event::ServerAck` carrying the ack's class/from/t, independent of
 /// waiter state; a nack carries its error code. Lets consumers measure
@@ -229,7 +303,7 @@ async fn test_ack_dispatches_server_ack_event() {
         .attr("from", "123456789@s.whatsapp.net")
         .attr("t", "1720000000")
         .build();
-    client.handle_ack_response(&ack_node.as_node_ref()).await;
+    client.handle_ack_response_arc(&Arc::new(to_owned_node(&ack_node)));
     assert!(
         collector.events().iter().any(|e| matches!(
             e.as_ref(),
@@ -249,7 +323,7 @@ async fn test_ack_dispatches_server_ack_event() {
         .attr("error", "479")
         .attr("from", SERVER_JID)
         .build();
-    client.handle_ack_response(&nack_node.as_node_ref()).await;
+    client.handle_ack_response_arc(&Arc::new(to_owned_node(&nack_node)));
     assert!(
         collector.events().iter().any(|e| matches!(
             e.as_ref(),
@@ -264,7 +338,7 @@ async fn test_ack_dispatches_server_ack_event() {
 
     // An ack without an id (e.g. non-message acks) dispatches nothing.
     let anon_ack = NodeBuilder::new("ack").attr("from", SERVER_JID).build();
-    client.handle_ack_response(&anon_ack.as_node_ref()).await;
+    client.handle_ack_response_arc(&Arc::new(to_owned_node(&anon_ack)));
     assert_eq!(
         collector
             .events()
@@ -287,7 +361,7 @@ async fn test_ack_dispatches_server_ack_event() {
         .attr("class", "message")
         .attr("from", SERVER_JID)
         .build();
-    let handled = client.handle_ack_response(&waited_ack.as_node_ref()).await;
+    let handled = client.handle_ack_response_arc(&Arc::new(to_owned_node(&waited_ack)));
     assert!(handled, "waiter for the id should have been resolved");
     let resolved = tokio::time::timeout(Duration::from_secs(1), rx)
         .await
@@ -907,6 +981,39 @@ async fn test_ensure_e2e_sessions_waits_for_offline_sync() {
     info!("✅ test_ensure_e2e_sessions_waits_for_offline_sync passed");
 }
 
+/// A warm session cache must satisfy the ensure without any network fetch:
+/// the client here is disconnected, so reaching the usync fetch would error.
+#[tokio::test]
+async fn ensure_sessions_warm_cache_short_circuits() {
+    use wacore::types::jid::JidExt;
+    let client = crate::test_utils::create_test_client().await;
+    let jid: Jid = "15550005555@s.whatsapp.net".parse().unwrap();
+
+    // Cold cache and disconnected: the probe misses, so the fetch runs and
+    // fails — proves the pre-filter does not silently skip unknown sessions.
+    assert!(
+        client
+            .ensure_e2e_sessions_resolved(std::slice::from_ref(&jid))
+            .await
+            .is_err(),
+        "unknown session must still attempt the fetch"
+    );
+
+    assert!(
+        client
+            .signal_cache
+            .try_put_session(
+                &jid.to_protocol_address(),
+                wacore::libsignal::protocol::SessionRecord::new_fresh(),
+            )
+            .is_ok()
+    );
+    client
+        .ensure_e2e_sessions_resolved(&[jid])
+        .await
+        .expect("cached session must satisfy ensure without network");
+}
+
 /// Integration test: Verify that the immediate session establishment does NOT
 /// wait for offline sync. This is critical for PDO to work during offline sync.
 ///
@@ -3177,19 +3284,15 @@ mod counting_alloc {
 async fn ack_miss_path_does_not_heap_allocate() {
     let client = crate::test_utils::create_test_client().await;
 
-    let ack_node = NodeBuilder::new("ack")
-        .attr("id", "3EB0A9252A8F12B7E2")
-        .attr("from", SERVER_JID)
-        .build();
-    let node_ref = ack_node.as_node_ref();
+    let node = Arc::new(owned_ack_node("3EB0A9252A8F12B7E2"));
 
     // Min-delta over many windows: sibling tests share the process-global
     // counter, but their allocations are sporadic. A per-call String shows up
     // in every window, so the minimum only reaches 0 when the path is clean.
     let mut min_delta = u64::MAX;
     for _ in 0..100 {
         let before = counting_alloc::ALLOCS.load(std::sync::atomic::Ordering::Relaxed);
-        let handled = client.handle_ack_response(&node_ref).await;
+        let handled = client.handle_ack_response_arc(&node);
         let after = counting_alloc::ALLOCS.load(std::sync::atomic::Ordering::Relaxed);
         assert!(!handled, "no waiter is registered for this id");
         min_delta = min_delta.min(after - before);

src/handlers/basic.rs

@@ -99,7 +99,7 @@ impl StanzaHandler for AckHandler {
         node: Arc<OwnedNodeRef>,
         _cancelled: &mut bool,
     ) -> bool {
-        client.handle_ack_response(node.get()).await;
+        client.handle_ack_response_arc(&node);
         true
     }
 }