starpc 0.52.0 → 0.52.2

package/srpc/transport.rs

@@ -12,6 +12,7 @@ use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tokio::sync::Mutex;
 use tokio_util::codec::{Encoder, FramedRead};
 
+use crate::client::PacketReceiver;
 use crate::codec::PacketCodec;
 use crate::error::{Error, Result};
 use crate::proto::Packet;
@@ -68,65 +69,13 @@ impl<W: AsyncWrite + Send + Unpin + 'static> PacketWriter for TransportPacketWri
     }
 }
 
-/// Receiver for incoming packets from a transport.
-pub type PacketReceiver = tokio::sync::mpsc::Receiver<Packet>;
+/// Default bounded request capacity for multiplexed transport admission.
+#[cfg(feature = "yamux")]
+pub(crate) const DEFAULT_CHANNEL_BUFFER: usize = 32;
 
-/// Sender for incoming packets to be processed.
-pub type PacketSender = tokio::sync::mpsc::Sender<Packet>;
-
-/// Default channel buffer size for packet channels.
-pub const DEFAULT_CHANNEL_BUFFER: usize = 32;
-
-/// Spawns a task that reads packets from a transport and sends them through a channel.
-///
-/// This is a common pattern used by both client and server to handle incoming packets.
-/// The task will run until the transport is closed or an error occurs.
-///
-/// # Arguments
-/// * `reader` - The async reader to read packets from
-/// * `sender` - The channel sender to forward packets to
-///
-/// # Returns
-/// A `JoinHandle` for the spawned task.
-pub fn spawn_packet_reader<R>(
-    reader: R,
-    sender: PacketSender,
-) -> tokio::task::JoinHandle<()>
-where
-    R: AsyncRead + Send + Unpin + 'static,
-{
-    tokio::spawn(async move {
-        let mut framed = FramedRead::new(reader, PacketCodec::new());
-        while let Some(result) = framed.next().await {
-            match result {
-                Ok(packet) => {
-                    if sender.send(packet).await.is_err() {
-                        // Receiver dropped, stop reading
-                        break;
-                    }
-                }
-                Err(_) => {
-                    // Read error, stop reading
-                    break;
-                }
-            }
-        }
-    })
-}
-
-/// Creates a packet writer and receiver from a split transport.
-///
-/// This is a convenience function that:
-/// 1. Creates a `TransportPacketWriter` from the write half
-/// 2. Spawns a packet reader task for the read half
-/// 3. Returns the writer and receiver channel
-///
-/// # Arguments
-/// * `read_half` - The read half of the transport
-/// * `write_half` - The write half of the transport
-///
-/// # Returns
-/// A tuple of (packet writer, packet receiver).
+/// Creates a packet writer and an owned incoming packet stream from a split transport.
+/// The receiver reads on demand; dropping it immediately releases the read half.
+/// No background task or additional packet queue is created.
 pub fn create_packet_channel<R, W>(
     read_half: R,
     write_half: W,
@@ -135,10 +84,18 @@ where
     R: AsyncRead + Send + Unpin + 'static,
     W: AsyncWrite + Send + Unpin + 'static,
 {
+    // Framing belongs to the returned receive capability for its complete lifetime.
+    let reader = FramedRead::new(read_half, PacketCodec::new());
+    let packets = futures::stream::unfold(reader, |mut reader| async move {
+        match reader.next().await {
+            Some(Ok(packet)) => Some((packet, reader)),
+            _ => None,
+        }
+    });
+
+    // The peer-facing writer retains its own half of this same transport.
     let writer: Arc<dyn PacketWriter> = Arc::new(TransportPacketWriter::new(write_half));
-    let (tx, rx) = tokio::sync::mpsc::channel(DEFAULT_CHANNEL_BUFFER);
-    spawn_packet_reader(read_half, tx);
-    (writer, rx)
+    (writer, Box::pin(packets))
 }
 
 /// Encodes optional data for protobuf messages.
@@ -147,9 +104,6 @@ where
 /// - `None` -> empty data, `data_is_zero = false`
 /// - `Some(empty)` -> empty data, `data_is_zero = true`
 /// - `Some(data)` -> data bytes, `data_is_zero = false`
-///
-/// # Returns
-/// A tuple of (data bytes, data_is_zero flag).
 pub fn encode_optional_data(data: Option<Bytes>) -> (Vec<u8>, bool) {
     match data {
         Some(d) if d.is_empty() => (vec![], true),
@@ -158,13 +112,9 @@ pub fn encode_optional_data(data: Option<Bytes>) -> (Vec<u8>, bool) {
     }
 }
 
-/// Decodes optional data from protobuf messages.
-///
-/// Inverse of `encode_optional_data`.
-///
-/// # Returns
-/// `Some(Bytes)` if data was present (including empty data with `data_is_zero`),
-/// `None` if no data was included.
+/// Decodes optional data from protobuf messages, inverting
+/// `encode_optional_data`: `Some(Bytes)` when data was present, including
+/// empty data with `data_is_zero` set.
 pub fn decode_optional_data(data: Vec<u8>, data_is_zero: bool) -> Option<Bytes> {
     if !data.is_empty() || data_is_zero {
         Some(Bytes::from(data))

package/srpc/websocket.rs

@@ -1,70 +1,119 @@
-//! WebSocket transport adapters for starpc.
+//! Owned WebSocket byte transport for StarPC.
 
-use bytes::Bytes;
-use futures::{SinkExt, StreamExt};
-use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, DuplexStream};
-use tokio_tungstenite::{tungstenite::Message, WebSocketStream};
+use bytes::{Buf, Bytes};
+use futures::{Sink, Stream};
+use std::io;
+use std::pin::Pin;
+use std::task::{ready, Context, Poll};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+use tokio_tungstenite::{tungstenite, tungstenite::Message, WebSocketStream};
 
-/// WEBSOCKET_BYTE_STREAM_BUFFER is the duplex buffer used by WebSocket byte adapters.
-pub const WEBSOCKET_BYTE_STREAM_BUFFER: usize = 64 * 1024;
-
-/// WEBSOCKET_MESSAGE_BUFFER is the maximum byte chunk sent in one WebSocket message.
+/// Maximum byte chunk sent in one WebSocket message.
 pub const WEBSOCKET_MESSAGE_BUFFER: usize = 16 * 1024;
 
-/// websocket_byte_stream adapts a binary WebSocket stream to Tokio byte I/O.
-///
-/// The adapter treats incoming binary WebSocket messages as consecutive bytes
-/// and writes outbound byte chunks as binary WebSocket messages. Control and
-/// text messages are ignored except that close terminates the byte stream.
-pub fn websocket_byte_stream<S>(socket: WebSocketStream<S>) -> DuplexStream
-where
-    S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
-{
-    let (mut socket_writer, mut socket_reader) = socket.split();
-    let (stream, peer) = tokio::io::duplex(WEBSOCKET_BYTE_STREAM_BUFFER);
-    let (mut peer_reader, mut peer_writer) = tokio::io::split(peer);
+/// Adapts binary WebSocket messages to byte I/O without background tasks.
+/// Incoming binary messages form one byte stream. Text and control messages
+/// carry no application bytes; a close message ends input. Dropping this
+/// capability drops its socket, including when input is idle.
+pub fn websocket_byte_stream<S>(socket: WebSocketStream<S>) -> WebSocketByteStream<S> {
+    WebSocketByteStream {
+        socket,
+        buffered: Bytes::new(),
+        input: Input::Open,
+    }
+}
 
-    tokio::spawn(async move {
-        while let Some(result) = socket_reader.next().await {
-            match result {
-                Ok(message) if message.is_binary() => {
-                    let data = message.into_data();
-                    if peer_writer.write_all(data.as_ref()).await.is_err() {
-                        break;
-                    }
-                    if peer_writer.flush().await.is_err() {
-                        break;
-                    }
-                }
-                Ok(message) if message.is_close() => break,
-                Ok(_) => {}
-                Err(_) => break,
-            }
+/// Byte I/O that owns its WebSocket and the unread part of one incoming message.
+pub struct WebSocketByteStream<S> {
+    socket: WebSocketStream<S>,
+    buffered: Bytes,
+    input: Input,
+}
+
+enum Input {
+    Open,
+    Closing,
+    Ended,
+}
+
+impl<S: AsyncRead + AsyncWrite + Unpin> AsyncRead for WebSocketByteStream<S> {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        output: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        let this = self.get_mut();
+        if output.remaining() == 0 {
+            return Poll::Ready(Ok(()));
         }
-        let _ = peer_writer.shutdown().await;
-    });
 
-    tokio::spawn(async move {
-        let mut buf = vec![0; WEBSOCKET_MESSAGE_BUFFER];
+        // Retain each message until all its bytes have reached the reader.
         loop {
-            match peer_reader.read(&mut buf).await {
-                Ok(0) => {
-                    let _ = socket_writer.close().await;
-                    break;
-                }
-                Ok(n) => {
-                    let message = Message::Binary(Bytes::copy_from_slice(&buf[..n]));
-                    if socket_writer.send(message).await.is_err() {
-                        break;
-                    }
-                }
-                Err(_) => {
-                    let _ = socket_writer.close().await;
-                    break;
+            if !this.buffered.is_empty() {
+                let count = output.remaining().min(this.buffered.len());
+                output.put_slice(&this.buffered[..count]);
+                this.buffered.advance(count);
+                return Poll::Ready(Ok(()));
+            }
+            match this.input {
+                Input::Closing => {
+                    // Flush Tungstenite's close acknowledgment before reporting EOF.
+                    ready!(Pin::new(&mut this.socket).poll_flush(cx)).map_err(socket_error)?;
+                    this.input = Input::Ended;
+                    return Poll::Ready(Ok(()));
                 }
+                Input::Ended => return Poll::Ready(Ok(())),
+                Input::Open => {}
+            }
+
+            // Tungstenite owns WebSocket control handling and pending control writes.
+            match ready!(Pin::new(&mut this.socket).poll_next(cx)) {
+                Some(Ok(Message::Binary(data))) => this.buffered = data,
+                Some(Ok(Message::Close(_))) => this.input = Input::Closing,
+                None => this.input = Input::Ended,
+                Some(Ok(_)) => {}
+                Some(Err(error)) => return Poll::Ready(Err(socket_error(error))),
             }
         }
-    });
+    }
+}
+
+impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for WebSocketByteStream<S> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        input: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        if input.is_empty() {
+            return Poll::Ready(Ok(0));
+        }
+
+        // Respect socket backpressure before accepting another bounded message.
+        let socket = &mut self.get_mut().socket;
+        ready!(Pin::new(&mut *socket).poll_ready(cx)).map_err(socket_error)?;
+        let count = input.len().min(WEBSOCKET_MESSAGE_BUFFER);
+        Pin::new(socket)
+            .start_send(Message::Binary(Bytes::copy_from_slice(&input[..count])))
+            .map_err(socket_error)?;
+        Poll::Ready(Ok(count))
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.get_mut().socket)
+            .poll_flush(cx)
+            .map_err(socket_error)
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.get_mut().socket)
+            .poll_close(cx)
+            .map_err(socket_error)
+    }
+}
 
-    stream
+fn socket_error(error: tungstenite::Error) -> io::Error {
+    match error {
+        tungstenite::Error::Io(error) => error,
+        other => io::Error::other(other),
+    }
 }

package/srpc/writer.hpp

@@ -1,60 +1,74 @@
 #pragma once
 
+#include "errors.hpp"
+
 #include <functional>
 #include <memory>
 
-#include "errors.hpp"
-
 namespace srpc {
 class Packet;
 }
 
 namespace starpc {
 
-// PacketWriter is the interface used to write messages to a PacketStream.
-// Matches Go interface in writer.go
+/*
+ * PacketWriter writes one packet to the transport. Close may run from a
+ * receive callback or concurrently with a write, and must interrupt blocked
+ * I/O without joining callbacks. Destruction joins any transport receive
+ * threads.
+ */
 class PacketWriter {
 public:
-  virtual ~PacketWriter() = default;
+	virtual ~PacketWriter() = default;
 
-  // WritePacket writes a packet to the remote.
-  virtual Error WritePacket(const srpc::Packet &pkt) = 0;
+	/* WritePacket serializes pkt onto the transport. */
+	virtual Error WritePacket(const srpc::Packet &pkt) = 0;
 
-  // Close closes the writer.
-  virtual Error Close() = 0;
+	/* Close interrupts the transport. Repeated calls are safe. */
+	virtual Error Close() = 0;
 };
 
-// PacketWriterWithClose wraps a PacketWriter with an additional close function.
-// Matches packetWriterWithClose in writer.go
+/*
+ * PacketWriterWithClose forwards writes to an owned writer and runs an extra
+ * callback on the first Close. Matches packetWriterWithClose in writer.go.
+ */
 class PacketWriterWithClose : public PacketWriter {
 public:
-  PacketWriterWithClose(std::unique_ptr<PacketWriter> inner,
-                        std::function<Error()> close_fn)
-      : inner_(std::move(inner)), close_fn_(std::move(close_fn)) {}
-
-  Error WritePacket(const srpc::Packet &pkt) override {
-    return inner_->WritePacket(pkt);
-  }
-
-  Error Close() override {
-    Error err = inner_->Close();
-    Error err2 = close_fn_();
-    if (err != Error::OK)
-      return err;
-    return err2;
-  }
+	/* close_fn runs once, after inner->Close. */
+	PacketWriterWithClose(std::unique_ptr<PacketWriter> inner,
+			      std::function<Error()> close_fn)
+		: inner(std::move(inner)),
+		  close_fn(std::move(close_fn))
+	{
+	}
+
+	Error WritePacket(const srpc::Packet &pkt) override
+	{
+		return inner->WritePacket(pkt);
+	}
+
+	Error Close() override
+	{
+		/* Report the inner close even when the callback also fails. */
+		Error err = inner->Close();
+		Error err2 = close_fn();
+		if (err != Error::OK)
+			return err;
+		return err2;
+	}
 
 private:
-  std::unique_ptr<PacketWriter> inner_;
-  std::function<Error()> close_fn_;
+	std::unique_ptr<PacketWriter> inner;
+	std::function<Error()> close_fn;
 };
 
-// NewPacketWriterWithClose wraps a PacketWriter with a close function.
+/* NewPacketWriterWithClose wraps prw so close_fn runs on the first Close. */
 inline std::unique_ptr<PacketWriter>
 NewPacketWriterWithClose(std::unique_ptr<PacketWriter> prw,
-                         std::function<Error()> close_fn) {
-  return std::make_unique<PacketWriterWithClose>(std::move(prw),
-                                                 std::move(close_fn));
+			 std::function<Error()> close_fn)
+{
+	return std::make_unique<PacketWriterWithClose>(std::move(prw),
+						       std::move(close_fn));
 }
 
 } // namespace starpc

package/srpc/yamux.rs

@@ -1,13 +1,15 @@
 //! Yamux transport adapters for starpc.
 
 use async_trait::async_trait;
-use futures::future::poll_fn;
+use futures::{future::poll_fn, stream::FuturesUnordered, StreamExt};
 use std::io;
 use std::pin::Pin;
 use std::sync::Arc;
 use tokio::io::{AsyncRead, AsyncWrite};
-use tokio::sync::{mpsc, oneshot};
+use tokio::sync::{mpsc, oneshot, Mutex};
+use tokio::task::JoinHandle;
 use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
+use tokio_util::sync::CancellationToken;
 
 use crate::client::{OpenStream, PacketReceiver};
 use crate::error::{Error, Result};
@@ -19,14 +21,30 @@ use crate::transport::{create_packet_channel, DEFAULT_CHANNEL_BUFFER};
 type OpenResult = std::result::Result<::yamux::Stream, ::yamux::ConnectionError>;
 type OpenRequest = oneshot::Sender<OpenResult>;
 
-/// YamuxStreamOpener opens Starpc packet streams over one yamux connection.
+/// Opens StarPC packet streams over one owned yamux connection.
+/// Clones share the connection. Dropping the last opener cancels its driver;
+/// explicit close cancels and joins the driver before returning.
 #[derive(Clone)]
 pub struct YamuxStreamOpener {
+    driver: Arc<ClientDriver>,
+}
+
+struct ClientDriver {
     requests: mpsc::Sender<OpenRequest>,
+    task: Mutex<Option<JoinHandle<()>>>,
+    closed: CancellationToken,
+}
+
+impl Drop for ClientDriver {
+    fn drop(&mut self) {
+        if let Some(task) = self.task.get_mut().take() {
+            task.abort();
+        }
+    }
 }
 
 impl YamuxStreamOpener {
-    /// client creates a client-mode yamux opener over a Tokio transport.
+    /// Creates a client-mode yamux opener over a Tokio transport.
     pub fn client<T>(transport: T) -> Self
     where
         T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
@@ -34,17 +52,41 @@ impl YamuxStreamOpener {
         Self::client_with_config(transport, ::yamux::Config::default())
     }
 
-    /// client_with_config creates a client-mode yamux opener with a custom config.
+    /// Creates a client-mode yamux opener with a custom configuration.
     pub fn client_with_config<T>(transport: T, config: ::yamux::Config) -> Self
     where
         T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
     {
         let (requests, request_rx) = mpsc::channel(DEFAULT_CHANNEL_BUFFER);
-        spawn_client_driver(transport.compat(), config, request_rx);
-        Self { requests }
+        let closed = CancellationToken::new();
+        let task = spawn_client_driver(transport.compat(), config, request_rx, closed.clone());
+        Self {
+            driver: Arc::new(ClientDriver {
+                requests,
+                task: Mutex::new(Some(task)),
+                closed,
+            }),
+        }
     }
 
-    /// client_websocket creates a client-mode yamux opener over a WebSocket.
+    /// Cancels the shared connection and joins its driver.
+    /// Concurrent callers wait for the same completed teardown.
+    pub async fn close(&self) {
+        let mut task = self.driver.task.lock().await;
+        if let Some(task) = task.as_mut() {
+            task.abort();
+            let _ = task.await;
+        }
+        task.take();
+    }
+
+    /// Waits for the connection driver to end, including peer closure and protocol failure.
+    /// Observers do not prevent a concurrent explicit close from joining the driver.
+    pub async fn closed(&self) {
+        self.driver.closed.cancelled().await;
+    }
+
+    /// Creates a client-mode yamux opener over a WebSocket.
     #[cfg(feature = "websocket")]
     pub fn client_websocket<S>(socket: tokio_tungstenite::WebSocketStream<S>) -> Self
     where
@@ -53,7 +95,7 @@ impl YamuxStreamOpener {
         Self::client(crate::websocket::websocket_byte_stream(socket))
     }
 
-    /// client_websocket_with_config creates a client-mode yamux WebSocket opener with a custom config.
+    /// Creates a client-mode yamux WebSocket opener with a custom configuration.
     #[cfg(feature = "websocket")]
     pub fn client_websocket_with_config<S>(
         socket: tokio_tungstenite::WebSocketStream<S>,
@@ -70,7 +112,8 @@ impl YamuxStreamOpener {
 impl OpenStream for YamuxStreamOpener {
     async fn open_stream(&self) -> Result<(Arc<dyn PacketWriter>, PacketReceiver)> {
         let (tx, rx) = oneshot::channel();
-        self.requests
+        self.driver
+            .requests
             .send(tx)
             .await
             .map_err(|_| Error::StreamClosed)?;
@@ -85,7 +128,7 @@ impl OpenStream for YamuxStreamOpener {
     }
 }
 
-/// handle_server_connection accepts yamux streams and serves each as Starpc.
+/// Serves yamux substreams within this connection's owned lifetime.
 pub async fn handle_server_connection<I, T>(
     server: &Server<I>,
     transport: T,
@@ -98,18 +141,23 @@ where
     let mut connection =
         ::yamux::Connection::new(transport.compat(), config, ::yamux::Mode::Server);
 
+    // Keep accepted calls within the connection's lifetime. Losing the connection
+    // or dropping this future drops every call and its cancellation guard.
+    let mut calls = FuturesUnordered::new();
     loop {
-        match poll_fn(|cx| connection.poll_next_inbound(cx)).await {
-            Some(Ok(stream)) => {
-                let server = server.clone_for_spawn();
-                tokio::spawn(async move {
-                    if let Err(err) = server.handle_stream(stream.compat()).await {
-                        server.report_error(err);
-                    }
-                });
+        tokio::select! {
+            incoming = poll_fn(|cx| connection.poll_next_inbound(cx)) => {
+                match incoming {
+                    Some(Ok(stream)) => calls.push(server.handle_stream(stream.compat())),
+                    Some(Err(err)) => return Err(connection_error(err)),
+                    None => return Ok(()),
+                }
+            }
+            Some(result) = calls.next(), if !calls.is_empty() => {
+                if let Err(err) = result {
+                    server.report_error(err);
+                }
             }
-            Some(Err(err)) => return Err(connection_error(err)),
-            None => return Ok(()),
         }
     }
 }
@@ -126,21 +174,29 @@ fn spawn_client_driver<T>(
     transport: T,
     config: ::yamux::Config,
     mut requests: mpsc::Receiver<OpenRequest>,
-) -> tokio::task::JoinHandle<()>
+    closed: CancellationToken,
+) -> JoinHandle<()>
 where
     T: futures::io::AsyncRead + futures::io::AsyncWrite + Send + Unpin + 'static,
 {
+    // Construct the guard before spawning so cancellation before the first poll also notifies.
+    let ended = closed.drop_guard();
     tokio::spawn(async move {
+        let _ended = ended;
         let mut connection = ::yamux::Connection::new(transport, config, ::yamux::Mode::Client);
         let mut pending = None;
-        let mut requests_closed = false;
 
         loop {
             let event = poll_fn(|cx| {
-                if pending.is_none() && !requests_closed {
+                if pending.is_none() {
                     match Pin::new(&mut requests).poll_recv(cx) {
                         std::task::Poll::Ready(Some(request)) => pending = Some(request),
-                        std::task::Poll::Ready(None) => requests_closed = true,
+                        std::task::Poll::Ready(None) => {
+                            return std::task::Poll::Ready(DriverEvent::Closed {
+                                pending: None,
+                                err: None,
+                            });
+                        }
                         std::task::Poll::Pending => {}
                     }
                 }
@@ -193,6 +249,6 @@ where
 fn connection_error(err: ::yamux::ConnectionError) -> Error {
     match err {
         ::yamux::ConnectionError::Io(err) => Error::Io(err),
-        other => Error::Io(io::Error::new(io::ErrorKind::Other, other.to_string())),
+        other => Error::Io(io::Error::other(other.to_string())),
     }
 }