diff --git a/src/client/messaging.rs b/src/client/messaging.rs index 501d1435e..5aeefd891 100644 --- a/src/client/messaging.rs +++ b/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) })?; diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 6d4b42d6e..b916ea018 100644 --- a/src/client/node_io.rs +++ b/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 `` 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) -> 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) { + 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>> { 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( diff --git a/src/client/sessions.rs b/src/client/sessions.rs index ac194af33..933e8de61 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -292,6 +292,23 @@ impl Client { async fn ensure_sessions_inner(&self, jids: Vec) -> 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 = 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 diff --git a/src/client/tests.rs b/src/client/tests.rs index e54dd393d..6935ad6da 100644 --- a/src/client/tests.rs +++ b/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 `` 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,11 +3284,7 @@ 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 @@ -3189,7 +3292,7 @@ async fn ack_miss_path_does_not_heap_allocate() { 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); diff --git a/src/handlers/basic.rs b/src/handlers/basic.rs index 44d58445d..947b6e34b 100644 --- a/src/handlers/basic.rs +++ b/src/handlers/basic.rs @@ -99,7 +99,7 @@ impl StanzaHandler for AckHandler { node: Arc, _cancelled: &mut bool, ) -> bool { - client.handle_ack_response(node.get()).await; + client.handle_ack_response_arc(&node); true } }