From 7dc850c1cd857058219644431c23e527d5bbdf17 Mon Sep 17 00:00:00 2001 From: Zach Heylmun Date: Mon, 29 Jun 2026 15:02:40 -0400 Subject: [PATCH 1/7] refactor: make TxChannelPayload response channel optional Adds Option so a future fire-and-forget send can skip the confirmation round-trip, plus shared take_terminal_error helpers. No behavior change; existing tests green. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/lib.rs | 8 ++------ src/socket_loop.rs | 35 ++++++++++++++++++++++++----------- 2 files changed, 26 insertions(+), 17 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 143ee45..2b5a2bd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -189,11 +189,7 @@ where /// The cause is surfaced once: the first caller to observe the closed /// connection consumes it, and subsequent observers see `WebsocketClosed`. fn take_terminal_error(&self) -> Error { - self.terminal_error - .lock() - .unwrap_or_else(PoisonError::into_inner) - .take() - .unwrap_or(Error::WebsocketClosed) + socket_loop::take_terminal_error(&self.terminal_error) } /// Encode and send a message via the connection's [`Codec`]. @@ -237,7 +233,7 @@ where self.sender .send(TxChannelPayload { message, - response_tx: tx, + response_tx: Some(tx), }) .await .map_err(|_| self.take_terminal_error())?; diff --git a/src/socket_loop.rs b/src/socket_loop.rs index 219aaa4..a9503d2 100644 --- a/src/socket_loop.rs +++ b/src/socket_loop.rs @@ -29,11 +29,11 @@ use tracing::{error, info, instrument, trace}; use crate::Error; /// A frame to send, paired with a one-shot the loop uses to report the result -/// of the underlying `sink.send`. +/// of the underlying `sink.send`. `None` for fire-and-forget sends. #[derive(Debug)] pub(crate) struct TxChannelPayload { pub(crate) message: Message, - pub(crate) response_tx: oneshot::Sender>, + pub(crate) response_tx: Option>>, } pub(crate) type WebSocketStreamType = WebSocketStream>; @@ -51,6 +51,19 @@ enum LoopState { /// channels drop. pub(crate) type TerminalError = Arc>>; +/// Take the recorded terminal error, if any. The cause is surfaced once: the +/// first caller consumes it, later callers see `None`. A graceful close records +/// nothing, so this returns `None`. +pub(crate) fn take_terminal_error_opt(slot: &TerminalError) -> Option { + slot.lock().unwrap_or_else(PoisonError::into_inner).take() +} + +/// Like [`take_terminal_error_opt`] but falls back to [`Error::WebsocketClosed`] +/// when no specific cause was recorded. +pub(crate) fn take_terminal_error(slot: &TerminalError) -> Error { + take_terminal_error_opt(slot).unwrap_or(Error::WebsocketClosed) +} + /// Send a close frame via the tx channel and wait for confirmation. pub(crate) async fn send_close(sender: &mpsc::Sender) -> Result<(), Error> { let (tx, rx) = oneshot::channel::>(); @@ -60,7 +73,7 @@ pub(crate) async fn send_close(sender: &mpsc::Sender) -> Resul code: tungstenite::protocol::frame::coding::CloseCode::Normal, reason: Utf8Bytes::from_static("Closing Connection"), })), - response_tx: tx, + response_tx: Some(tx), }) .await .map_err(|_| Error::WebsocketClosed)?; @@ -123,15 +136,15 @@ async fn send_socket_message( trace!("Sending message: {:?}", message); let send_result = sink.send(message.message).await.map_err(Error::from); let socket_error = send_result.is_err(); - match message.response_tx.send(send_result) { - Ok(()) => { - if socket_error { - LoopState::Error(Error::WebsocketClosed) - } else { - LoopState::Running - } + if let Some(response_tx) = message.response_tx { + if response_tx.send(send_result).is_err() { + return LoopState::Error(Error::SocketeerDroppedWithoutClosing); } - Err(_) => LoopState::Error(Error::SocketeerDroppedWithoutClosing), + } + if socket_error { + LoopState::Error(Error::WebsocketClosed) + } else { + LoopState::Running } } else { #[cfg(feature = "tracing")] From dc7f9ff7a4c5bcca8b2beb7f49fe603bcd9c2801 Mon Sep 17 00:00:00 2001 From: Zach Heylmun Date: Mon, 29 Jun 2026 15:07:27 -0400 Subject: [PATCH 2/7] feat: add Socketeer::split into cloneable Tx and receive Rx halves SocketeerTx (Clone) with send/send_raw/close; SocketeerRx with next_message/next_raw_message. Adds Socketeer::from_parts and a shared send_confirmed helper. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/lib.rs | 75 +++++++++++++++++++++----- src/socket_loop.rs | 20 +++++++ src/split.rs | 126 +++++++++++++++++++++++++++++++++++++++++++ tests/integration.rs | 41 ++++++++++++++ 4 files changed, 248 insertions(+), 14 deletions(-) create mode 100644 src/split.rs diff --git a/src/lib.rs b/src/lib.rs index 2b5a2bd..358a4c5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,7 @@ mod handler; #[cfg(feature = "mocking")] mod mock_server; mod socket_loop; +mod split; #[cfg(feature = "msgpack")] pub use codec::MsgPackCodec; @@ -18,14 +19,15 @@ pub use handler::{ConnectionHandler, HandshakeContext, NoopHandler}; pub use mock_server::msgpack_echo_server; #[cfg(feature = "mocking")] pub use mock_server::{EchoControlMessage, auth_echo_server, echo_server, get_mock_address}; +pub use split::{SocketeerRx, SocketeerTx}; pub(crate) use socket_loop::WebSocketStreamType; -use socket_loop::{TerminalError, TxChannelPayload, send_close, socket_loop_split}; +use socket_loop::{TerminalError, TxChannelPayload, send_close, send_confirmed, socket_loop_split}; use std::sync::{Arc, Mutex, PoisonError}; use futures::StreamExt; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::mpsc; use tokio_tungstenite::{connect_async, tungstenite::Message}; #[cfg(feature = "tracing")] @@ -165,6 +167,31 @@ where }) } + /// Reassemble a `Socketeer` from its parts. Used by + /// [`SocketeerRx::reunite`](crate::SocketeerRx::reunite). + #[allow(clippy::too_many_arguments)] + pub(crate) fn from_parts( + url: Url, + options: ConnectOptions, + codec: C, + handler: Handler, + receiver: mpsc::Receiver, + sender: mpsc::Sender, + socket_handle: tokio::task::JoinHandle<()>, + terminal_error: TerminalError, + ) -> Self { + Self { + url, + options, + codec, + handler, + receiver, + sender, + socket_handle, + terminal_error, + } + } + /// Wait for the next message from the WebSocket connection, decoded by the /// connection's [`Codec`]. /// @@ -229,18 +256,7 @@ where /// /// - If the WebSocket connection is closed, or otherwise errored pub async fn send_raw(&self, message: Message) -> Result<(), Error> { - let (tx, rx) = oneshot::channel::>(); - self.sender - .send(TxChannelPayload { - message, - response_tx: Some(tx), - }) - .await - .map_err(|_| self.take_terminal_error())?; - match rx.await { - Ok(result) => result, - Err(_) => unreachable!("Socket loop always sends response before dropping one-shot"), - } + send_confirmed(&self.sender, &self.terminal_error, message).await } /// Consume self, closing down any remaining send/receive, and return a new Socketeer instance if successful. @@ -302,3 +318,34 @@ where } } } + +impl Socketeer +where + C: Codec + Clone, + Handler: ConnectionHandler, +{ + /// Split into an owned, cloneable send half and an owned receive half. + /// + /// A pure move — no new tasks are spawned and the existing channels are + /// reused. The codec is cloned into both halves (the send half encodes, + /// the receive half decodes). Recombine with + /// [`SocketeerRx::reunite`](crate::SocketeerRx::reunite). + #[must_use] + pub fn split(self) -> (SocketeerTx, SocketeerRx) { + let tx = SocketeerTx { + sender: self.sender, + codec: self.codec.clone(), + terminal_error: Arc::clone(&self.terminal_error), + }; + let rx = SocketeerRx { + receiver: self.receiver, + codec: self.codec, + terminal_error: self.terminal_error, + socket_handle: self.socket_handle, + url: self.url, + options: self.options, + handler: self.handler, + }; + (tx, rx) + } +} diff --git a/src/socket_loop.rs b/src/socket_loop.rs index a9503d2..fb04243 100644 --- a/src/socket_loop.rs +++ b/src/socket_loop.rs @@ -64,6 +64,26 @@ pub(crate) fn take_terminal_error(slot: &TerminalError) -> Error { take_terminal_error_opt(slot).unwrap_or(Error::WebsocketClosed) } +/// Enqueue a frame and await the loop's result for this specific send. +pub(crate) async fn send_confirmed( + sender: &mpsc::Sender, + terminal_error: &TerminalError, + message: Message, +) -> Result<(), Error> { + let (tx, rx) = oneshot::channel::>(); + sender + .send(TxChannelPayload { + message, + response_tx: Some(tx), + }) + .await + .map_err(|_| take_terminal_error(terminal_error))?; + match rx.await { + Ok(result) => result, + Err(_) => unreachable!("Socket loop always sends response before dropping one-shot"), + } +} + /// Send a close frame via the tx channel and wait for confirmation. pub(crate) async fn send_close(sender: &mpsc::Sender) -> Result<(), Error> { let (tx, rx) = oneshot::channel::>(); diff --git a/src/split.rs b/src/split.rs new file mode 100644 index 0000000..55b4319 --- /dev/null +++ b/src/split.rs @@ -0,0 +1,126 @@ +//! Owned, splittable handles for concurrent send/receive use. +//! +//! [`Socketeer::split`](crate::Socketeer::split) divides a connected client +//! into a cloneable [`SocketeerTx`] send half and a [`SocketeerRx`] receive +//! half that can move into separate tasks. [`SocketeerRx::reunite`] recombines +//! them to regain the full [`Socketeer`](crate::Socketeer) handle. + +use std::sync::Arc; + +use tokio::sync::mpsc; +use tokio::task::JoinHandle; +use tokio_tungstenite::tungstenite::Message; +use url::Url; + +use crate::socket_loop::{ + TerminalError, TxChannelPayload, send_close, send_confirmed, take_terminal_error, +}; +use crate::{Codec, ConnectOptions, ConnectionHandler, Error, NoopHandler}; + +/// The cloneable send half of a [`Socketeer`](crate::Socketeer), produced by +/// [`Socketeer::split`](crate::Socketeer::split). Clone it to send from +/// multiple tasks concurrently. +pub struct SocketeerTx { + pub(crate) sender: mpsc::Sender, + pub(crate) codec: C, + pub(crate) terminal_error: TerminalError, +} + +impl Clone for SocketeerTx { + fn clone(&self) -> Self { + Self { + sender: self.sender.clone(), + codec: self.codec.clone(), + terminal_error: Arc::clone(&self.terminal_error), + } + } +} + +impl std::fmt::Debug for SocketeerTx { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SocketeerTx").finish_non_exhaustive() + } +} + +impl SocketeerTx { + /// Encode and send a message, awaiting confirmation that this specific send + /// reached the socket. See [`Socketeer::send`](crate::Socketeer::send). + /// # Errors + /// - If the codec fails to encode the value + /// - If the connection is closed or errored + pub async fn send(&self, message: C::Tx) -> Result<(), Error> { + let encoded = self.codec.encode(&message)?; + self.send_raw(encoded).await + } + + /// Send a raw [`Message`], awaiting confirmation. + /// # Errors + /// - If the connection is closed or errored + pub async fn send_raw(&self, message: Message) -> Result<(), Error> { + send_confirmed(&self.sender, &self.terminal_error, message).await + } + + /// Send a graceful close frame. The receive half observes its stream + /// ending as the signal that the connection has closed. + /// # Errors + /// - If the connection is already closed + pub async fn close(&self) -> Result<(), Error> { + send_close(&self.sender).await + } +} + +/// The receive half of a [`Socketeer`](crate::Socketeer), produced by +/// [`Socketeer::split`](crate::Socketeer::split). Implements +/// [`Stream`](futures::Stream) and can be recombined with a [`SocketeerTx`] via +/// [`reunite`](Self::reunite). +pub struct SocketeerRx +where + Handler: ConnectionHandler, +{ + pub(crate) receiver: mpsc::Receiver, + pub(crate) codec: C, + pub(crate) terminal_error: TerminalError, + pub(crate) socket_handle: JoinHandle<()>, + pub(crate) url: Url, + pub(crate) options: ConnectOptions, + pub(crate) handler: Handler, +} + +impl std::fmt::Debug + for SocketeerRx +where + Handler: ConnectionHandler, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SocketeerRx") + .field("url", &self.url) + .finish_non_exhaustive() + } +} + +impl SocketeerRx +where + Handler: ConnectionHandler, +{ + /// Wait for the next decoded message. See + /// [`Socketeer::next_message`](crate::Socketeer::next_message). + /// # Errors + /// - If the connection is closed or errored + /// - If the codec fails to decode the frame + pub async fn next_message(&mut self) -> Result { + match self.receiver.recv().await { + Some(message) => self.codec.decode(&message), + None => Err(take_terminal_error(&self.terminal_error)), + } + } + + /// Wait for the next raw [`Message`] without decoding. + /// # Errors + /// - If the connection is closed or errored + pub async fn next_raw_message(&mut self) -> Result { + match self.receiver.recv().await { + Some(message) => Ok(message), + None => Err(take_terminal_error(&self.terminal_error)), + } + } +} diff --git a/tests/integration.rs b/tests/integration.rs index f99db15..dfcf634 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -569,6 +569,47 @@ async fn test_handshake_recv_text_rejects_binary() { socketeer.close_connection().await.unwrap(); } +#[tokio::test] +async fn test_split_round_trip() { + let server_address = get_mock_address(echo_server).await; + let socketeer: Socketeer = Socketeer::connect(&format!("ws://{server_address}")) + .await + .unwrap(); + let (tx, mut rx) = socketeer.split(); + let message = EchoControlMessage::Message("split hello".into()); + tx.send(message.clone()).await.unwrap(); + assert_eq!(rx.next_message().await.unwrap(), message); + tx.close().await.unwrap(); +} + +#[tokio::test] +async fn test_split_cloned_sender_concurrent() { + let server_address = get_mock_address(echo_server).await; + let socketeer: Socketeer = Socketeer::connect(&format!("ws://{server_address}")) + .await + .unwrap(); + let (tx, mut rx) = socketeer.split(); + let tx2 = tx.clone(); + let message = EchoControlMessage::Message("from clone".into()); + tx2.send(message.clone()).await.unwrap(); + assert_eq!(rx.next_message().await.unwrap(), message); + tx.close().await.unwrap(); +} + +#[tokio::test] +async fn test_tx_close_ends_receive() { + let server_address = get_mock_address(echo_server).await; + let socketeer: Socketeer = Socketeer::connect(&format!("ws://{server_address}")) + .await + .unwrap(); + let (tx, mut rx) = socketeer.split(); + tx.close().await.unwrap(); + assert!(matches!( + rx.next_message().await.unwrap_err(), + Error::WebsocketClosed + )); +} + #[tokio::test] async fn test_binary_custom_keepalive() { // The widening of custom_keepalive_message from Option to From 35a9b8bddc0a42ade8f1e8cce58bb25092b587e8 Mon Sep 17 00:00:00 2001 From: Zach Heylmun Date: Mon, 29 Jun 2026 15:12:31 -0400 Subject: [PATCH 3/7] feat: add fire-and-forget send_unconfirmed to SocketeerTx Backpressure-respecting send that skips the confirmation round-trip; errors surface via the terminal error rather than per-send. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/socket_loop.rs | 17 +++++++++++++++++ src/split.rs | 22 +++++++++++++++++++++- tests/integration.rs | 19 +++++++++++++++++++ 3 files changed, 57 insertions(+), 1 deletion(-) diff --git a/src/socket_loop.rs b/src/socket_loop.rs index fb04243..1d7c932 100644 --- a/src/socket_loop.rs +++ b/src/socket_loop.rs @@ -84,6 +84,23 @@ pub(crate) async fn send_confirmed( } } +/// Enqueue a frame without awaiting the loop's per-send result. Applies +/// backpressure (awaits channel capacity) but returns `Ok` as soon as the +/// frame is queued; a socket-send failure surfaces via the terminal error. +pub(crate) async fn send_unconfirmed( + sender: &mpsc::Sender, + terminal_error: &TerminalError, + message: Message, +) -> Result<(), Error> { + sender + .send(TxChannelPayload { + message, + response_tx: None, + }) + .await + .map_err(|_| take_terminal_error(terminal_error)) +} + /// Send a close frame via the tx channel and wait for confirmation. pub(crate) async fn send_close(sender: &mpsc::Sender) -> Result<(), Error> { let (tx, rx) = oneshot::channel::>(); diff --git a/src/split.rs b/src/split.rs index 55b4319..424452c 100644 --- a/src/split.rs +++ b/src/split.rs @@ -13,7 +13,8 @@ use tokio_tungstenite::tungstenite::Message; use url::Url; use crate::socket_loop::{ - TerminalError, TxChannelPayload, send_close, send_confirmed, take_terminal_error, + TerminalError, TxChannelPayload, send_close, send_confirmed, send_unconfirmed, + take_terminal_error, }; use crate::{Codec, ConnectOptions, ConnectionHandler, Error, NoopHandler}; @@ -60,6 +61,25 @@ impl SocketeerTx { send_confirmed(&self.sender, &self.terminal_error, message).await } + /// Encode and send a message fire-and-forget: applies backpressure but does + /// not await the loop's per-send result, skipping the round-trip. Returns + /// `Ok(())` once enqueued; a later socket-send failure surfaces via the + /// receive stream / terminal error. + /// # Errors + /// - If the codec fails to encode the value + /// - If the connection is already closed + pub async fn send_unconfirmed(&self, message: C::Tx) -> Result<(), Error> { + let encoded = self.codec.encode(&message)?; + self.send_raw_unconfirmed(encoded).await + } + + /// Send a raw [`Message`] fire-and-forget. See [`Self::send_unconfirmed`]. + /// # Errors + /// - If the connection is already closed + pub async fn send_raw_unconfirmed(&self, message: Message) -> Result<(), Error> { + send_unconfirmed(&self.sender, &self.terminal_error, message).await + } + /// Send a graceful close frame. The receive half observes its stream /// ending as the signal that the connection has closed. /// # Errors diff --git a/tests/integration.rs b/tests/integration.rs index dfcf634..c3ae9f1 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -610,6 +610,25 @@ async fn test_tx_close_ends_receive() { )); } +#[tokio::test] +async fn test_send_unconfirmed_delivers() { + let server_address = get_mock_address(echo_server).await; + let socketeer: Socketeer = Socketeer::connect(&format!("ws://{server_address}")) + .await + .unwrap(); + let (tx, mut rx) = socketeer.split(); + let message = EchoControlMessage::Message("unconfirmed".into()); + // Fire-and-forget: returns once enqueued, no round-trip. + tx.send_unconfirmed(message.clone()).await.unwrap(); + // The echo server still delivers it. + assert_eq!(rx.next_message().await.unwrap(), message); + // Liveness: a confirmed send still round-trips afterward. + let confirmed = EchoControlMessage::Message("after".into()); + tx.send(confirmed.clone()).await.unwrap(); + assert_eq!(rx.next_message().await.unwrap(), confirmed); + tx.close().await.unwrap(); +} + #[tokio::test] async fn test_binary_custom_keepalive() { // The widening of custom_keepalive_message from Option to From 7a3f80e34e0609a6468d7c598e3a4e642f612294 Mon Sep 17 00:00:00 2001 From: Zach Heylmun Date: Mon, 29 Jun 2026 15:15:52 -0400 Subject: [PATCH 4/7] feat: implement Stream for SocketeerRx Yields Result; surfaces the real terminal cause once on abrupt close then ends, ends cleanly on graceful close, and continues past a per-message decode error. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/split.rs | 28 +++++++++++++++++++++++++++- tests/integration.rs | 43 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/src/split.rs b/src/split.rs index 424452c..b275a23 100644 --- a/src/split.rs +++ b/src/split.rs @@ -5,8 +5,11 @@ //! half that can move into separate tasks. [`SocketeerRx::reunite`] recombines //! them to regain the full [`Socketeer`](crate::Socketeer) handle. +use std::pin::Pin; use std::sync::Arc; +use std::task::{Context, Poll}; +use futures::Stream; use tokio::sync::mpsc; use tokio::task::JoinHandle; use tokio_tungstenite::tungstenite::Message; @@ -14,7 +17,7 @@ use url::Url; use crate::socket_loop::{ TerminalError, TxChannelPayload, send_close, send_confirmed, send_unconfirmed, - take_terminal_error, + take_terminal_error, take_terminal_error_opt, }; use crate::{Codec, ConnectOptions, ConnectionHandler, Error, NoopHandler}; @@ -144,3 +147,26 @@ where } } } + +impl Stream + for SocketeerRx +where + Handler: ConnectionHandler + Unpin, +{ + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + match this.receiver.poll_recv(cx) { + Poll::Ready(Some(message)) => Poll::Ready(Some(this.codec.decode(&message))), + Poll::Ready(None) => { + // Channel closed: surface the recorded cause once, then end. + match take_terminal_error_opt(&this.terminal_error) { + Some(error) => Poll::Ready(Some(Err(error))), + None => Poll::Ready(None), + } + } + Poll::Pending => Poll::Pending, + } + } +} diff --git a/tests/integration.rs b/tests/integration.rs index c3ae9f1..49552ee 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -9,6 +9,7 @@ use bytes::Bytes; use tokio::time::sleep; use tokio_tungstenite::tungstenite::Message; +use futures::StreamExt; use socketeer::{ Codec, ConnectOptions, ConnectionHandler, EchoControlMessage, Error, HandshakeContext, JsonCodec, RawCodec, Socketeer, auth_echo_server, echo_server, get_mock_address, @@ -629,6 +630,48 @@ async fn test_send_unconfirmed_delivers() { tx.close().await.unwrap(); } +#[tokio::test] +async fn test_rx_stream_round_trip() { + let server_address = get_mock_address(echo_server).await; + let socketeer: Socketeer = Socketeer::connect(&format!("ws://{server_address}")) + .await + .unwrap(); + let (tx, mut rx) = socketeer.split(); + let message = EchoControlMessage::Message("stream hello".into()); + tx.send(message.clone()).await.unwrap(); + let item = rx.next().await.unwrap().unwrap(); + assert_eq!(item, message); + tx.close().await.unwrap(); +} + +#[tokio::test] +async fn test_rx_stream_abrupt_close_yields_error_then_ends() { + let server_address = get_mock_address(echo_server).await; + let socketeer: Socketeer = Socketeer::connect(&format!("ws://{server_address}")) + .await + .unwrap(); + let (tx, mut rx) = socketeer.split(); + tx.send(EchoControlMessage::Abort).await.unwrap(); + let err = rx.next().await.unwrap().unwrap_err(); + assert!( + matches!(err, Error::WebsocketError(_)), + "expected WebsocketError, got {err:?}" + ); + assert!(rx.next().await.is_none(), "stream must end after the error"); +} + +#[tokio::test] +async fn test_rx_stream_graceful_close_ends_cleanly() { + let server_address = get_mock_address(echo_server).await; + let socketeer: Socketeer = Socketeer::connect(&format!("ws://{server_address}")) + .await + .unwrap(); + let (tx, mut rx) = socketeer.split(); + tx.close().await.unwrap(); + // Graceful close records no cause, so the stream just ends. + assert!(rx.next().await.is_none()); +} + #[tokio::test] async fn test_binary_custom_keepalive() { // The widening of custom_keepalive_message from Option to From 5eb9aa152014e688def240f425c7fad731cee1f5 Mon Sep 17 00:00:00 2001 From: Zach Heylmun Date: Mon, 29 Jun 2026 15:20:43 -0400 Subject: [PATCH 5/7] feat: add SocketeerRx::reunite and ReuniteError Recombines the split halves into a full Socketeer, validated via Arc::ptr_eq on the shared terminal-error slot; mismatched halves are returned in ReuniteError. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/lib.rs | 2 +- src/split.rs | 72 +++++++++++++++++++++++++++++++++++++++++++- tests/integration.rs | 39 ++++++++++++++++++++++++ 3 files changed, 111 insertions(+), 2 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 358a4c5..1e68036 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,7 +19,7 @@ pub use handler::{ConnectionHandler, HandshakeContext, NoopHandler}; pub use mock_server::msgpack_echo_server; #[cfg(feature = "mocking")] pub use mock_server::{EchoControlMessage, auth_echo_server, echo_server, get_mock_address}; -pub use split::{SocketeerRx, SocketeerTx}; +pub use split::{ReuniteError, SocketeerRx, SocketeerTx}; pub(crate) use socket_loop::WebSocketStreamType; use socket_loop::{TerminalError, TxChannelPayload, send_close, send_confirmed, socket_loop_split}; diff --git a/src/split.rs b/src/split.rs index b275a23..7adcac2 100644 --- a/src/split.rs +++ b/src/split.rs @@ -19,7 +19,7 @@ use crate::socket_loop::{ TerminalError, TxChannelPayload, send_close, send_confirmed, send_unconfirmed, take_terminal_error, take_terminal_error_opt, }; -use crate::{Codec, ConnectOptions, ConnectionHandler, Error, NoopHandler}; +use crate::{Codec, ConnectOptions, ConnectionHandler, Error, NoopHandler, Socketeer}; /// The cloneable send half of a [`Socketeer`](crate::Socketeer), produced by /// [`Socketeer::split`](crate::Socketeer::split). Clone it to send from @@ -146,6 +146,76 @@ where None => Err(take_terminal_error(&self.terminal_error)), } } + + /// Recombine with a [`SocketeerTx`] to restore the full + /// [`Socketeer`](crate::Socketeer) handle (and thus `close_connection` / + /// `reconnect`). + /// + /// # Errors + /// Returns [`ReuniteError`] (carrying both halves back) if `tx` did not come + /// from the same [`split`](crate::Socketeer::split) as `self`. + // ReuniteError intentionally carries both halves back to the caller; boxing + // it would change the public API and force callers to dereference to destructure. + #[allow(clippy::result_large_err)] + pub fn reunite( + self, + tx: SocketeerTx, + ) -> Result, ReuniteError> { + if Arc::ptr_eq(&self.terminal_error, &tx.terminal_error) { + Ok(Socketeer::from_parts( + self.url, + self.options, + self.codec, + self.handler, + self.receiver, + tx.sender, + self.socket_handle, + self.terminal_error, + )) + } else { + Err(ReuniteError { tx, rx: self }) + } + } +} + +/// Error returned by [`SocketeerRx::reunite`] when the two halves did not come +/// from the same [`Socketeer::split`](crate::Socketeer::split). Carries both +/// halves back so the caller can recover them. +pub struct ReuniteError +where + Handler: ConnectionHandler, +{ + /// The send half passed to `reunite`. + pub tx: SocketeerTx, + /// The receive half `reunite` was called on. + pub rx: SocketeerRx, +} + +impl std::fmt::Debug + for ReuniteError +where + Handler: ConnectionHandler, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ReuniteError").finish_non_exhaustive() + } +} + +impl std::fmt::Display + for ReuniteError +where + Handler: ConnectionHandler, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("tried to reunite a SocketeerTx and SocketeerRx from different connections") + } +} + +impl std::error::Error + for ReuniteError +where + Handler: ConnectionHandler, +{ } impl Stream diff --git a/tests/integration.rs b/tests/integration.rs index 49552ee..86c6219 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -672,6 +672,45 @@ async fn test_rx_stream_graceful_close_ends_cleanly() { assert!(rx.next().await.is_none()); } +#[tokio::test] +async fn test_reunite_then_reconnect() { + let server_address = get_mock_address(echo_server).await; + let socketeer: Socketeer = Socketeer::connect(&format!("ws://{server_address}")) + .await + .unwrap(); + let (tx, rx) = socketeer.split(); + let socketeer = rx.reunite(tx).expect("matching halves reunite"); + // Full handle restored: reconnect works. + let mut socketeer = socketeer.reconnect().await.unwrap(); + let message = EchoControlMessage::Message("after reunite".into()); + socketeer.send(message.clone()).await.unwrap(); + assert_eq!(socketeer.next_message().await.unwrap(), message); + socketeer.close_connection().await.unwrap(); +} + +#[tokio::test] +async fn test_reunite_mismatch_returns_halves() { + // Two separate servers: the mock server's accept loop is serial (it blocks + // in echo_server until the connection closes), so two concurrent + // connections must target two addresses or the second connect hangs. + let addr_a = get_mock_address(echo_server).await; + let addr_b = get_mock_address(echo_server).await; + let sock_a: Socketeer = Socketeer::connect(&format!("ws://{addr_a}")).await.unwrap(); + let sock_b: Socketeer = Socketeer::connect(&format!("ws://{addr_b}")).await.unwrap(); + let (tx_a, rx_a) = sock_a.split(); + let (tx_b, rx_b) = sock_b.split(); + // Mismatched halves: rx from A, tx from B. + let err = rx_a + .reunite(tx_b) + .expect_err("mismatched halves must not reunite"); + let socketeer::ReuniteError { tx: tx_b, rx: rx_a } = err; + // Both halves survive and reunite with their correct partners. + let sock_a = rx_a.reunite(tx_a).expect("correct halves reunite"); + let sock_b = rx_b.reunite(tx_b).expect("correct halves reunite"); + sock_a.close_connection().await.unwrap(); + sock_b.close_connection().await.unwrap(); +} + #[tokio::test] async fn test_binary_custom_keepalive() { // The widening of custom_keepalive_message from Option to From 27f16d81e27a9ce94f07e7e15548fce3a5a609ca Mon Sep 17 00:00:00 2001 From: Zach Heylmun Date: Mon, 29 Jun 2026 15:24:33 -0400 Subject: [PATCH 6/7] feat: implement Stream for Socketeer Lets the un-split handle be consumed with .next()/combinators, same disconnect semantics as SocketeerRx. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/lib.rs | 24 +++++++++++++++++++++++- tests/integration.rs | 12 ++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/src/lib.rs b/src/lib.rs index 1e68036..dc01125 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -24,9 +24,11 @@ pub use split::{ReuniteError, SocketeerRx, SocketeerTx}; pub(crate) use socket_loop::WebSocketStreamType; use socket_loop::{TerminalError, TxChannelPayload, send_close, send_confirmed, socket_loop_split}; +use std::pin::Pin; use std::sync::{Arc, Mutex, PoisonError}; +use std::task::{Context, Poll}; -use futures::StreamExt; +use futures::{Stream, StreamExt}; use tokio::sync::mpsc; use tokio_tungstenite::{connect_async, tungstenite::Message}; @@ -73,6 +75,26 @@ where } } +impl Stream + for Socketeer +where + Handler: ConnectionHandler + Unpin, +{ + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + match this.receiver.poll_recv(cx) { + Poll::Ready(Some(message)) => Poll::Ready(Some(this.codec.decode(&message))), + Poll::Ready(None) => match socket_loop::take_terminal_error_opt(&this.terminal_error) { + Some(error) => Poll::Ready(Some(Err(error))), + None => Poll::Ready(None), + }, + Poll::Pending => Poll::Pending, + } + } +} + impl Socketeer where C: Codec + Default, diff --git a/tests/integration.rs b/tests/integration.rs index 86c6219..379dc79 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -711,6 +711,18 @@ async fn test_reunite_mismatch_returns_halves() { sock_b.close_connection().await.unwrap(); } +#[tokio::test] +async fn test_socketeer_as_stream() { + let server_address = get_mock_address(echo_server).await; + let mut socketeer: Socketeer = Socketeer::connect(&format!("ws://{server_address}")) + .await + .unwrap(); + let message = EchoControlMessage::Message("direct stream".into()); + socketeer.send(message.clone()).await.unwrap(); + let item = socketeer.next().await.unwrap().unwrap(); + assert_eq!(item, message); +} + #[tokio::test] async fn test_binary_custom_keepalive() { // The widening of custom_keepalive_message from Option to From 9f0aaabf0988746b22b3b18eea74e722e902cc8c Mon Sep 17 00:00:00 2001 From: Zach Heylmun Date: Mon, 29 Jun 2026 15:32:09 -0400 Subject: [PATCH 7/7] test: cover concurrent cloned-Tx sends and ReuniteError Display/Error Replace the sequential test_split_cloned_sender_concurrent with test_split_concurrent_sends_from_clones, which spawns two independent tokio tasks each owning a cloned SocketeerTx and asserts that both echoed payloads are received (order-insensitive via HashSet). Also exercises ReuniteError's Display and std::error::Error impls before the destructuring move in test_reunite_mismatch_returns_halves. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/integration.rs | 40 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 36 insertions(+), 4 deletions(-) diff --git a/tests/integration.rs b/tests/integration.rs index 379dc79..364e56f 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -584,16 +584,45 @@ async fn test_split_round_trip() { } #[tokio::test] -async fn test_split_cloned_sender_concurrent() { +async fn test_split_concurrent_sends_from_clones() { + // Multiple cloned Tx handles send concurrently; all deliveries observed. + use std::collections::HashSet; + let server_address = get_mock_address(echo_server).await; let socketeer: Socketeer = Socketeer::connect(&format!("ws://{server_address}")) .await .unwrap(); let (tx, mut rx) = socketeer.split(); + let tx1 = tx.clone(); let tx2 = tx.clone(); - let message = EchoControlMessage::Message("from clone".into()); - tx2.send(message.clone()).await.unwrap(); - assert_eq!(rx.next_message().await.unwrap(), message); + + let payload1 = "concurrent-a".to_string(); + let payload2 = "concurrent-b".to_string(); + + let msg1 = EchoControlMessage::Message(payload1.clone()); + let msg2 = EchoControlMessage::Message(payload2.clone()); + + // Spawn two independent tasks, each owning a clone; SocketeerTx is Send + 'static. + let h1 = tokio::spawn(async move { tx1.send(msg1).await.unwrap() }); + let h2 = tokio::spawn(async move { tx2.send(msg2).await.unwrap() }); + + h1.await.unwrap(); + h2.await.unwrap(); + + // Drain exactly two echoed messages; order across senders is unspecified. + let r1 = rx.next_message().await.unwrap(); + let r2 = rx.next_message().await.unwrap(); + + let received: HashSet = [r1, r2] + .into_iter() + .map(|m| match m { + EchoControlMessage::Message(s) => s, + other => panic!("unexpected echo: {other:?}"), + }) + .collect(); + let expected: HashSet = [payload1, payload2].into_iter().collect(); + assert_eq!(received, expected); + tx.close().await.unwrap(); } @@ -703,6 +732,9 @@ async fn test_reunite_mismatch_returns_halves() { let err = rx_a .reunite(tx_b) .expect_err("mismatched halves must not reunite"); + // Exercise Display and Error trait impls before consuming err by value. + assert!(!err.to_string().is_empty()); + let _: &dyn std::error::Error = &err; let socketeer::ReuniteError { tx: tx_b, rx: rx_a } = err; // Both halves survive and reunite with their correct partners. let sock_a = rx_a.reunite(tx_a).expect("correct halves reunite");