Adds Socketeer split API for concurrent send/receive

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,17 @@ 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::{ReuniteError, 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::pin::Pin;
 use std::sync::{Arc, Mutex, PoisonError};
+use std::task::{Context, Poll};
 
-use futures::StreamExt;
-use tokio::sync::{mpsc, oneshot};
+use futures::{Stream, StreamExt};
+use tokio::sync::mpsc;
 use tokio_tungstenite::{connect_async, tungstenite::Message};
 
 #[cfg(feature = "tracing")]
@@ -71,6 +75,26 @@ where
     }
 }
 
+impl<C: Codec + Unpin, Handler, const CHANNEL_SIZE: usize> Stream
+    for Socketeer<C, Handler, CHANNEL_SIZE>
+where
+    Handler: ConnectionHandler<C> + Unpin,
+{
+    type Item = Result<C::Rx, Error>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        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<C, const CHANNEL_SIZE: usize> Socketeer<C, NoopHandler, CHANNEL_SIZE>
 where
     C: Codec + Default,
@@ -165,6 +189,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<Message>,
+        sender: mpsc::Sender<TxChannelPayload>,
+        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`].
     ///
@@ -189,11 +238,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`].
@@ -233,18 +278,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::<Result<(), Error>>();
-        self.sender
-            .send(TxChannelPayload {
-                message,
-                response_tx: 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.
@@ -306,3 +340,34 @@ where
         }
     }
 }
+
+impl<C, Handler, const CHANNEL_SIZE: usize> Socketeer<C, Handler, CHANNEL_SIZE>
+where
+    C: Codec + Clone,
+    Handler: ConnectionHandler<C>,
+{
+    /// 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<C>, SocketeerRx<C, Handler, CHANNEL_SIZE>) {
+        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)
+    }
+}

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<Result<(), Error>>,
+    pub(crate) response_tx: Option<oneshot::Sender<Result<(), Error>>>,
 }
 
 pub(crate) type WebSocketStreamType = WebSocketStream<MaybeTlsStream<TcpStream>>;
@@ -51,6 +51,56 @@ enum LoopState {
 /// channels drop.
 pub(crate) type TerminalError = Arc<Mutex<Option<Error>>>;
 
+/// 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<Error> {
+    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)
+}
+
+/// Enqueue a frame and await the loop's result for this specific send.
+pub(crate) async fn send_confirmed(
+    sender: &mpsc::Sender<TxChannelPayload>,
+    terminal_error: &TerminalError,
+    message: Message,
+) -> Result<(), Error> {
+    let (tx, rx) = oneshot::channel::<Result<(), Error>>();
+    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"),
+    }
+}
+
+/// 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<TxChannelPayload>,
+    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<TxChannelPayload>) -> Result<(), Error> {
     let (tx, rx) = oneshot::channel::<Result<(), Error>>();
@@ -60,7 +110,7 @@ pub(crate) async fn send_close(sender: &mpsc::Sender<TxChannelPayload>) -> 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 +173,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")]