starpc 0.52.1 → 0.52.2

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())),
     }
 }

package/srpc/tests/integration_test.rs

@@ -1,495 +0,0 @@
-//! Integration tests for starpc, mirroring the Go/JS test patterns.
-//!
-//! These tests verify:
-//! 1. Unary RPC
-//! 2. Server streaming (5 messages)
-//! 3. Client streaming
-//! 4. Bidirectional streaming
-//! 5. Error handling
-//! 6. Wire format compatibility
-
-use async_trait::async_trait;
-use prost::Message;
-use std::sync::Arc;
-use std::time::Duration;
-
-use starpc::error::{Error, Result};
-use starpc::handler::Handler;
-use starpc::invoker::Invoker;
-use starpc::mux::Mux;
-use starpc::server::Server;
-use starpc::stream::{Stream, StreamExt};
-use starpc::testing::{create_test_pair, SingleInMemoryOpener};
-use starpc::Client;
-
-// Simple EchoMsg for testing
-#[derive(Clone, PartialEq, Message)]
-struct EchoMsg {
-    #[prost(string, tag = "1")]
-    body: String,
-}
-
-const BODY_TXT: &str = "hello world via starpc e2e test";
-
-/// Echo server implementation matching Go's echo/server.go
-struct EchoServer;
-
-#[async_trait]
-impl Invoker for EchoServer {
-    async fn invoke_method(
-        &self,
-        _service_id: &str,
-        method_id: &str,
-        stream: Box<dyn Stream>,
-    ) -> (bool, Result<()>) {
-        match method_id {
-            "Echo" => (true, self.echo(stream).await),
-            "EchoServerStream" => (true, self.echo_server_stream(stream).await),
-            "EchoClientStream" => (true, self.echo_client_stream(stream).await),
-            "EchoBidiStream" => (true, self.echo_bidi_stream(stream).await),
-            _ => (false, Err(Error::Unimplemented)),
-        }
-    }
-}
-
-impl Handler for EchoServer {
-    fn service_id(&self) -> &'static str {
-        "echo.Echoer"
-    }
-
-    fn method_ids(&self) -> &'static [&'static str] {
-        &[
-            "Echo",
-            "EchoServerStream",
-            "EchoClientStream",
-            "EchoBidiStream",
-        ]
-    }
-}
-
-impl EchoServer {
-    /// Unary echo - returns the same message
-    async fn echo(&self, stream: Box<dyn Stream>) -> Result<()> {
-        let msg: EchoMsg = stream.msg_recv().await?;
-        stream.msg_send(&msg).await?;
-        Ok(())
-    }
-
-    /// Server streaming - sends 5 copies of the message
-    async fn echo_server_stream(&self, stream: Box<dyn Stream>) -> Result<()> {
-        let msg: EchoMsg = stream.msg_recv().await?;
-
-        // Send 5 responses with delays
-        for _ in 0..5 {
-            if stream.context().is_cancelled() {
-                return Err(Error::Cancelled);
-            }
-            stream.msg_send(&msg).await?;
-            tokio::time::sleep(Duration::from_millis(10)).await;
-        }
-        Ok(())
-    }
-
-    /// Client streaming - returns first message received
-    async fn echo_client_stream(&self, stream: Box<dyn Stream>) -> Result<()> {
-        let msg: EchoMsg = stream.msg_recv().await?;
-        stream.msg_send(&msg).await?;
-        Ok(())
-    }
-
-    /// Bidirectional streaming - server sends first, then echoes all messages
-    async fn echo_bidi_stream(&self, stream: Box<dyn Stream>) -> Result<()> {
-        // Server sends initial message
-        let initial = EchoMsg {
-            body: "hello from server".to_string(),
-        };
-        stream.msg_send(&initial).await?;
-
-        // Echo all received messages
-        loop {
-            match stream.msg_recv::<EchoMsg>().await {
-                Ok(msg) => {
-                    if msg.body.is_empty() {
-                        return Err(Error::Remote("got message with empty body".to_string()));
-                    }
-                    stream.msg_send(&msg).await?;
-                }
-                Err(Error::StreamClosed) => break,
-                Err(e) => return Err(e),
-            }
-        }
-        Ok(())
-    }
-}
-
-/// Test infrastructure: creates connected client and server
-async fn setup_e2e() -> (starpc::SrpcClient<SingleInMemoryOpener>, tokio::task::JoinHandle<()>) {
-    let (opener, server_stream) = create_test_pair();
-
-    // Set up the server
-    let mux = Arc::new(Mux::new());
-    mux.register(Arc::new(EchoServer)).unwrap();
-    let server = Server::with_arc(mux);
-
-    // Spawn server handler
-    let server_handle = tokio::spawn(async move {
-        let _ = server.handle_stream(server_stream).await;
-    });
-
-    // Create client
-    let client = starpc::SrpcClient::new(opener);
-
-    (client, server_handle)
-}
-
-// ============================================================================
-// Tests matching Go's server_test.go
-// ============================================================================
-
-#[tokio::test]
-async fn test_e2e_unary() {
-    let (client, server_handle) = setup_e2e().await;
-
-    // Make unary call
-    let request = EchoMsg {
-        body: BODY_TXT.to_string(),
-    };
-    let response: EchoMsg = client
-        .exec_call("echo.Echoer", "Echo", &request)
-        .await
-        .expect("exec_call failed");
-
-    assert_eq!(response.body, BODY_TXT);
-
-    server_handle.abort();
-}
-
-#[tokio::test]
-async fn test_e2e_server_stream() {
-    let (client, server_handle) = setup_e2e().await;
-
-    // Send request and open stream
-    let request = EchoMsg {
-        body: BODY_TXT.to_string(),
-    };
-    let data = request.encode_to_vec();
-    let stream = client
-        .new_stream("echo.Echoer", "EchoServerStream", Some(&data))
-        .await
-        .expect("new_stream failed");
-
-    // Close send side
-    stream.close_send().await.expect("close_send failed");
-
-    // Expect to receive 5 messages
-    let expected_rx = 5;
-    let mut received = 0;
-
-    loop {
-        match stream.msg_recv::<EchoMsg>().await {
-            Ok(msg) => {
-                assert_eq!(msg.body, BODY_TXT);
-                received += 1;
-            }
-            Err(Error::StreamClosed) => break,
-            Err(e) => panic!("unexpected error: {}", e),
-        }
-    }
-
-    assert_eq!(
-        received, expected_rx,
-        "expected {} messages, got {}",
-        expected_rx, received
-    );
-
-    server_handle.abort();
-}
-
-#[tokio::test]
-async fn test_e2e_client_stream() {
-    let (client, server_handle) = setup_e2e().await;
-
-    // Open stream without initial message
-    let stream = client
-        .new_stream("echo.Echoer", "EchoClientStream", None)
-        .await
-        .expect("new_stream failed");
-
-    // Send a message
-    let request = EchoMsg {
-        body: BODY_TXT.to_string(),
-    };
-    stream.msg_send(&request).await.expect("msg_send failed");
-
-    // Close send side
-    stream.close_send().await.expect("close_send failed");
-
-    // Receive response
-    let response: EchoMsg = stream.msg_recv().await.expect("msg_recv failed");
-    assert_eq!(response.body, BODY_TXT);
-
-    stream.close().await.ok();
-    server_handle.abort();
-}
-
-#[tokio::test]
-async fn test_e2e_bidi_stream() {
-    let (client, server_handle) = setup_e2e().await;
-
-    // Open bidirectional stream
-    let stream = client
-        .new_stream("echo.Echoer", "EchoBidiStream", None)
-        .await
-        .expect("new_stream failed");
-
-    // Receive server's initial message
-    let initial: EchoMsg = stream.msg_recv().await.expect("msg_recv failed");
-    assert_eq!(initial.body, "hello from server");
-
-    // Send a message from client
-    let client_msg = EchoMsg {
-        body: "hello from client".to_string(),
-    };
-    stream.msg_send(&client_msg).await.expect("msg_send failed");
-
-    // Receive echoed message
-    let echoed: EchoMsg = stream.msg_recv().await.expect("msg_recv failed");
-    assert_eq!(echoed.body, "hello from client");
-
-    // Close the stream
-    stream.close().await.expect("close failed");
-    server_handle.abort();
-}
-
-#[tokio::test]
-async fn test_e2e_multiple_bidi_messages() {
-    let (client, server_handle) = setup_e2e().await;
-
-    let stream = client
-        .new_stream("echo.Echoer", "EchoBidiStream", None)
-        .await
-        .expect("new_stream failed");
-
-    // Receive server's initial message
-    let _: EchoMsg = stream.msg_recv().await.expect("initial recv failed");
-
-    // Send and receive multiple messages
-    for i in 0..10 {
-        let msg = EchoMsg {
-            body: format!("message {}", i),
-        };
-        stream.msg_send(&msg).await.expect("msg_send failed");
-
-        let echoed: EchoMsg = stream.msg_recv().await.expect("msg_recv failed");
-        assert_eq!(echoed.body, format!("message {}", i));
-    }
-
-    stream.close().await.expect("close failed");
-    server_handle.abort();
-}
-
-#[tokio::test]
-async fn test_e2e_unary_empty_message() {
-    let (client, server_handle) = setup_e2e().await;
-
-    // Send empty message
-    let request = EchoMsg {
-        body: String::new(),
-    };
-    let response: EchoMsg = client
-        .exec_call("echo.Echoer", "Echo", &request)
-        .await
-        .expect("exec_call failed");
-
-    assert_eq!(response.body, "");
-
-    server_handle.abort();
-}
-
-#[tokio::test]
-async fn test_e2e_unimplemented_method() {
-    let (client, server_handle) = setup_e2e().await;
-
-    let request = EchoMsg {
-        body: "test".to_string(),
-    };
-    let result: Result<EchoMsg> = client
-        .exec_call("echo.Echoer", "NonExistentMethod", &request)
-        .await;
-
-    assert!(result.is_err());
-
-    server_handle.abort();
-}
-
-#[tokio::test]
-async fn test_codec_wire_format() {
-    use starpc::codec::PacketCodec;
-    use starpc::proto::{packet::Body, CallData, CallStart, Packet};
-    use tokio_util::codec::{Decoder, Encoder};
-
-    let mut codec = PacketCodec::new();
-    let mut buf = bytes::BytesMut::new();
-
-    // Test CallStart encoding
-    let call_start = Packet {
-        body: Some(Body::CallStart(CallStart {
-            rpc_service: "test.Service".into(),
-            rpc_method: "TestMethod".into(),
-            data: vec![1, 2, 3, 4],
-            data_is_zero: false,
-        })),
-    };
-
-    codec
-        .encode(call_start.clone(), &mut buf)
-        .expect("encode failed");
-
-    // Verify length prefix (little-endian u32)
-    let len = u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]) as usize;
-    assert_eq!(len, buf.len() - 4);
-
-    // Decode and verify
-    let decoded = codec
-        .decode(&mut buf)
-        .expect("decode failed")
-        .expect("no packet");
-    assert_eq!(decoded, call_start);
-
-    // Test CallData encoding
-    buf.clear();
-    let call_data = Packet {
-        body: Some(Body::CallData(CallData {
-            data: vec![5, 6, 7, 8],
-            data_is_zero: false,
-            complete: true,
-            error: String::new(),
-        })),
-    };
-
-    codec
-        .encode(call_data.clone(), &mut buf)
-        .expect("encode failed");
-    let decoded = codec
-        .decode(&mut buf)
-        .expect("decode failed")
-        .expect("no packet");
-    assert_eq!(decoded, call_data);
-
-    // Test empty data with data_is_zero flag
-    buf.clear();
-    let empty_data = Packet {
-        body: Some(Body::CallData(CallData {
-            data: vec![],
-            data_is_zero: true,
-            complete: false,
-            error: String::new(),
-        })),
-    };
-
-    codec
-        .encode(empty_data.clone(), &mut buf)
-        .expect("encode failed");
-    let decoded = codec
-        .decode(&mut buf)
-        .expect("decode failed")
-        .expect("no packet");
-    assert_eq!(decoded, empty_data);
-}
-
-#[tokio::test]
-async fn test_packet_validation() {
-    use starpc::packet::Validate;
-    use starpc::proto::{packet::Body, CallData, CallStart, Packet};
-
-    // Valid CallStart
-    let valid_start = Packet {
-        body: Some(Body::CallStart(CallStart {
-            rpc_service: "svc".into(),
-            rpc_method: "method".into(),
-            data: vec![],
-            data_is_zero: false,
-        })),
-    };
-    assert!(valid_start.validate().is_ok());
-
-    // Invalid CallStart - empty method
-    let invalid_start = Packet {
-        body: Some(Body::CallStart(CallStart {
-            rpc_service: "svc".into(),
-            rpc_method: String::new(),
-            data: vec![],
-            data_is_zero: false,
-        })),
-    };
-    assert!(invalid_start.validate().is_err());
-
-    // Valid CallData with data
-    let valid_data = Packet {
-        body: Some(Body::CallData(CallData {
-            data: vec![1, 2, 3],
-            data_is_zero: false,
-            complete: false,
-            error: String::new(),
-        })),
-    };
-    assert!(valid_data.validate().is_ok());
-
-    // Invalid CallData - empty everything
-    let invalid_data = Packet {
-        body: Some(Body::CallData(CallData {
-            data: vec![],
-            data_is_zero: false,
-            complete: false,
-            error: String::new(),
-        })),
-    };
-    assert!(invalid_data.validate().is_err());
-
-    // Empty packet
-    let empty_packet = Packet { body: None };
-    assert!(empty_packet.validate().is_err());
-}
-
-#[tokio::test]
-async fn test_mux_registration_and_lookup() {
-    let mux = Mux::new();
-
-    // Register handler
-    mux.register(Arc::new(EchoServer)).unwrap();
-
-    // Check service exists
-    assert!(mux.has_service("echo.Echoer"));
-    assert!(!mux.has_service("nonexistent"));
-
-    // Check methods exist
-    assert!(mux.has_service_method("echo.Echoer", "Echo"));
-    assert!(mux.has_service_method("echo.Echoer", "EchoServerStream"));
-    assert!(!mux.has_service_method("echo.Echoer", "NonExistent"));
-
-    // Empty strings should return false
-    assert!(!mux.has_service(""));
-    assert!(!mux.has_service_method("", "Echo"));
-    assert!(!mux.has_service_method("echo.Echoer", ""));
-}
-
-#[tokio::test]
-async fn test_error_types() {
-    use starpc::error::codes;
-
-    // Test error predicates
-    assert!(Error::Aborted.is_abort());
-    assert!(Error::Cancelled.is_abort());
-    assert!(!Error::StreamClosed.is_abort());
-
-    assert!(Error::StreamClosed.is_closed());
-    assert!(Error::Cancelled.is_closed());
-
-    assert!(Error::StreamIdle.is_timeout());
-
-    assert!(Error::Unimplemented.is_unimplemented());
-
-    // Test error codes
-    assert_eq!(codes::ERR_RPC_ABORT, "ERR_RPC_ABORT");
-    assert_eq!(codes::ERR_STREAM_IDLE, "ERR_STREAM_IDLE");
-}