starpc 0.52.1 → 0.52.2

package/srpc/testing.rs

@@ -13,22 +13,8 @@ use crate::error::{Error, Result};
 use crate::rpc::PacketWriter;
 use crate::transport::create_packet_channel;
 
-/// Creates a pair of connected in-memory streams.
-///
-/// Returns (client_stream, server_stream) that can be used for testing
-/// client-server communication.
-///
-/// # Arguments
-/// * `buffer_size` - Size of the internal buffer for each direction
-///
-/// # Example
-///
-/// ```ignore
-/// let (client_stream, server_stream) = create_pipe(64 * 1024);
-///
-/// // Use client_stream with a client
-/// // Use server_stream with a server
-/// ```
+/// Creates a pair of connected in-memory streams for testing client-server
+/// communication, with `buffer_size` bytes of buffer in each direction.
 pub fn create_pipe(buffer_size: usize) -> (DuplexStream, DuplexStream) {
     duplex(buffer_size)
 }
@@ -38,21 +24,8 @@ pub fn create_pipe_default() -> (DuplexStream, DuplexStream) {
     create_pipe(64 * 1024)
 }
 
-/// In-memory stream opener for testing.
-///
-/// This opener can hold multiple streams and returns them one at a time
-/// on each call to `open_stream`.
-///
-/// # Example
-///
-/// ```ignore
-/// let (client_stream, server_stream) = create_pipe_default();
-///
-/// let opener = InMemoryOpener::new(vec![client_stream]);
-/// let client = SrpcClient::new(opener);
-///
-/// // The client will use client_stream for its RPC
-/// ```
+/// In-memory stream opener for testing that holds multiple streams and
+/// returns them one at a time on each call to `open_stream`.
 pub struct InMemoryOpener {
     streams: Arc<Mutex<Vec<DuplexStream>>>,
 }
@@ -108,40 +81,15 @@ impl SingleInMemoryOpener {
 #[async_trait]
 impl OpenStream for SingleInMemoryOpener {
     async fn open_stream(&self) -> Result<(Arc<dyn PacketWriter>, PacketReceiver)> {
-        let stream = self
-            .stream
-            .lock()
-            .await
-            .take()
-            .ok_or(Error::StreamClosed)?;
+        let stream = self.stream.lock().await.take().ok_or(Error::StreamClosed)?;
 
         let (read_half, write_half) = tokio::io::split(stream);
         Ok(create_packet_channel(read_half, write_half))
     }
 }
 
-/// Creates a connected client-server test setup.
-///
-/// Returns (client_opener, server_stream) where:
-/// - `client_opener` can be used to create a `SrpcClient`
-/// - `server_stream` can be passed to `Server::handle_stream`
-///
-/// # Example
-///
-/// ```ignore
-/// let (opener, server_stream) = create_test_pair();
-///
-/// let client = SrpcClient::new(opener);
-/// let server = Server::new(mux);
-///
-/// // Spawn server
-/// tokio::spawn(async move {
-///     server.handle_stream(server_stream).await
-/// });
-///
-/// // Use client
-/// let response = client.exec_call("svc", "method", &request).await?;
-/// ```
+/// Creates a connected client-server test setup: an opener for a
+/// `SrpcClient` and the server stream for `Server::handle_stream`.
 pub fn create_test_pair() -> (SingleInMemoryOpener, DuplexStream) {
     let (client_stream, server_stream) = create_pipe_default();
     (SingleInMemoryOpener::new(client_stream), server_stream)
@@ -225,13 +173,10 @@ mod tests {
         stream.close_send().await.unwrap();
 
         // Read response with a timeout to handle potential races.
-        let response = tokio::time::timeout(
-            std::time::Duration::from_secs(5),
-            stream.recv_bytes(),
-        )
-        .await
-        .expect("timeout")
-        .expect("recv_bytes failed");
+        let response = tokio::time::timeout(std::time::Duration::from_secs(5), stream.recv_bytes())
+            .await
+            .expect("timeout")
+            .expect("recv_bytes failed");
         assert_eq!(&response[..], b"hello");
 
         // Wait for server to complete.

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