diff --git a/src/client/messaging.rs b/src/client/messaging.rs index f442e63f3..a27749848 100644 --- a/src/client/messaging.rs +++ b/src/client/messaging.rs @@ -21,8 +21,60 @@ 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. + /// + /// 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: &mut Vec>, + ) -> Result, ClientError> { + 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 + .drain(..) + .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 +82,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..e2e3478e0 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -651,6 +651,28 @@ 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 + /// 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 /// client alive. @@ -667,16 +689,104 @@ impl Client { let client = Arc::downgrade(self); self.runtime .spawn(Box::pin(async move { - while let Ok((node, guard)) = rx.recv().await { + // 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; }; - 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. + batch.push(first); + while batch.len() < Self::MAX_ACK_BURST + && let Ok(next) = rx.try_recv() + { + batch.push(next); + } + + // 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. + for (node, guard) in batch.drain(..) { + 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; + } + + // 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. 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:?}"); + } + } + } + }; + #[cfg(feature = "tracing")] { - warn!("Failed to send ack: {e:?}"); + use tracing::Instrument; + send_and_report + .instrument(tracing::trace_span!( + "wa.conn.ack_burst", + frames = frame_count + )) + .await; } - drop(guard); + #[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 5a37dd835..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] @@ -3555,6 +3781,82 @@ 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" + ); +} + +/// 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 6f832c01f..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)))] @@ -127,6 +137,10 @@ impl Client { let _ = tx.try_send((Arc::clone(info), guard)); } + /// 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 /// keeping the client alive. @@ -138,15 +152,135 @@ impl Client { let client = Arc::downgrade(self); self.runtime .spawn(Box::pin(async move { - while let Ok((info, guard)) = rx.recv().await { + // 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; }; - 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. + batch.push(first); + while batch.len() < Self::MAX_RECEIPT_BURST + && let Ok(next) = rx.try_recv() + { + 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. + 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) { + frames.push(frame); + guards.push(guard); + } + } + if frames.is_empty() { + continue; + } + + // 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:?}"); + } + } + } + }; + #[cfg(feature = "tracing")] + { + use tracing::Instrument; + send_and_report + .instrument(tracing::debug_span!( + "wa.receipt.delivery_burst", + frames = frame_count + )) + .await; + } + #[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:?}" + ); + } + } +} 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/socket/noise_socket.rs b/src/socket/noise_socket.rs index 56fd79216..94df69741 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,121 @@ 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"); + // 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}" + ); + } + } + + /// 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 diff --git a/src/transport.rs b/src/transport.rs index a8c289343..e0e3f880a 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,29 @@ 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(split_framed).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() }