From 80f793fc5d78776935fee805e76a7f94d10a25f9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sun, 26 Jul 2026 11:29:27 -0300 Subject: [PATCH 1/8] fix(socket): hand a lone waiter the transport error untouched A batch of one is the overwhelmingly common case, and wrapping its failure in SharedSendFailure buried the transport's own error type one level down: downcast_ref inspects the concrete type rather than walking the chain, so a caller with a custom Transport lost the typed error it used to recover. Sharing only pays for itself when there is more than one waiter. --- src/socket/noise_socket.rs | 92 +++++++++++++++++++++++++++++++------- 1 file changed, 77 insertions(+), 15 deletions(-) diff --git a/src/socket/noise_socket.rs b/src/socket/noise_socket.rs index 56fd79216..e1b26b6e3 100644 --- a/src/socket/noise_socket.rs +++ b/src/socket/noise_socket.rs @@ -259,21 +259,34 @@ impl NoiseSocket { } // Every frame in this batch shares the fate of the single write. - // EncryptSendError carries an anyhow::Error and is not Clone, and - // its Display renders only the kind, so flattening the failure to a - // string would hand every caller "transport error" with the actual - // cause discarded. Sharing one Arc keeps the whole chain intact for - // all of them at the cost of a refcount bump each. - let failure = match outcome { - Ok(()) => None, - Err(err) => Some(Arc::new(err)), - }; - for (response_tx, _) in waiters.drain(..) { - let result = match &failure { - None => Ok(()), - Some(err) => Err(EncryptSendError::transport(SharedSendFailure(err.clone()))), - }; - let _ = response_tx.send(result); + match outcome { + Ok(()) => { + for (response_tx, _) in waiters.drain(..) { + let _ = response_tx.send(Ok(())); + } + } + // One waiter owns the failure outright. This is the overwhelmingly + // common case, and handing over the error untouched is what keeps + // `err.source.downcast_ref::()` working for a + // caller with its own Transport: wrapping would bury the typed + // error one level down for no benefit, since there is nobody to + // share it with. + Err(err) if waiters.len() == 1 => { + let (response_tx, _) = waiters.drain(..).next().expect("length checked"); + let _ = response_tx.send(Err(err)); + } + // Several waiters, and EncryptSendError is not Clone: they share + // one Arc. Display renders only the kind, so re-wording per waiter + // would hand each caller "transport error" with the cause gone; + // sharing keeps the whole chain reachable for a refcount bump each. + Err(err) => { + let shared = Arc::new(err); + for (response_tx, _) in waiters.drain(..) { + let _ = response_tx.send(Err(EncryptSendError::transport( + SharedSendFailure(shared.clone()), + ))); + } + } } if let Some((response_tx, err)) = encrypt_failure { let _ = response_tx.send(Err(err)); @@ -674,6 +687,55 @@ mod tests { } } + /// A single-frame send hands its caller the transport's own error, not a + /// wrapper. Callers with a custom `Transport` downcast to their own error + /// type to decide whether a failure is retryable, and `downcast_ref` looks + /// at the concrete type rather than walking the chain, so wrapping the + /// common case would silently break that. + #[tokio::test] + async fn a_lone_waiter_gets_the_transport_error_untouched() { + #[derive(Debug)] + struct TypedTransportError; + impl std::fmt::Display for TypedTransportError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "typed transport error") + } + } + impl std::error::Error for TypedTransportError {} + + struct TypedFailTransport; + + #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] + #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] + impl Transport for TypedFailTransport { + async fn send(&self, _data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> { + Err(anyhow::Error::new(TypedTransportError)) + } + async fn disconnect(&self) {} + } + + let key = [0x77u8; 32]; + let runtime: Arc = Arc::new(crate::runtime_impl::TokioRuntime); + let socket = NoiseSocket::new( + runtime, + Arc::new(TypedFailTransport), + NoiseCipher::new(&key).expect("32-byte key"), + NoiseCipher::new(&key).expect("32-byte key"), + ); + + let err = socket + .encrypt_and_send(bytes::Bytes::from(vec![9u8; 32])) + .await + .expect_err("the transport always fails"); + + assert!(matches!(err.kind, EncryptSendErrorKind::Transport)); + assert!( + err.source.downcast_ref::().is_some(), + "a lone waiter must receive the transport's own error type, got: {:?}", + err.source + ); + } + /// The byte ceiling must hold across a burst. Checking it after appending /// would let a nearly-full batch overshoot by a whole frame, which for a /// large stanza is the difference between a bounded buffer and an unbounded From 413505cb90aab82b1063be091365a2ffffaae9ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sun, 26 Jul 2026 11:38:07 -0300 Subject: [PATCH 2/8] perf(client): send queued acks and receipts as one burst The noise sender coalesces whatever is queued when it wakes, but both workers awaited each send before reading the next, so neither ever had two frames queued at once. Batching only fired when two different producers happened to interleave, which in the pingpong harness was ~9% of frames. Each worker now drains what is already waiting (try_recv only, never a wait for work that has not arrived), marshals the whole burst synchronously, and hands it over in one go. The socket is resolved once before the burst, so the only await left inside each send is the channel push, which resolves on its first poll: arrival order survives, which the ack worker documents and callers rely on. CapturingMockTransport::sent() now splits writes into frames. Its callers all assume one frame per write and decrypt under the write's index as the counter, which silently stopped holding once batches actually formed; sent_writes() and write_count() expose the raw writes for tests that care about transport behaviour instead. --- src/client/messaging.rs | 43 ++++++++++++++++++++++++---- src/client/node_io.rs | 62 +++++++++++++++++++++++++++++++++++++---- src/message/dispatch.rs | 51 +++++++++++++++++++++++++++++++-- src/receipt.rs | 27 ++++++++++++++---- src/socket/error.rs | 3 ++ src/transport.rs | 46 ++++++++++++++++++++++++++++++ 6 files changed, 213 insertions(+), 19 deletions(-) diff --git a/src/client/messaging.rs b/src/client/messaging.rs index f442e63f3..34a8179ea 100644 --- a/src/client/messaging.rs +++ b/src/client/messaging.rs @@ -21,8 +21,43 @@ impl Client { Ok(()) } + /// Send several pre-marshaled stanzas as one burst, returning a result per + /// stanza in the order given. + /// + /// The noise sender coalesces whatever is queued when it wakes, but a + /// worker that awaits each send before starting the next never has two + /// frames queued at once, so the coalescing it was built for never fires. + /// Handing over the whole burst is what turns batching from incidental into + /// the normal case. + /// + /// Order is preserved, which the ack worker depends on. The socket is + /// resolved once up front, so the only await left inside each send is the + /// channel push, and a channel with room resolves that on its first poll: + /// polling the joined sends in order therefore queues them in order. + /// Resolving the socket per send instead would put a contended mutex + /// between the futures and let them queue in any order. + pub(crate) async fn send_raw_bytes_burst( + &self, + frames: Vec>, + ) -> Result, ClientError> { + let noise_socket = self.get_noise_socket().await?; + let sends = frames + .into_iter() + .map(|plaintext| noise_socket.encrypt_and_send(bytes::Bytes::from(plaintext))); + Ok(futures::future::join_all(sends).await) + } + #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.node", level = "debug", skip_all, fields(tag = %node.tag), err(Debug)))] pub async fn send_node(&self, node: Node) -> Result<(), ClientError> { + let plaintext_buf = self.marshal_node_for_send(node)?; + self.send_raw_bytes(plaintext_buf).await + } + + /// Everything [`send_node`](Client::send_node) does short of the send: + /// logging, waiter resolution and marshalling. Split out so a burst can + /// marshal its whole batch before touching the socket, which is what keeps + /// the sends orderable. + pub(crate) fn marshal_node_for_send(&self, node: Node) -> Result, ClientError> { debug!(target: "Client/Send", "{}", DisplayableNode(&node)); if self.sent_node_waiter_count.load(Ordering::Acquire) > 0 { self.resolve_sent_node_waiters(&Arc::new(node.clone())); @@ -30,12 +65,10 @@ impl Client { // 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| { + wacore_binary::marshal::marshal_exact(&node).map_err(|e| { error!("Failed to marshal node: {e:?}"); - SocketError::Marshal(e) - })?; - - self.send_raw_bytes(plaintext_buf).await + SocketError::Marshal(e).into() + }) } #[cfg_attr( diff --git a/src/client/node_io.rs b/src/client/node_io.rs index f61c6bc19..3bac9fae2 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -651,6 +651,11 @@ impl Client { let _ = tx.try_send((node, guard)); } + /// How many queued acks one burst may take. Matches the noise sender's own + /// per-batch frame ceiling: sending more in one go cannot coalesce further, + /// it would only hold flush guards for longer. + const MAX_ACK_BURST: usize = 16; + /// Worker shared by every deferred ack. Holds a `Weak`, so a dropped /// `Client` closes the channel and ends the task instead of keeping the /// client alive. @@ -667,16 +672,63 @@ impl Client { let client = Arc::downgrade(self); self.runtime .spawn(Box::pin(async move { - while let Ok((node, guard)) = rx.recv().await { + while let Ok(first) = rx.recv().await { let Some(client) = client.upgrade() else { break; }; - if let Err(e) = client.send_ack_for(node.get()).await - && !e.is_transport_unavailable() + + // Take everything already waiting, not just the one job that + // woke us. Awaiting each ack before reading the next is what + // kept the noise sender from ever seeing two frames at once, + // so its batching only fired when some *other* producer + // happened to interleave. `try_recv` only: this never waits + // for work that has not arrived. + let mut batch = vec![first]; + while batch.len() < Self::MAX_ACK_BURST + && let Ok(next) = rx.try_recv() { - warn!("Failed to send ack: {e:?}"); + batch.push(next); + } + + // Encoding is synchronous, so the whole burst is marshalled + // before anything is sent and arrival order survives. + let mut frames = Vec::with_capacity(batch.len()); + let mut guards = Vec::with_capacity(batch.len()); + for (node, guard) in batch { + match client.encode_ack_from_snapshot( + node.get(), + AckParticipantPolicy::OmitReceiptDestinationDuplicate, + ) { + Ok(buf) => { + frames.push(buf); + guards.push(guard); + } + // Matches the single-ack path: log and drop this one + // rather than failing the rest of the burst. + Err(e) => warn!("Failed to encode ack: {e}"), + } + } + if frames.is_empty() { + continue; + } + + match client.send_raw_bytes_burst(frames).await { + Ok(results) => { + for result in results { + if let Err(e) = result + && !e.is_transport_unavailable() + { + warn!("Failed to send ack: {e:?}"); + } + } + } + Err(e) => { + if !matches!(e, ClientError::NotConnected) { + warn!("Failed to send ack burst: {e:?}"); + } + } } - drop(guard); + drop(guards); } })) .detach(); diff --git a/src/message/dispatch.rs b/src/message/dispatch.rs index 6f832c01f..8da7ac42f 100644 --- a/src/message/dispatch.rs +++ b/src/message/dispatch.rs @@ -127,6 +127,10 @@ impl Client { let _ = tx.try_send((Arc::clone(info), guard)); } + /// How many queued receipts one burst may take; mirrors the ack worker and + /// the noise sender's own per-batch frame ceiling. + const MAX_RECEIPT_BURST: usize = 16; + /// Worker task shared by every live delivery receipt. Holds only a `Weak` /// so a dropped `Client` closes the channel and ends the task instead of /// keeping the client alive. @@ -138,12 +142,53 @@ impl Client { let client = Arc::downgrade(self); self.runtime .spawn(Box::pin(async move { - while let Ok((info, guard)) = rx.recv().await { + while let Ok(first) = rx.recv().await { let Some(client) = client.upgrade() else { break; }; - client.send_delivery_receipt(&info).await; - drop(guard); + + // Same reasoning as the ack worker: awaiting each receipt + // before reading the next means the noise sender never has + // two frames to coalesce. `try_recv` only, so nothing waits + // on work that has not arrived. + let mut batch = vec![first]; + while batch.len() < Self::MAX_RECEIPT_BURST + && let Ok(next) = rx.try_recv() + { + batch.push(next); + } + + let mut frames = Vec::with_capacity(batch.len()); + let mut guards = Vec::with_capacity(batch.len()); + for (info, guard) in batch { + // Building the node is synchronous, so the burst is + // fully prepared before anything reaches the socket. + if let Some(frame) = client.prepare_delivery_receipt(&info) { + frames.push(frame); + guards.push(guard); + } + } + if frames.is_empty() { + continue; + } + + match client.send_raw_bytes_burst(frames).await { + Ok(results) => { + for result in results { + if let Err(e) = result + && !e.is_transport_unavailable() + { + log::warn!(target: "Client/Receipt", "Failed to send delivery receipt: {e:?}"); + } + } + } + Err(e) => { + if !matches!(e, crate::client::ClientError::NotConnected) { + log::warn!(target: "Client/Receipt", "Failed to send delivery receipt burst: {e:?}"); + } + } + } + drop(guards); } })) .detach(); diff --git a/src/receipt.rs b/src/receipt.rs index b862e5d60..3da400a82 100644 --- a/src/receipt.rs +++ b/src/receipt.rs @@ -591,8 +591,23 @@ impl Client { /// handled by the ack gate, not here). #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.receipt.send_delivery", level = "debug", skip_all, fields(chat = %info.source.chat.observe(), sender = %info.source.sender.observe(), msg_id = %info.id)))] pub(crate) async fn send_delivery_receipt(&self, info: &MessageInfo) { - if !Self::should_send_delivery_receipt(info) { + let Some(frame) = self.prepare_delivery_receipt(info) else { return; + }; + if let Err(e) = self.send_raw_bytes(frame).await + && !matches!(e, crate::client::ClientError::NotConnected) + { + log::warn!(target: "Client/Receipt", "Failed to send delivery receipt for message {}: {:?}", info.id, e); + } + } + + /// Everything [`Self::send_delivery_receipt`] does short of the send: the + /// eligibility gate, node construction, logging and marshalling. Returns + /// `None` when no receipt is owed. Split out so the receipt worker can + /// prepare a whole burst before touching the socket. + pub(crate) fn prepare_delivery_receipt(&self, info: &MessageInfo) -> Option> { + if !Self::should_send_delivery_receipt(info) { + return None; } let receipt_node = build_delivery_receipt_node(info, self.receipts_are_active()); @@ -611,11 +626,11 @@ impl Client { debug!(target: "Client/Receipt", "Sending {} receipt for message {} to {}", receipt_kind.as_wire_str(), info.id, info.source.sender.observe()); - if let Err(e) = self.send_node(receipt_node).await - && !matches!(e, crate::client::ClientError::NotConnected) - { - log::warn!(target: "Client/Receipt", "Failed to send delivery receipt for message {}: {:?}", info.id, e); - } + self.marshal_node_for_send(receipt_node) + .inspect_err(|e| { + log::warn!(target: "Client/Receipt", "Failed to marshal delivery receipt for message {}: {:?}", info.id, e); + }) + .ok() } /// Buffer an offline-drained message's delivery receipt for the aggregate diff --git a/src/socket/error.rs b/src/socket/error.rs index afc027493..10ac76992 100644 --- a/src/socket/error.rs +++ b/src/socket/error.rs @@ -17,6 +17,9 @@ pub enum SocketError { pub type Result = std::result::Result; +/// Outcome of one frame's trip through the noise sender. +pub type EncryptSendResult = std::result::Result<(), EncryptSendError>; + #[derive(Debug, thiserror::Error)] #[non_exhaustive] pub enum EncryptSendErrorKind { diff --git a/src/transport.rs b/src/transport.rs index a8c289343..353120596 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -47,6 +47,31 @@ pub mod mock { } } + /// Splits one transport write into the length-prefixed frames it carries, + /// each returned with its 3-byte prefix intact so callers keep seeing what + /// a single-frame write used to look like. A trailing partial frame (which + /// this sender never produces) is returned as-is rather than dropped. + fn split_framed(write: &bytes::Bytes) -> Vec { + const PREFIX: usize = 3; + let mut frames = Vec::new(); + let mut offset = 0usize; + while offset + PREFIX <= write.len() { + let len = ((write[offset] as usize) << 16) + | ((write[offset + 1] as usize) << 8) + | (write[offset + 2] as usize); + let end = offset + PREFIX + len; + if end > write.len() { + break; + } + frames.push(write.slice(offset..end)); + offset = end; + } + if offset < write.len() { + frames.push(write.slice(offset..)); + } + frames + } + /// Records every `send()` payload so a unit test can assert what the /// client wrote to the wire. pub struct CapturingMockTransport { @@ -64,11 +89,32 @@ pub mod mock { } } + /// Every frame written, one entry each, in write-counter order. + /// + /// The noise sender coalesces queued frames into a single `send()`, so + /// a captured write is not necessarily one frame. Splitting here keeps + /// the whole assertion surface ("the Nth frame is ...", decrypted under + /// counter N) valid whether or not a batch happened to form. pub fn sent(&self) -> Vec { + self.sent_writes() + .iter() + .flat_map(|write| split_framed(write)) + .collect() + } + + /// The raw `send()` payloads, batches included. Use this to assert on + /// transport-level behaviour (how many writes, how large); use + /// [`Self::sent`] to assert on frames. + pub fn sent_writes(&self) -> Vec { self.sent.lock().expect("capturing mutex").clone() } pub fn sent_count(&self) -> usize { + self.sent().len() + } + + /// Number of `send()` calls, as opposed to frames. + pub fn write_count(&self) -> usize { self.sent.lock().expect("capturing mutex").len() } From cb942badf49b932a8fa3d026a0ab87887931281b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sun, 26 Jul 2026 11:39:42 -0300 Subject: [PATCH 3/8] style: drop the redundant closure --- src/transport.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/transport.rs b/src/transport.rs index 353120596..e0e3f880a 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -96,10 +96,7 @@ pub mod mock { /// the whole assertion surface ("the Nth frame is ...", decrypted under /// counter N) valid whether or not a batch happened to form. pub fn sent(&self) -> Vec { - self.sent_writes() - .iter() - .flat_map(|write| split_framed(write)) - .collect() + self.sent_writes().iter().flat_map(split_framed).collect() } /// The raw `send()` payloads, batches included. Use this to assert on From f4d0a520ec11880eadec2429b97c68f40a7d152c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sun, 26 Jul 2026 12:10:47 -0300 Subject: [PATCH 4/8] perf(client): cap the ack and receipt burst at 4 The send-job channel holds 8, so a burst of 16 fills it and makes unrelated producers wait for a slot: the harness showed 29% fewer writes but 3.7% worse pong latency (paired t = 2.8) at that size. At 4 the write saving is ~16% with latency no worse than main. Raising the channel instead was measured too. It recovers the latency but gives back most of the coalescing (-17.6% writes), because a sender that never has to wait consumes jobs one at a time - the queueing pressure is part of what creates something to coalesce. --- src/client/node_io.rs | 14 ++++++++++---- src/message/dispatch.rs | 6 +++--- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 3bac9fae2..b5f4fb7c8 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -651,10 +651,16 @@ impl Client { let _ = tx.try_send((node, guard)); } - /// How many queued acks one burst may take. Matches the noise sender's own - /// per-batch frame ceiling: sending more in one go cannot coalesce further, - /// it would only hold flush guards for longer. - const MAX_ACK_BURST: usize = 16; + /// How many queued acks one burst may take. + /// + /// Measured, not guessed: the send-job channel holds 8, so a larger burst + /// fills it and makes unrelated producers (a reply, a receipt) wait for a + /// slot. At 16 the harness showed 29% fewer writes but 3.7% worse pong + /// latency (paired t = 2.8); at 4 the write saving is ~16% and latency is + /// no worse than main. Raising the channel instead recovers the latency but + /// gives back most of the coalescing, because a sender that never waits + /// consumes jobs one at a time. + const MAX_ACK_BURST: usize = 4; /// Worker shared by every deferred ack. Holds a `Weak`, so a dropped /// `Client` closes the channel and ends the task instead of keeping the diff --git a/src/message/dispatch.rs b/src/message/dispatch.rs index 8da7ac42f..55a468856 100644 --- a/src/message/dispatch.rs +++ b/src/message/dispatch.rs @@ -127,9 +127,9 @@ impl Client { let _ = tx.try_send((Arc::clone(info), guard)); } - /// How many queued receipts one burst may take; mirrors the ack worker and - /// the noise sender's own per-batch frame ceiling. - const MAX_RECEIPT_BURST: usize = 16; + /// How many queued receipts one burst may take; mirrors the ack worker's + /// [`MAX_ACK_BURST`](Client::MAX_ACK_BURST), where the tradeoff is measured. + const MAX_RECEIPT_BURST: usize = 4; /// Worker task shared by every live delivery receipt. Holds only a `Weak` /// so a dropped `Client` closes the channel and ends the task instead of From 0cb61a7c405d4ee083d5a89bc435cce54adfd40a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sun, 26 Jul 2026 12:17:38 -0300 Subject: [PATCH 5/8] fix(observability): restore the spans the burst path skipped Both workers stopped going through the function that carried their tracing span: receipts through send_delivery_receipt (wa.receipt.send_delivery) and acks through send_ack_for (wa.conn.ack). Since the burst path is the normal one, tracing users lost the span for essentially every live receipt and ack. The receipt span moves to prepare_delivery_receipt, which both paths share. Acks get a wa.conn.ack_burst span reporting the burst size instead of N per-ack spans, applied with instrument() because an EnteredSpan is not Send and cannot be held across the await. Also adds order_survives_a_full_job_channel: the ordering guarantee rests on parked senders being woken in queue order once the job channel fills, which nothing covered. --- src/client/node_io.rs | 15 +++++++++- src/receipt.rs | 6 +++- src/socket/noise_socket.rs | 56 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 75 insertions(+), 2 deletions(-) diff --git a/src/client/node_io.rs b/src/client/node_io.rs index b5f4fb7c8..b87c2af96 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -718,7 +718,20 @@ impl Client { continue; } - match client.send_raw_bytes_burst(frames).await { + // The per-ack `wa.conn.ack` span lived in `send_ack_for`, + // which this path no longer calls; a burst reports itself + // once, with its size, rather than N times. `instrument` + // rather than `entered()`: an EnteredSpan is not Send and + // cannot be held across the await. + #[cfg(feature = "tracing")] + let burst = { + use tracing::Instrument; + let span = tracing::trace_span!("wa.conn.ack_burst", frames = frames.len()); + client.send_raw_bytes_burst(frames).instrument(span).await + }; + #[cfg(not(feature = "tracing"))] + let burst = client.send_raw_bytes_burst(frames).await; + match burst { Ok(results) => { for result in results { if let Err(e) = result diff --git a/src/receipt.rs b/src/receipt.rs index 3da400a82..6e5205722 100644 --- a/src/receipt.rs +++ b/src/receipt.rs @@ -589,7 +589,6 @@ impl Client { /// `Send/DeliveryReceiptJob.js`); these are NOT skipped anymore. /// - Newsletters and messages without an ID are skipped (newsletters are /// handled by the ack gate, not here). - #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.receipt.send_delivery", level = "debug", skip_all, fields(chat = %info.source.chat.observe(), sender = %info.source.sender.observe(), msg_id = %info.id)))] pub(crate) async fn send_delivery_receipt(&self, info: &MessageInfo) { let Some(frame) = self.prepare_delivery_receipt(info) else { return; @@ -605,6 +604,11 @@ impl Client { /// eligibility gate, node construction, logging and marshalling. Returns /// `None` when no receipt is owed. Split out so the receipt worker can /// prepare a whole burst before touching the socket. + /// + /// Carries the per-receipt tracing span, because this is the step both the + /// single-receipt and the burst path go through - the span used to sit on + /// [`Self::send_delivery_receipt`], which the worker no longer calls. + #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.receipt.send_delivery", level = "debug", skip_all, fields(chat = %info.source.chat.observe(), sender = %info.source.sender.observe(), msg_id = %info.id)))] pub(crate) fn prepare_delivery_receipt(&self, info: &MessageInfo) -> Option> { if !Self::should_send_delivery_receipt(info) { return None; diff --git a/src/socket/noise_socket.rs b/src/socket/noise_socket.rs index e1b26b6e3..46401880c 100644 --- a/src/socket/noise_socket.rs +++ b/src/socket/noise_socket.rs @@ -687,6 +687,62 @@ mod tests { } } + /// Order must survive a full job channel, not just an empty one. + /// + /// A burst larger than the channel leaves some sends parked waiting for a + /// slot, and the whole ordering guarantee (`send_raw_bytes_burst` promises + /// arrival order, and the ack worker relies on it) then rests on those + /// parked senders being woken in the order they queued. Frame N decrypts + /// only under counter N, so any reordering fails here. + #[tokio::test] + async fn order_survives_a_full_job_channel() { + let key = [0x88u8; 32]; + let transport = GatedTransport::closed(); + let runtime: Arc = Arc::new(crate::runtime_impl::TokioRuntime); + let socket = Arc::new(NoiseSocket::new( + runtime, + transport.clone(), + NoiseCipher::new(&key).expect("32-byte key"), + NoiseCipher::new(&key).expect("32-byte key"), + )); + + // Comfortably past the channel's capacity, so later sends must park. + const FRAMES: usize = 20; + let sends: Vec = (0..FRAMES) + .map(|i| { + let socket = socket.clone(); + Box::pin(async move { + socket + .encrypt_and_send(bytes::Bytes::from(vec![i as u8; 32])) + .await + }) as BoxSend + }) + .collect(); + let mut joined = futures::future::join_all(sends); + assert!( + futures::FutureExt::now_or_never(&mut joined).is_none(), + "the gate is closed, so nothing can have completed" + ); + + transport.gate.add_permits(FRAMES); + for result in joined.await { + result.expect("send must succeed"); + } + + let read_key = NoiseCipher::new(&key).expect("32-byte key"); + let bodies: Vec> = transport + .writes() + .iter() + .flat_map(|w| split_frames(w)) + .collect(); + assert_eq!(bodies.len(), FRAMES, "every frame must reach the wire"); + for (counter, mut body) in bodies.into_iter().enumerate() { + read_key + .decrypt_in_place_with_counter(counter as u32, &mut body) + .expect("a frame written out of counter order cannot authenticate"); + } + } + /// A single-frame send hands its caller the transport's own error, not a /// wrapper. Callers with a custom `Transport` downcast to their own error /// type to decide whether a failure is retryable, and `downcast_ref` looks From 0c2d9584e5900f870afd8d9f3e70f9b01a558c78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sun, 26 Jul 2026 12:36:29 -0300 Subject: [PATCH 6/8] fix(client): keep the teardown gate and the send-side spans on the burst path Three review findings, all real: send_ack_for returns early on expected_disconnect and errors when not connected; bursting called encode+send directly and skipped both, so a 515 or an intentional disconnect could still write stale acks and hold the outbound flush until its timeout. The queue is still drained, exactly as the one-at-a-time worker did. The receipt span was moved onto the synchronous preparation, which closes it before the await: every receipt would look instant and transport stalls would fall outside it. It goes back on the async single-receipt path, and the worker gets an instrumented burst span like the ack worker already has. order_survives_a_full_job_channel asserted only that each frame decrypts under its position's counter, which reordered jobs would also satisfy since they would be encrypted in the order they woke. It now asserts the payload. --- src/client/node_io.rs | 12 ++++++++++++ src/message/dispatch.rs | 15 ++++++++++++++- src/receipt.rs | 6 +----- src/socket/noise_socket.rs | 10 ++++++++++ 4 files changed, 37 insertions(+), 6 deletions(-) diff --git a/src/client/node_io.rs b/src/client/node_io.rs index b87c2af96..34c283261 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -696,6 +696,18 @@ impl Client { batch.push(next); } + // The gate `send_ack_for` applies, which bursting would + // otherwise skip: during an expected teardown (an + // intentional disconnect, or a 515) queued acks are + // deliberately dropped rather than raced against the + // disconnect, and sending them here would also hold the + // outbound flush until its timeout. The queue is still + // drained, exactly as the one-at-a-time worker did. + if client.expected_disconnect.load(Ordering::Relaxed) || !client.is_connected() + { + continue; + } + // Encoding is synchronous, so the whole burst is marshalled // before anything is sent and arrival order survives. let mut frames = Vec::with_capacity(batch.len()); diff --git a/src/message/dispatch.rs b/src/message/dispatch.rs index 55a468856..c01f1da28 100644 --- a/src/message/dispatch.rs +++ b/src/message/dispatch.rs @@ -172,7 +172,20 @@ impl Client { continue; } - match client.send_raw_bytes_burst(frames).await { + // Spans the await, not just the preparation: a receipt that + // stalls in the transport has to show up inside the span. + // The per-receipt `wa.receipt.send_delivery` span stays on + // the single-receipt path, which this one does not use. + #[cfg(feature = "tracing")] + let burst = { + use tracing::Instrument; + let span = + tracing::debug_span!("wa.receipt.delivery_burst", frames = frames.len()); + client.send_raw_bytes_burst(frames).instrument(span).await + }; + #[cfg(not(feature = "tracing"))] + let burst = client.send_raw_bytes_burst(frames).await; + match burst { Ok(results) => { for result in results { if let Err(e) = result diff --git a/src/receipt.rs b/src/receipt.rs index 6e5205722..3da400a82 100644 --- a/src/receipt.rs +++ b/src/receipt.rs @@ -589,6 +589,7 @@ impl Client { /// `Send/DeliveryReceiptJob.js`); these are NOT skipped anymore. /// - Newsletters and messages without an ID are skipped (newsletters are /// handled by the ack gate, not here). + #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.receipt.send_delivery", level = "debug", skip_all, fields(chat = %info.source.chat.observe(), sender = %info.source.sender.observe(), msg_id = %info.id)))] pub(crate) async fn send_delivery_receipt(&self, info: &MessageInfo) { let Some(frame) = self.prepare_delivery_receipt(info) else { return; @@ -604,11 +605,6 @@ impl Client { /// eligibility gate, node construction, logging and marshalling. Returns /// `None` when no receipt is owed. Split out so the receipt worker can /// prepare a whole burst before touching the socket. - /// - /// Carries the per-receipt tracing span, because this is the step both the - /// single-receipt and the burst path go through - the span used to sit on - /// [`Self::send_delivery_receipt`], which the worker no longer calls. - #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.receipt.send_delivery", level = "debug", skip_all, fields(chat = %info.source.chat.observe(), sender = %info.source.sender.observe(), msg_id = %info.id)))] pub(crate) fn prepare_delivery_receipt(&self, info: &MessageInfo) -> Option> { if !Self::should_send_delivery_receipt(info) { return None; diff --git a/src/socket/noise_socket.rs b/src/socket/noise_socket.rs index 46401880c..94df69741 100644 --- a/src/socket/noise_socket.rs +++ b/src/socket/noise_socket.rs @@ -740,6 +740,16 @@ mod tests { read_key .decrypt_in_place_with_counter(counter as u32, &mut body) .expect("a frame written out of counter order cannot authenticate"); + // Decrypting alone would not catch a reorder: jobs that woke out of + // FIFO order would be encrypted in that order too, so their + // counters would still line up. The payload is what pins it - + // unlike the concurrent-producer test, these sends are polled in + // order by one joined future, so submission order is deterministic. + assert_eq!( + body, + vec![counter as u8; 32], + "frame {counter} must carry the payload submitted at position {counter}" + ); } } From 8610a0839faad3d7b99e923a105d0f4104d36fbd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sun, 26 Jul 2026 12:41:00 -0300 Subject: [PATCH 7/8] refactor(client): name the outbound teardown gate and test it The gate the burst path applies was an inline condition duplicating what send_ack_for checks, so nothing would fail if a future edit dropped it. It is now a named predicate with a test covering both signals it folds in. Also records why the receipt worker deliberately has no such gate: the single-receipt path never had one either, and adding it for symmetry would start dropping receipts that today still go out. --- src/client/node_io.rs | 23 ++++++++++++++--------- src/client/tests.rs | 29 +++++++++++++++++++++++++++++ src/message/dispatch.rs | 5 +++++ 3 files changed, 48 insertions(+), 9 deletions(-) diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 34c283261..e13817349 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -651,6 +651,17 @@ impl Client { let _ = tx.try_send((node, guard)); } + /// Whether queued outbound work should be dropped rather than sent. + /// + /// This is the gate [`Self::send_ack_for`] applies before every ack, hoisted + /// so the burst path applies it too: during an expected teardown (an + /// intentional disconnect, or a 515) queued acks are deliberately dropped + /// rather than raced against the disconnect, and sending them anyway would + /// also hold the outbound flush open until its timeout. + pub(crate) fn outbound_teardown_in_progress(&self) -> bool { + self.expected_disconnect.load(Ordering::Relaxed) || !self.is_connected() + } + /// How many queued acks one burst may take. /// /// Measured, not guessed: the send-job channel holds 8, so a larger burst @@ -696,15 +707,9 @@ impl Client { batch.push(next); } - // The gate `send_ack_for` applies, which bursting would - // otherwise skip: during an expected teardown (an - // intentional disconnect, or a 515) queued acks are - // deliberately dropped rather than raced against the - // disconnect, and sending them here would also hold the - // outbound flush until its timeout. The queue is still - // drained, exactly as the one-at-a-time worker did. - if client.expected_disconnect.load(Ordering::Relaxed) || !client.is_connected() - { + // The queue is still drained, exactly as the + // one-at-a-time worker did; only the send is skipped. + if client.outbound_teardown_in_progress() { continue; } diff --git a/src/client/tests.rs b/src/client/tests.rs index 5a37dd835..f2cd704b8 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -3555,6 +3555,35 @@ async fn test_send_ack_for_returns_error_when_disconnected() { ); } +/// The gate that `send_ack_for` applies per ack, and that the burst path +/// applies once per burst, must agree on what counts as teardown. A burst that +/// missed it would write stale acks into a socket that is being torn down and +/// hold the outbound flush open until its timeout. +#[tokio::test] +async fn outbound_teardown_gate_covers_both_disconnect_signals() { + let client = crate::test_utils::create_test_client().await; + + client.set_connected_for_test(true); + client.expected_disconnect.store(false, Ordering::Relaxed); + assert!( + !client.outbound_teardown_in_progress(), + "a live connection must not be treated as tearing down" + ); + + client.expected_disconnect.store(true, Ordering::Relaxed); + assert!( + client.outbound_teardown_in_progress(), + "an expected disconnect (an intentional close, or a 515) must gate sends" + ); + + client.expected_disconnect.store(false, Ordering::Relaxed); + client.set_connected_for_test(false); + assert!( + client.outbound_teardown_in_progress(), + "a disconnected client must gate sends even without the expected flag" + ); +} + /// Verifies that `send_ack_for` returns Ok when expected_disconnect is set, /// since this is an intentional shutdown path. #[tokio::test] diff --git a/src/message/dispatch.rs b/src/message/dispatch.rs index c01f1da28..b64e71d54 100644 --- a/src/message/dispatch.rs +++ b/src/message/dispatch.rs @@ -158,6 +158,11 @@ impl Client { batch.push(next); } + // No teardown gate here, unlike the ack worker: the + // single-receipt path never had one either (it relies on + // the socket reporting NotConnected), and adding one for + // symmetry would silently start dropping receipts that + // today still go out. let mut frames = Vec::with_capacity(batch.len()); let mut guards = Vec::with_capacity(batch.len()); for (info, guard) in batch { From 11c79b76e2ff418fd476744fe95f7b3ddf86e7f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sun, 26 Jul 2026 13:04:11 -0300 Subject: [PATCH 8/8] perf(client): reuse the workers' burst buffers, and keep failures inside the span The two workers allocated batch/frames/guards per burst and handed the frame vector away by value; they now keep all three for the worker's lifetime. Worth it on measurement, not on principle: 173.67 -> 166.83 allocator calls per message (-3.94%, t = -13.0) and -0.92% bytes, harness pingpong at 12k/s. send_raw_bytes_burst takes &mut and always drains, including when no socket is installed. Leaving frames behind would let a caller that forgets to clear resend the same acks on its next burst, which is a duplicate on the wire rather than a wasted allocation; a debug_assert pins the contract at both call sites. The single-frame case skips join_all entirely. The burst result inspection also moves inside the instrumented future. It was running after .instrument(...).await had closed the span, so a failed burst left the span with no error recorded and emitted its warning outside it, unlike the send_ack_for path it replaced. --- src/client/messaging.rs | 23 +++- src/client/node_io.rs | 76 +++++++---- src/client/tests.rs | 273 ++++++++++++++++++++++++++++++++++++++++ src/message/dispatch.rs | 129 ++++++++++++++----- 4 files changed, 442 insertions(+), 59 deletions(-) diff --git a/src/client/messaging.rs b/src/client/messaging.rs index 34a8179ea..a27749848 100644 --- a/src/client/messaging.rs +++ b/src/client/messaging.rs @@ -36,13 +36,30 @@ impl Client { /// polling the joined sends in order therefore queues them in order. /// Resolving the socket per send instead would put a contended mutex /// between the futures and let them queue in any order. + /// + /// Always drains `frames`, including when no socket is installed, while + /// retaining its outer allocation for the persistent workers to reuse. pub(crate) async fn send_raw_bytes_burst( &self, - frames: Vec>, + frames: &mut Vec>, ) -> Result, ClientError> { - let noise_socket = self.get_noise_socket().await?; + let noise_socket = match self.get_noise_socket().await { + Ok(socket) => socket, + Err(error) => { + frames.clear(); + return Err(error); + } + }; + if frames.len() == 1 { + let plaintext = frames.pop().expect("length checked"); + return Ok(vec![ + noise_socket + .encrypt_and_send(bytes::Bytes::from(plaintext)) + .await, + ]); + } let sends = frames - .into_iter() + .drain(..) .map(|plaintext| noise_socket.encrypt_and_send(bytes::Bytes::from(plaintext))); Ok(futures::future::join_all(sends).await) } diff --git a/src/client/node_io.rs b/src/client/node_io.rs index e13817349..e2e3478e0 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -689,6 +689,12 @@ impl Client { let client = Arc::downgrade(self); self.runtime .spawn(Box::pin(async move { + // Reuse the bounded control buffers for the worker's lifetime. + // Encoded payload allocations still move into `Bytes`; only + // the outer storage stays here. + 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); while let Ok(first) = rx.recv().await { let Some(client) = client.upgrade() else { break; @@ -700,7 +706,7 @@ impl Client { // so its batching only fired when some *other* producer // happened to interleave. `try_recv` only: this never waits // for work that has not arrived. - let mut batch = vec![first]; + batch.push(first); while batch.len() < Self::MAX_ACK_BURST && let Ok(next) = rx.try_recv() { @@ -710,14 +716,13 @@ impl Client { // The queue is still drained, exactly as the // one-at-a-time worker did; only the send is skipped. if client.outbound_teardown_in_progress() { + batch.clear(); continue; } // Encoding is synchronous, so the whole burst is marshalled // before anything is sent and arrival order survives. - let mut frames = Vec::with_capacity(batch.len()); - let mut guards = Vec::with_capacity(batch.len()); - for (node, guard) in batch { + for (node, guard) in batch.drain(..) { match client.encode_ack_from_snapshot( node.get(), AckParticipantPolicy::OmitReceiptDestinationDuplicate, @@ -737,34 +742,51 @@ impl Client { // The per-ack `wa.conn.ack` span lived in `send_ack_for`, // which this path no longer calls; a burst reports itself - // once, with its size, rather than N times. `instrument` - // rather than `entered()`: an EnteredSpan is not Send and - // cannot be held across the await. - #[cfg(feature = "tracing")] - let burst = { - use tracing::Instrument; - let span = tracing::trace_span!("wa.conn.ack_burst", frames = frames.len()); - client.send_raw_bytes_burst(frames).instrument(span).await - }; - #[cfg(not(feature = "tracing"))] - let burst = client.send_raw_bytes_burst(frames).await; - match burst { - Ok(results) => { - for result in results { - if let Err(e) = result - && !e.is_transport_unavailable() - { - warn!("Failed to send ack: {e:?}"); + // once, with its size, rather than N times. The result + // inspection is inside the instrumented future, not after + // it: a failure has to be recorded while the span is open, + // the way `send_ack_for`'s `err(Debug)` used to. And + // `instrument` rather than `entered()`, because an + // 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 { + if let Err(e) = result + && !e.is_transport_unavailable() + { + warn!("Failed to send ack: {e:?}"); + } } } - } - Err(e) => { - if !matches!(e, ClientError::NotConnected) { - warn!("Failed to send ack burst: {e:?}"); + Err(e) => { + if !matches!(e, ClientError::NotConnected) { + warn!("Failed to send ack burst: {e:?}"); + } } } + }; + #[cfg(feature = "tracing")] + { + use tracing::Instrument; + send_and_report + .instrument(tracing::trace_span!( + "wa.conn.ack_burst", + frames = frame_count + )) + .await; } - drop(guards); + #[cfg(not(feature = "tracing"))] + { + let _ = frame_count; + send_and_report.await; + } + debug_assert!( + frames.is_empty(), + "send_raw_bytes_burst must always drain its input" + ); + guards.clear(); } })) .detach(); diff --git a/src/client/tests.rs b/src/client/tests.rs index f2cd704b8..60cf8977c 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -3333,6 +3333,25 @@ async fn disconnect_does_not_signal_connection_cleanup_before_outbound_flush() { ); } +async fn install_test_noise_socket( + client: &Arc, + transport: Arc, + runtime: Arc, +) { + use crate::socket::NoiseSocket; + use wacore::handshake::NoiseCipher; + + let key = [0u8; 32]; + let noise_socket = NoiseSocket::new( + runtime, + transport, + NoiseCipher::new(&key).expect("valid key"), + NoiseCipher::new(&key).expect("valid key"), + ); + *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + client.set_connected_for_test(true); +} + fn receipt_test_info(id: &str) -> Arc { Arc::new(crate::types::message::MessageInfo { id: id.to_string(), @@ -3345,6 +3364,180 @@ fn receipt_test_info(id: &str) -> Arc { }) } +#[derive(Debug)] +struct DropSpawnRuntime; + +#[async_trait::async_trait] +impl Runtime for DropSpawnRuntime { + fn spawn( + &self, + _future: std::pin::Pin + Send + 'static>>, + ) -> wacore::runtime::AbortHandle { + // Dropping the sender future closes its receiver synchronously. + wacore::runtime::AbortHandle::noop() + } + + fn sleep(&self, _duration: Duration) -> std::pin::Pin + Send>> { + Box::pin(async {}) + } + + fn spawn_blocking( + &self, + operation: Box, + ) -> std::pin::Pin + Send>> { + Box::pin(async move { operation() }) + } + + fn yield_now(&self) -> Option + Send>>> { + None + } +} + +#[tokio::test] +async fn raw_bytes_burst_drains_and_reuses_input_on_happy_paths() { + use crate::transport::mock::CapturingMockTransport; + + let client = crate::test_utils::create_test_client().await; + let transport = Arc::new(CapturingMockTransport::new()); + install_test_noise_socket( + &client, + transport.clone(), + Arc::new(crate::runtime_impl::TokioRuntime), + ) + .await; + + 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) + .await + .expect("installed socket"); + assert_eq!(single.len(), 1); + assert!(single.into_iter().all(|result| result.is_ok())); + 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) + .await + .expect("installed socket"); + assert_eq!(burst.len(), 4); + assert!(burst.into_iter().all(|result| result.is_ok())); + 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"); + assert_eq!( + transport.write_count(), + 2, + "the four-frame call must remain one coalesced transport write" + ); +} + +#[tokio::test] +async fn raw_bytes_burst_drains_input_when_disconnected() { + let client = crate::test_utils::create_test_client().await; + let mut frames = Vec::with_capacity(4); + let retained_capacity = frames.capacity(); + frames.extend([vec![0x21; 32], vec![0x22; 32]]); + + let result = client.send_raw_bytes_burst(&mut frames).await; + assert!( + matches!(result, Err(ClientError::NotConnected)), + "a missing socket must remain an outer NotConnected error: {result:?}" + ); + assert!(frames.is_empty(), "the outer-error path must also drain"); + assert_eq!(frames.capacity(), retained_capacity); +} + +#[tokio::test] +async fn raw_bytes_burst_surfaces_transport_then_poisoned_per_frame() { + use crate::socket::error::EncryptSendErrorKind; + use crate::transport::mock::CapturingMockTransport; + + let client = crate::test_utils::create_test_client().await; + let transport = Arc::new(CapturingMockTransport::new()); + transport.fail_next_sends(1); + install_test_noise_socket( + &client, + transport.clone(), + Arc::new(crate::runtime_impl::TokioRuntime), + ) + .await; + + 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) + .await + .expect("the socket lookup itself succeeds"); + let transport_error = failed + .pop() + .expect("one result") + .expect_err("the transport is configured to fail"); + assert!(matches!( + transport_error.kind, + EncryptSendErrorKind::Transport + )); + assert!(transport_error.is_transport_unavailable()); + assert!(frames.is_empty()); + assert_eq!(frames.capacity(), retained_capacity); + + frames.push(vec![0x32; 32]); + let mut poisoned = client + .send_raw_bytes_burst(&mut frames) + .await + .expect("the installed socket remains reachable"); + let poisoned_error = poisoned + .pop() + .expect("one result") + .expect_err("the sender must reject work after an ambiguous write"); + assert!(matches!( + poisoned_error.kind, + EncryptSendErrorKind::Poisoned + )); + assert!(poisoned_error.is_transport_unavailable()); + assert!(frames.is_empty()); + assert_eq!(frames.capacity(), retained_capacity); + assert_eq!(transport.failed_sends(), 1); + assert_eq!( + transport.write_count(), + 0, + "a poisoned sender must not attempt another transport write" + ); +} + +#[tokio::test] +async fn raw_bytes_burst_surfaces_a_closed_sender_per_frame() { + use crate::socket::error::EncryptSendErrorKind; + + let client = crate::test_utils::create_test_client().await; + install_test_noise_socket( + &client, + Arc::new(crate::transport::mock::MockTransport), + Arc::new(DropSpawnRuntime), + ) + .await; + + 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) + .await + .expect("the installed socket remains reachable"); + let error = results + .pop() + .expect("one result") + .expect_err("the sender receiver was dropped at construction"); + assert!(matches!(error.kind, EncryptSendErrorKind::ChannelClosed)); + assert!(error.is_transport_unavailable()); + assert!(frames.is_empty()); + assert_eq!(frames.capacity(), retained_capacity); +} + /// Live delivery receipts flow through the persistent worker: the receipt /// reaches the transport and the flush counter returns to zero afterwards. #[tokio::test] @@ -3406,6 +3599,39 @@ async fn delivery_receipt_worker_sends_and_releases_flush() { ); } +/// Transport loss and the poisoned follow-up are reconnect signals, not +/// receipt-worker stalls: both must release their flush guards without a +/// second write attempt. +#[tokio::test] +async fn delivery_receipt_worker_releases_flush_after_transport_and_poisoned_failures() { + use crate::transport::mock::CapturingMockTransport; + + let client = crate::test_utils::create_test_client().await; + let transport = Arc::new(CapturingMockTransport::new()); + transport.fail_next_sends(1); + install_test_noise_socket( + &client, + transport.clone(), + Arc::new(crate::runtime_impl::TokioRuntime), + ) + .await; + + client.ack_received_message(&receipt_test_info("RCPT-FAIL-1")); + crate::test_utils::wait_for_outbound_tasks(&client).await; + assert_eq!(client.outbound_flush.pending(), 0); + assert_eq!(transport.failed_sends(), 1); + + client.ack_received_message(&receipt_test_info("RCPT-POISONED-2")); + crate::test_utils::wait_for_outbound_tasks(&client).await; + assert_eq!(client.outbound_flush.pending(), 0); + assert_eq!( + transport.failed_sends(), + 1, + "the poisoned sender must reject locally instead of touching transport" + ); + assert_eq!(transport.write_count(), 0); +} + /// A closed flush scope (disconnect in progress) drops live receipts without /// leaking the flush counter — mirroring the previous spawn-per-receipt path. #[tokio::test] @@ -3584,6 +3810,53 @@ async fn outbound_teardown_gate_covers_both_disconnect_signals() { ); } +/// Exercise the actual deferred-ack worker, not only its predicate. Dropped +/// teardown batches must release guards, and reusing the batch buffer must not +/// leak either dropped ack into the next live burst. +#[tokio::test] +async fn deferred_ack_worker_drops_teardown_batches_and_recovers_cleanly() { + use crate::transport::mock::CapturingMockTransport; + + let client = crate::test_utils::create_test_client().await; + let transport = Arc::new(CapturingMockTransport::new()); + install_test_noise_socket( + &client, + transport.clone(), + Arc::new(crate::runtime_impl::TokioRuntime), + ) + .await; + let receipt = |id| { + let node = NodeBuilder::new("receipt") + .attr("from", "15550001111@s.whatsapp.net") + .attr("id", id) + .build(); + crate::test_utils::node_to_owned_ref(&node) + }; + + client.expected_disconnect.store(true, Ordering::Relaxed); + client + .process_node(receipt("ACK-EXPECTED-DISCONNECT")) + .await; + crate::test_utils::wait_for_outbound_tasks(&client).await; + assert_eq!(transport.sent_count(), 0); + + client.expected_disconnect.store(false, Ordering::Relaxed); + client.set_connected_for_test(false); + client.process_node(receipt("ACK-DISCONNECTED")).await; + crate::test_utils::wait_for_outbound_tasks(&client).await; + assert_eq!(transport.sent_count(), 0); + + client.set_connected_for_test(true); + client.process_node(receipt("ACK-LIVE")).await; + crate::test_utils::wait_for_outbound_tasks(&client).await; + assert_eq!( + transport.sent_count(), + 1, + "only the live ack may survive into the reusable batch" + ); + assert_eq!(client.outbound_flush.pending(), 0); +} + /// Verifies that `send_ack_for` returns Ok when expected_disconnect is set, /// since this is an intentional shutdown path. #[tokio::test] diff --git a/src/message/dispatch.rs b/src/message/dispatch.rs index b64e71d54..9b25a7b73 100644 --- a/src/message/dispatch.rs +++ b/src/message/dispatch.rs @@ -2,6 +2,16 @@ use super::*; +#[inline] +fn delivery_receipt_burst_warning( + result: &crate::socket::error::EncryptSendResult, +) -> Option<&crate::socket::error::EncryptSendError> { + match result { + Err(error) if !error.is_transport_unavailable() => Some(error), + Ok(()) | Err(_) => None, + } +} + impl Client { /// Dispatches a successfully parsed message to the event bus and sends a delivery receipt. #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.recv.dispatch", level = "debug", skip_all, fields(chat = %info.source.chat.observe(), sender = %info.source.sender.observe(), msg_id = %info.id)))] @@ -142,6 +152,12 @@ impl Client { let client = Arc::downgrade(self); self.runtime .spawn(Box::pin(async move { + // Reuse the bounded control buffers for the worker's lifetime. + // Encoded payload allocations still move into `Bytes`; only + // the outer storage stays here. + 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); while let Ok(first) = rx.recv().await { let Some(client) = client.upgrade() else { break; @@ -151,7 +167,7 @@ impl Client { // before reading the next means the noise sender never has // two frames to coalesce. `try_recv` only, so nothing waits // on work that has not arrived. - let mut batch = vec![first]; + batch.push(first); while batch.len() < Self::MAX_RECEIPT_BURST && let Ok(next) = rx.try_recv() { @@ -163,9 +179,7 @@ impl Client { // the socket reporting NotConnected), and adding one for // symmetry would silently start dropping receipts that // today still go out. - let mut frames = Vec::with_capacity(batch.len()); - let mut guards = Vec::with_capacity(batch.len()); - for (info, guard) in batch { + for (info, guard) in batch.drain(..) { // Building the node is synchronous, so the burst is // fully prepared before anything reaches the socket. if let Some(frame) = client.prepare_delivery_receipt(&info) { @@ -177,39 +191,96 @@ impl Client { continue; } - // Spans the await, not just the preparation: a receipt that - // stalls in the transport has to show up inside the span. - // The per-receipt `wa.receipt.send_delivery` span stays on - // the single-receipt path, which this one does not use. - #[cfg(feature = "tracing")] - let burst = { - use tracing::Instrument; - let span = - tracing::debug_span!("wa.receipt.delivery_burst", frames = frames.len()); - client.send_raw_bytes_burst(frames).instrument(span).await - }; - #[cfg(not(feature = "tracing"))] - let burst = client.send_raw_bytes_burst(frames).await; - match burst { - Ok(results) => { - for result in results { - if let Err(e) = result - && !e.is_transport_unavailable() - { - log::warn!(target: "Client/Receipt", "Failed to send delivery receipt: {e:?}"); + // Spans the await *and* the result inspection: a receipt + // that stalls or fails in the transport has to show up + // inside the span, not after it closed. The per-receipt + // `wa.receipt.send_delivery` span stays on the + // 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 { + if let Some(error) = delivery_receipt_burst_warning(&result) { + log::warn!(target: "Client/Receipt", "Failed to send delivery receipt: {error:?}"); + } } } - } - Err(e) => { - if !matches!(e, crate::client::ClientError::NotConnected) { - log::warn!(target: "Client/Receipt", "Failed to send delivery receipt burst: {e:?}"); + Err(e) => { + if !matches!(e, crate::client::ClientError::NotConnected) { + log::warn!(target: "Client/Receipt", "Failed to send delivery receipt burst: {e:?}"); + } } } + }; + #[cfg(feature = "tracing")] + { + use tracing::Instrument; + send_and_report + .instrument(tracing::debug_span!( + "wa.receipt.delivery_burst", + frames = frame_count + )) + .await; } - drop(guards); + #[cfg(not(feature = "tracing"))] + { + let _ = frame_count; + send_and_report.await; + } + debug_assert!( + frames.is_empty(), + "send_raw_bytes_burst must always drain its input" + ); + guards.clear(); } })) .detach(); tx } } + +#[cfg(test)] +mod tests { + use super::delivery_receipt_burst_warning; + use crate::socket::error::EncryptSendError; + + #[test] + fn receipt_burst_logging_is_quiet_for_success_and_reconnect_failures() { + let reconnect_failures = [ + Ok(()), + Err(EncryptSendError::transport(anyhow::anyhow!( + "transport unavailable" + ))), + Err(EncryptSendError::channel_closed()), + Err(EncryptSendError::poisoned()), + ]; + + for result in &reconnect_failures { + assert!( + delivery_receipt_burst_warning(result).is_none(), + "success and reconnect-related failures must stay quiet: {result:?}" + ); + } + } + + #[test] + fn receipt_burst_logging_keeps_actionable_local_failures_visible() { + let actionable_failures = [ + Err(EncryptSendError::crypto(anyhow::anyhow!( + "encryption failed" + ))), + Err(EncryptSendError::framing(anyhow::anyhow!("framing failed"))), + Err(EncryptSendError::join(anyhow::anyhow!( + "sender join failed" + ))), + ]; + + for result in &actionable_failures { + assert!( + delivery_receipt_burst_warning(result).is_some(), + "local send failures must remain visible: {result:?}" + ); + } + } +}