starpc 0.52.1 → 0.52.2

package/srpc/client.rs

@@ -6,6 +6,7 @@
 
 use async_trait::async_trait;
 use bytes::Bytes;
+use futures::StreamExt as _;
 use prost::Message;
 use std::sync::Arc;
 
@@ -14,8 +15,8 @@ use crate::rpc::{ClientRpc, PacketWriter};
 use crate::stream::{Context, Stream, StreamExt};
 use crate::transport::create_packet_channel;
 
-/// Receiver for incoming packets.
-pub type PacketReceiver = tokio::sync::mpsc::Receiver<crate::proto::Packet>;
+/// Owned incoming packets; dropping this stream releases its transport or nested call.
+pub type PacketReceiver = futures::stream::BoxStream<'static, crate::proto::Packet>;
 
 /// Trait for opening streams to a remote server.
 ///
@@ -40,12 +41,8 @@ pub trait Client: Send + Sync {
         I: Message + Send + Sync,
         O: Message + Default;
 
-    /// Opens a new stream for a streaming RPC.
-    ///
-    /// # Arguments
-    /// * `service` - The service ID
-    /// * `method` - The method ID
-    /// * `first_msg` - Optional initial message data
+    /// Opens a new stream for a streaming RPC, sending `first_msg` as the
+    /// initial data when present.
     async fn new_stream(
         &self,
         service: &str,
@@ -95,38 +92,36 @@ impl<T: OpenStream + 'static> Client for SrpcClient<T> {
             writer,
         ));
 
-        // Start the RPC with the input data.
+        // Cancellation and every early return release the directly owned incoming transport.
+        let _cancel = ctx.cancel_token().drop_guard_ref();
         rpc.start(Some(Bytes::from(input_data))).await?;
 
-        // Spawn a task to handle incoming packets.
-        let rpc_clone = rpc.clone();
-        let packet_handler = tokio::spawn(async move {
-            while let Some(packet) = receiver.recv().await {
-                if rpc_clone.handle_packet(packet).await.is_err() {
+        // Read packets alongside the unary response without a second task or lifetime.
+        let read = async {
+            while let Some(packet) = receiver.next().await {
+                if rpc.handle_packet(packet).await.is_err() {
                     break;
                 }
             }
-            let _ = rpc_clone.handle_stream_close(None).await;
-        });
-
-        // Send close to indicate we're done sending.
-        rpc.close_send().await?;
-
-        // Receive the response.
-        let output: O = rpc.msg_recv().await?;
-
-        // Wait for the RPC to complete properly (receive any trailing packets).
-        // This ensures we process any completion/error packets from the server
-        // before cleaning up.
-        let _ = rpc.wait().await;
-
-        // Close the RPC to signal completion.
-        let _ = rpc.close().await;
-
-        // Clean up the packet handler.
-        packet_handler.abort();
-
-        Ok(output)
+            let _ = rpc.handle_stream_close(None).await;
+        };
+        let response = async {
+            // A failed half-close means the call ended; msg_recv reports its outcome.
+            let _ = rpc.close_send().await;
+            let output: O = rpc.msg_recv().await?;
+
+            // Process trailing completion before releasing the writer and receive capability.
+            let _ = rpc.wait().await;
+            let _ = rpc.close().await;
+            Ok(output)
+        };
+        tokio::pin!(read, response);
+
+        // Transport closure settles the same RPC that the response future is observing.
+        tokio::select! {
+            output = &mut response => output,
+            () = &mut read => response.await,
+        }
     }
 
     async fn new_stream(
@@ -154,7 +149,7 @@ impl<T: OpenStream + 'static> Client for SrpcClient<T> {
         // Spawn a task to handle incoming packets.
         let rpc_clone = rpc.clone();
         let packet_handler = tokio::spawn(async move {
-            while let Some(packet) = receiver.recv().await {
+            while let Some(packet) = receiver.next().await {
                 if rpc_clone.handle_packet(packet).await.is_err() {
                     break;
                 }
@@ -178,6 +173,15 @@ struct ClientStream {
     packet_handler: tokio::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
 }
 
+impl Drop for ClientStream {
+    fn drop(&mut self) {
+        self.rpc.context().cancel();
+        if let Some(handler) = self.packet_handler.get_mut().take() {
+            handler.abort();
+        }
+    }
+}
+
 #[async_trait]
 impl Stream for ClientStream {
     fn context(&self) -> &Context {
@@ -199,9 +203,12 @@ impl Stream for ClientStream {
     async fn close(&self) -> Result<()> {
         let _ = self.rpc.close().await;
         // Abort the background packet handler to ensure cleanup.
-        if let Some(handle) = self.packet_handler.lock().await.take() {
+        let mut handler = self.packet_handler.lock().await;
+        if let Some(handle) = handler.as_mut() {
             handle.abort();
+            let _ = handle.await;
         }
+        handler.take();
         Ok(())
     }
 }
@@ -309,7 +316,12 @@ mod tests {
             (
                 Self {
                     writer: Arc::new(MockWriter::new()),
-                    receiver: Mutex::new(Some(rx)),
+                    receiver: Mutex::new(Some(Box::pin(futures::stream::unfold(
+                        rx,
+                        |mut receiver| async move {
+                            receiver.recv().await.map(|packet| (packet, receiver))
+                        },
+                    )))),
                 },
                 tx,
             )
@@ -340,6 +352,9 @@ mod tests {
             .unwrap();
 
         assert!(!stream.context().is_cancelled());
+        let context = stream.context().clone();
+        drop(stream);
+        assert!(context.is_cancelled());
     }
 
     #[tokio::test]
@@ -357,4 +372,29 @@ mod tests {
         let result2 = opener.open_stream().await;
         assert!(matches!(result2, Err(Error::StreamClosed)));
     }
+
+    #[tokio::test]
+    async fn malformed_unary_response_releases_receiver() {
+        #[derive(Clone, PartialEq, prost::Message)]
+        struct Reply {
+            #[prost(string, tag = "1")]
+            text: String,
+        }
+
+        let (opener, sender) = MockOpener::new();
+        let client = SrpcClient::new(opener);
+        sender
+            .send(crate::packet::new_call_data_full(
+                Some(Bytes::from_static(&[0xff])),
+                false,
+                None,
+            ))
+            .await
+            .unwrap();
+        let result: Result<Reply> = client.exec_call("test", "call", &Reply::default()).await;
+        assert!(result.is_err());
+        tokio::time::timeout(std::time::Duration::from_secs(2), sender.closed())
+            .await
+            .expect("failed unary call retained its packet receiver");
+    }
 }

package/srpc/client_half_close_test.cc

@@ -0,0 +1,331 @@
+//go:build deps_only
+
+/*
+ * Half-close checks for the generated C++ callers. Once the request is
+ * written, the half-close is only a notice: a caller that half-closes after
+ * the remote ended the call still reads the buffered responses and the
+ * call's outcome.
+ */
+
+#include "echo/echo_srpc.pb.hpp"
+#include "srpc/rpcproto.pb.h"
+#include "srpc/server-rpc.hpp"
+
+#include <cstdlib>
+#include <iostream>
+#include <latch>
+
+namespace {
+using starpc::Error;
+
+const char *body = "hello world via starpc half-close test";
+
+void Check(bool value, const char *message)
+{
+	if (value)
+		return;
+	std::cerr << message << '\n';
+	std::exit(1);
+}
+
+/*
+ * OnceEchoServer answers each streaming call once and completes it.
+ */
+class OnceEchoServer final : public echo::SRPCEchoerServer {
+public:
+	Error Echo(const echo::EchoMsg &req, echo::EchoMsg *resp) override
+	{
+		*resp = req;
+		return Error::OK;
+	}
+
+	Error
+	EchoServerStream(const echo::EchoMsg &req,
+			 echo::SRPCEchoer_EchoServerStreamStream *strm) override
+	{
+		return strm->Send(req);
+	}
+
+	/*
+	 * EchoClientStream replies to the first message without waiting for
+	 * the half-close.
+	 */
+	Error EchoClientStream(echo::SRPCEchoer_EchoClientStreamStream *strm,
+			       echo::EchoMsg *resp) override
+	{
+		return strm->Recv(resp);
+	}
+
+	Error EchoBidiStream(echo::SRPCEchoer_EchoBidiStreamStream *) override
+	{
+		return Error::Unimplemented;
+	}
+
+	Error RpcStream(echo::SRPCEchoer_RpcStreamStream *) override
+	{
+		return Error::Unimplemented;
+	}
+
+	Error DoNothing(const google::protobuf::Empty &,
+			google::protobuf::Empty *) override
+	{
+		return Error::OK;
+	}
+};
+
+/* FuncWriter forwards packets and the transport close to callbacks. */
+class FuncWriter final : public starpc::PacketWriter {
+public:
+	FuncWriter(std::function<Error(const srpc::Packet &)> write,
+		   std::function<void()> close)
+		: write(std::move(write)),
+		  close(std::move(close))
+	{
+	}
+
+	Error WritePacket(const srpc::Packet &packet) override
+	{
+		return write(packet);
+	}
+
+	Error Close() override
+	{
+		close();
+		return Error::OK;
+	}
+
+private:
+	std::function<Error(const srpc::Packet &)> write;
+	std::function<void()> close;
+};
+
+/*
+ * EchoPipe serves OnceEchoServer to one client stream in memory. closed
+ * counts down after the client's transport close callback returns.
+ */
+class EchoPipe {
+public:
+	EchoPipe()
+	{
+		auto [handler, err] =
+			echo::SRPCRegisterEchoer(mux.get(), &impl);
+		Check(err == Error::OK, "register echoer");
+		this->handler = std::move(handler);
+	}
+
+	/* OpenStream connects the client stream to the server. */
+	starpc::OpenStreamFunc OpenStream()
+	{
+		return [this](starpc::PacketDataHandler on_packet,
+			      starpc::CloseHandler on_close) {
+			server_writer = std::make_unique<FuncWriter>(
+				[on_packet](const srpc::Packet &packet) {
+					return on_packet(
+						packet.SerializeAsString());
+				},
+				[this, on_close] {
+					on_close(Error::OK);
+					closed.count_down();
+				});
+			server = starpc::NewServerRPC(mux.get(),
+						      server_writer.get());
+			auto client_writer = std::make_unique<FuncWriter>(
+				[this](const srpc::Packet &packet) {
+					return server->HandlePacket(packet);
+				},
+				[this] {
+					server->HandleStreamClose(Error::OK);
+				});
+			return std::make_pair(
+				std::unique_ptr<starpc::PacketWriter>(
+					std::move(client_writer)),
+				Error::OK);
+		};
+	}
+
+	std::latch closed{1};
+
+private:
+	OnceEchoServer impl;
+	std::unique_ptr<starpc::Mux> mux = starpc::NewMux();
+	std::unique_ptr<echo::SRPCEchoerHandler> handler;
+	std::unique_ptr<FuncWriter> server_writer;
+	std::unique_ptr<starpc::ServerRPC> server;
+};
+
+/*
+ * DroppingPipe accepts CallStart, then drops the transport without a
+ * response. closed counts down after the client's transport close callback
+ * returns.
+ */
+class DroppingPipe {
+public:
+	starpc::OpenStreamFunc OpenStream()
+	{
+		return [this](starpc::PacketDataHandler,
+			      starpc::CloseHandler on_close) {
+			auto writer = std::make_unique<FuncWriter>(
+				[this, on_close](const srpc::Packet &) {
+					dropper =
+						std::jthread([this, on_close] {
+							on_close(Error::EOF_);
+							closed.count_down();
+						});
+					return Error::OK;
+				},
+				[] {});
+			return std::make_pair(
+				std::unique_ptr<starpc::PacketWriter>(
+					std::move(writer)),
+				Error::OK);
+		};
+	}
+
+	std::latch closed{1};
+
+private:
+	std::jthread dropper;
+};
+
+/*
+ * EndedClient returns each new stream only after its transport has closed,
+ * so the generated caller half-closes after the remote already ended the
+ * call.
+ */
+class EndedClient final : public starpc::Client {
+public:
+	EndedClient(starpc::OpenStreamFunc open_stream, std::latch *closed)
+		: client(std::move(open_stream)),
+		  closed(closed)
+	{
+	}
+
+	Error ExecCall(const std::string &service, const std::string &method,
+		       const starpc::Message &in, starpc::Message *out) override
+	{
+		return client.ExecCall(service, method, in, out);
+	}
+
+	std::pair<std::unique_ptr<starpc::Stream>, Error>
+	NewStream(const std::string &service, const std::string &method,
+		  const starpc::Message *first_msg) override
+	{
+		auto opened = client.NewStream(service, method, first_msg);
+		if (opened.second == Error::OK)
+			closed->wait();
+		return opened;
+	}
+
+private:
+	starpc::ClientImpl client;
+	std::latch *closed;
+};
+
+/* ServerStreamCompletedBeforeCloseSend reads a reply after remote completion.
+ */
+void ServerStreamCompletedBeforeCloseSend()
+{
+	/* Complete the server stream before the generated caller half-closes.
+	 */
+	EchoPipe pipe;
+	EndedClient client(pipe.OpenStream(), &pipe.closed);
+	echo::SRPCEchoerClientImpl echoer(&client);
+	echo::EchoMsg req;
+	req.set_body(body);
+	auto [strm, err] = echoer.EchoServerStream(req);
+	Check(err == Error::OK,
+	      "server stream discarded after remote completion");
+
+	/* Read the buffered response, then the clean completion. */
+	echo::EchoMsg msg;
+	Check(strm->Recv(&msg) == Error::OK, "buffered response lost");
+	Check(msg.body() == body, "response body");
+	Check(strm->Recv(&msg) == Error::EOF_, "expected EOF after completion");
+}
+
+/* ClientStreamCompletedBeforeCloseSend reads a reply after a failed half-close.
+ */
+void ClientStreamCompletedBeforeCloseSend()
+{
+	EchoPipe pipe;
+	auto client = starpc::NewClient(pipe.OpenStream());
+	echo::SRPCEchoerClientImpl echoer(client.get());
+	auto [strm, err] = echoer.EchoClientStream();
+	Check(err == Error::OK, "open client stream");
+
+	/* Send the one request and wait for the call to end. */
+	echo::EchoMsg req;
+	req.set_body(body);
+	Check(strm->Send(req) == Error::OK, "send request");
+	pipe.closed.wait();
+
+	/* The half-close fails, but the reply is still readable. */
+	echo::EchoMsg msg;
+	Check(strm->CloseAndRecv(&msg) == Error::OK, "buffered reply lost");
+	Check(msg.body() == body, "reply body");
+}
+
+/* ServerStreamTransportFailureBeforeCloseSend reports a dropped transport. */
+void ServerStreamTransportFailureBeforeCloseSend()
+{
+	DroppingPipe pipe;
+	EndedClient client(pipe.OpenStream(), &pipe.closed);
+	echo::SRPCEchoerClientImpl echoer(&client);
+	echo::EchoMsg req;
+	req.set_body(body);
+	auto [strm, err] = echoer.EchoServerStream(req);
+	Check(err == Error::OK,
+	      "server stream discarded after transport failure");
+
+	/* Recv reports that the call ended without a verdict. */
+	echo::EchoMsg msg;
+	Check(strm->Recv(&msg) == Error::ClosedBeforeCompletion,
+	      "expected ClosedBeforeCompletion");
+}
+
+/*
+ * FailedHalfCloseKeepsPendingCompletion keeps the transport's buffered reply
+ * when the half-close write fails.
+ */
+void FailedHalfCloseKeepsPendingCompletion()
+{
+	/*
+	 * The half-close write fails while the transport still holds the
+	 * reply and completion it read but has not delivered.
+	 */
+	starpc::ClientRPC rpc("test", "call");
+	FuncWriter writer(
+		[](const srpc::Packet &packet) {
+			return packet.has_call_data() ? Error::Canceled
+						      : Error::OK;
+		},
+		[] {});
+	Check(rpc.Start(&writer, false, "") == Error::OK, "start");
+	auto strm = starpc::NewMsgStream(&rpc, [&] { rpc.Close(); });
+	Check(strm->CloseSend() == Error::Canceled, "half-close write fails");
+
+	/* The transport delivers what it read, then closes. */
+	srpc::CallData reply;
+	reply.set_data("reply");
+	reply.set_complete(true);
+	Check(rpc.HandleCallData(reply) == Error::OK,
+	      "deliver pending completion");
+	rpc.HandleStreamClose(Error::OK);
+
+	std::string data;
+	Check(rpc.ReadOne(&data) == Error::OK && data == "reply",
+	      "pending reply lost");
+	Check(rpc.ReadOne(&data) == Error::EOF_,
+	      "expected EOF after completion");
+}
+
+} // namespace
+
+int main()
+{
+	ServerStreamCompletedBeforeCloseSend();
+	ClientStreamCompletedBeforeCloseSend();
+	ServerStreamTransportFailureBeforeCloseSend();
+	FailedHalfCloseKeepsPendingCompletion();
+	std::cout << "Half-close checks passed\n";
+}

package/srpc/client_invoker.rs

@@ -0,0 +1,75 @@
+//! Proxy invocations through an existing client while retaining both directions.
+
+use async_trait::async_trait;
+
+use crate::{Client, Error, Invoker, Result, Stream};
+
+/// Routes incoming methods through one client, preserving the remote terminal result.
+pub struct ClientInvoker<C> {
+    client: C,
+}
+
+impl<C: Client> ClientInvoker<C> {
+    /// Wraps a client whose transport and routing remain under its existing owner.
+    pub fn new(client: C) -> Self {
+        Self { client }
+    }
+}
+
+#[async_trait]
+impl<C: Client> Invoker for ClientInvoker<C> {
+    async fn invoke_method(
+        &self,
+        service: &str,
+        method: &str,
+        stream: Box<dyn Stream>,
+    ) -> (bool, Result<()>) {
+        // Opening a remote stream remains within the incoming call's cancellation.
+        let remote = tokio::select! {
+            () = stream.context().cancelled() => return (true, Err(Error::Cancelled)),
+            remote = self.client.new_stream(service, method, None) => match remote {
+                Ok(remote) => remote,
+                Err(error) => return (true, Err(error)),
+            },
+        };
+
+        // Request forwarding may half-close before the remote sends its result.
+        let requests = async {
+            match forward_messages(stream.as_ref(), remote.as_ref()).await {
+                Ok(()) => remote.close_send().await,
+                Err(error) => {
+                    let _ = remote.close().await;
+                    Err(error)
+                }
+            }
+        };
+        let responses = forward_messages(remote.as_ref(), stream.as_ref());
+        tokio::pin!(requests, responses);
+
+        // A remote result settles the call even if its caller has not half-closed.
+        let exchange = async {
+            tokio::select! {
+                result = &mut responses => result,
+                _ = &mut requests => responses.await,
+            }
+        };
+        let result = tokio::select! {
+            result = exchange => result,
+            () = stream.context().cancelled() => Err(Error::Cancelled),
+        };
+        let _ = remote.close().await;
+        (true, result)
+    }
+}
+
+/// Copies every message, including empty protobufs, until the source half-closes.
+async fn forward_messages(source: &dyn Stream, target: &dyn Stream) -> Result<()> {
+    loop {
+        let data = match source.recv_bytes().await {
+            Ok(data) => data,
+            Err(Error::StreamClosed) => return Ok(()),
+            Err(error) => return Err(error),
+        };
+        target.send_bytes(data).await?;
+    }
+}