starpc 0.52.0 → 0.52.2

package/srpc/server-rpc.go

@@ -6,15 +6,15 @@ import (
 	"github.com/pkg/errors"
 )
 
-// ServerRPC represents the server side of an on-going RPC call message stream.
+// ServerRPC runs one remote method and retains its resources until the handler exits.
 type ServerRPC struct {
 	commonRPC
-	// invoker is the rpc call invoker
+
+	// invoker dispatches the selected service and method.
 	invoker Invoker
 }
 
-// NewServerRPC constructs a new ServerRPC session.
-// note: call SetWriter before handling any incoming messages.
+// NewServerRPC attaches the handler and transport for one incoming call.
 func NewServerRPC(ctx context.Context, invoker Invoker, writer PacketWriter) *ServerRPC {
 	rpc := &ServerRPC{invoker: invoker}
 	initCommonRPC(ctx, &rpc.commonRPC)
@@ -33,6 +33,7 @@ func (r *ServerRPC) HandlePacketData(data []byte) error {
 
 // HandlePacket handles an incoming parsed message packet.
 func (r *ServerRPC) HandlePacket(msg *Packet) error {
+	// Ignore an absent packet and reject malformed messages before dispatch.
 	if msg == nil {
 		return nil
 	}
@@ -40,6 +41,7 @@ func (r *ServerRPC) HandlePacket(msg *Packet) error {
 		return err
 	}
 
+	// Route each packet through the shared stream state.
 	switch b := msg.GetBody().(type) {
 	case *Packet_CallStart:
 		return r.HandleCallStart(b.CallStart)
@@ -57,25 +59,20 @@ func (r *ServerRPC) HandlePacket(msg *Packet) error {
 
 // HandleCallStart handles the call start packet.
 func (r *ServerRPC) HandleCallStart(pkt *CallStart) error {
-	var err error
-
+	// Accept one start before any terminal stream transition.
 	locked := r.bcast.Lock()
-	// process start: method and service
 	if r.method != "" || r.service != "" {
-		err = errors.New("call start must be sent only once")
 		locked.Unlock()
-		return err
+		return errors.New("call start must be sent only once")
 	}
 	if r.dataClosed {
-		err = ErrCompleted
 		locked.Unlock()
-		return err
+		return ErrCompleted
 	}
 
+	// Retain the request identity and optional first message for the handler.
 	service, method := pkt.GetRpcService(), pkt.GetRpcMethod()
 	r.service, r.method = service, method
-
-	// process first data packet, if included
 	if data := pkt.GetData(); len(data) != 0 || pkt.GetDataIsZero() {
 		r.dataQueue = append(r.dataQueue, data)
 	}
@@ -84,27 +81,32 @@ func (r *ServerRPC) HandleCallStart(pkt *CallStart) error {
 	// Mark the method active before scheduling it so cancellation cannot make
 	// Wait return and release a mux while invokeRPC is still running user code.
 	r.localActive = true
-
-	// invoke the rpc
 	locked.Broadcast()
 	startServerRPCInvoke(func() {
 		r.invokeRPC(service, method)
 	})
 	locked.Unlock()
 
-	return err
+	return nil
 }
 
 // invokeRPC invokes the RPC after CallStart is received.
 func (r *ServerRPC) invokeRPC(serviceID, methodID string) {
-	// on the server side, the writer is closed by invokeRPC.
+	// Run the selected handler before publishing its terminal result.
 	strm := NewMsgStream(r.ctx, r, r.cancelContext)
 	ok, err := r.invoker.InvokeMethod(serviceID, methodID, strm)
 	if err == nil && !ok {
 		err = ErrUnimplemented
 	}
-	r.beginLocalCompletion()
+
+	// Cancellation has its own packet; a text error would lose its identity.
 	outPkt := NewCallDataPacket(nil, false, true, err)
-	_ = r.writer.WritePacket(outPkt)
+	if errors.Is(err, context.Canceled) {
+		outPkt = NewCallCancelPacket()
+	}
+	r.beginLocalCompletion()
+	if err := r.writer.WritePacket(outPkt); err != nil {
+		r.HandleStreamClose(err)
+	}
 	r.finishLocalCompletion()
 }

package/srpc/server-rpc.hpp

@@ -1,14 +1,14 @@
 #pragma once
 
-#include <string>
-#include <thread>
-
 #include "common-rpc.hpp"
 #include "errors.hpp"
 #include "invoker.hpp"
 #include "msg-stream.hpp"
 #include "writer.hpp"
 
+#include <string>
+#include <thread>
+
 namespace srpc {
 class Packet;
 class CallStart;
@@ -16,48 +16,67 @@ class CallStart;
 
 namespace starpc {
 
-// ServerRPC represents the server side of an on-going RPC call message stream.
-// Matches Go ServerRPC struct in server-rpc.go
+/*
+ * ServerRPC is the remote's side of one call: it runs the handler the
+ * Invoker selects and cancels and joins it on destruction. The Invoker and
+ * PacketWriter must outlive it. HandlePacket calls are serialized by the
+ * transport; handlers observe cancellation through Stream::StopToken.
+ */
 class ServerRPC : public CommonRPC, public MsgStreamRw {
 public:
-  // Constructor matching NewServerRPC in Go
-  ServerRPC(Invoker *invoker, PacketWriter *writer);
-  ~ServerRPC() override;
-
-  // HandlePacketData handles an incoming unparsed message packet.
-  // Matches Go HandlePacketData in server-rpc.go
-  Error HandlePacketData(const std::string &data);
-
-  // HandlePacket handles an incoming parsed message packet.
-  // Matches Go HandlePacket in server-rpc.go
-  Error HandlePacket(const srpc::Packet &msg);
-
-  // HandleCallStart handles the call start packet.
-  // Matches Go HandleCallStart in server-rpc.go
-  Error HandleCallStart(const srpc::CallStart &pkt);
-
-  // MsgStreamRw interface implementation
-  Error ReadOne(std::string *out) override { return CommonRPC::ReadOne(out); }
-  Error WriteCallData(const std::string &data, bool data_is_zero, bool complete,
-                      Error err) override {
-    return CommonRPC::WriteCallData(data, data_is_zero, complete, err);
-  }
-  Error WriteCallCancel() override { return CommonRPC::WriteCallCancel(); }
+	/* ServerRPC serves one call through invoker, writing via writer. */
+	ServerRPC(Invoker *invoker, PacketWriter *writer);
+	~ServerRPC() override;
+
+	/* HandlePacketData parses one serialized packet and handles it. */
+	Error HandlePacketData(const std::string &data);
+
+	/* HandlePacket validates and dispatches one parsed packet. */
+	Error HandlePacket(const srpc::Packet &msg);
+
+	/*
+	 * HandleCallStart accepts the call's first packet and runs the
+	 * handler on its own thread.
+	 */
+	Error HandleCallStart(const srpc::CallStart &pkt);
+
+	/* MsgStreamRw over the shared call state. */
+	Error ReadOne(std::string *out) override
+	{
+		return CommonRPC::ReadOne(out);
+	}
+	Error WriteCallData(const std::string &data, bool data_is_zero,
+			    bool complete, Error err) override
+	{
+		return CommonRPC::WriteCallData(data, data_is_zero, complete,
+						err);
+	}
+	Error WriteCallCancel() override
+	{
+		return CommonRPC::WriteCallCancel();
+	}
+	std::string RemoteErrorMessage() const override
+	{
+		return CommonRPC::RemoteErrorMessage();
+	}
 
 private:
-  // InvokeRPC invokes the RPC after CallStart is received.
-  // Matches Go invokeRPC in server-rpc.go
-  void InvokeRPC(const std::string &service_id, const std::string &method_id);
+	/* InvokeRPC runs the handler and publishes its verdict. */
+	void InvokeRPC(const std::string &service_id,
+		       const std::string &method_id);
+
+	/* invoker is borrowed; the caller owns it. */
+	Invoker *invoker;
 
-  Invoker *invoker_;
-  std::thread invoke_thread_;
+	/* invoke_thread is last so it stops before the state it reads. */
+	std::jthread invoke_thread;
 };
 
-// NewServerRPC constructs a new ServerRPC session.
-// Matches Go NewServerRPC function in server-rpc.go
+/* NewServerRPC constructs a ServerRPC for one call. */
 inline std::unique_ptr<ServerRPC> NewServerRPC(Invoker *invoker,
-                                               PacketWriter *writer) {
-  return std::make_unique<ServerRPC>(invoker, writer);
+					       PacketWriter *writer)
+{
+	return std::make_unique<ServerRPC>(invoker, writer);
 }
 
 } // namespace starpc

package/srpc/server.rs

@@ -4,8 +4,7 @@
 //! The server supports all streaming patterns: unary, client streaming,
 //! server streaming, and bidirectional streaming.
 
-use bytes::Bytes;
-use futures::StreamExt;
+use futures::{stream::FuturesUnordered, StreamExt};
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -128,86 +127,48 @@ impl<I: Invoker + 'static> Server<I> {
         // Create framed reader.
         let mut framed = FramedRead::new(read_half, PacketCodec::new());
 
-        // Wait for the first packet (CallStart).
+        // Wait for the first packet and take its CallStart body.
         let first_packet = framed.next().await;
-        let call_start = match first_packet {
-            Some(Ok(packet)) => {
-                // Validate the packet
-                packet.validate()?;
-
-                match packet.body {
-                    Some(Body::CallStart(cs)) => cs,
-                    _ => return Err(Error::ExpectedCallStart),
-                }
-            }
-            Some(Err(e)) => return Err(e),
-            None => return Err(Error::StreamClosed),
+        let Some(Ok(packet)) = first_packet else {
+            // A decode error ends the stream; a closed transport has no error.
+            return first_packet.map_or(Err(Error::StreamClosed), |res| res.map(|_| ()));
+        };
+        packet.validate()?;
+        let Some(Body::CallStart(call_start)) = packet.body else {
+            return Err(Error::ExpectedCallStart);
         };
 
-        // Validate service and method IDs (already done by packet.validate(),
-        // but double-check for clarity).
-        if call_start.rpc_service.is_empty() {
-            return Err(Error::EmptyServiceId);
-        }
-        if call_start.rpc_method.is_empty() {
-            return Err(Error::EmptyMethodId);
-        }
-
-        let service_id = call_start.rpc_service.clone();
-        let method_id = call_start.rpc_method.clone();
-
-        // Create the server RPC.
-        let ctx = Context::new();
-        let rpc = Arc::new(ServerRpc::from_call_start(ctx, call_start, writer));
+        // Retain the call's transport context until this borrowing operation ends.
+        let rpc = Arc::new(ServerRpc::from_call_start(
+            Context::new(),
+            call_start,
+            writer,
+        ));
+        let _cancel = rpc.context().cancel_token().drop_guard_ref();
 
-        // Spawn a task to read remaining packets.
-        let rpc_clone = rpc.clone();
-        let mut read_task = tokio::spawn(async move {
+        // The packet reader is a sibling future, so dropping this call cannot detach transport work.
+        let read = async {
             while let Some(result) = framed.next().await {
                 match result {
                     Ok(packet) => {
-                        if rpc_clone.handle_packet(packet).await.is_err() {
+                        if rpc.handle_packet(packet).await.is_err() {
                             break;
                         }
                     }
                     Err(_) => break,
                 }
             }
-            let _ = rpc_clone.handle_stream_close(None).await;
-        });
-
-        // Invoke the method.
-        let stream: Box<dyn Stream> = Box::new(ServerStream { rpc: rpc.clone() });
-        let (found, result) = self
-            .invoker
-            .invoke_method(&service_id, &method_id, stream)
-            .await;
-
-        // Handle the result.
-        if !found {
-            // Send unimplemented error.
-            let _ = rpc.send_error("method not implemented".to_string()).await;
-        } else if let Err(e) = result {
-            // Send error response.
-            let _ = rpc.send_error(e.to_string()).await;
-        } else {
-            // Close send side on success.
-            let _ = rpc.close_send().await;
-        }
-
-        // Explicitly close the RPC to release the writer/transport.
-        // This ensures the connection doesn't remain open after the terminal response.
-        let _ = rpc.close().await;
+            let _ = rpc.handle_stream_close(None).await;
+        };
+        let invoke = serve_rpc(self.invoker.as_ref(), &rpc);
+        tokio::pin!(read, invoke);
 
-        // Wait for the read task to finish or timeout, then abort it.
-        // This gives time for the client to receive our response.
+        // Remote closure cancels the same context that the active service receives.
+        // After a reply, retain the existing bounded grace period for the peer's transport close.
         tokio::select! {
-            _ = &mut read_task => {
-                // Read task finished naturally
-            }
-            _ = tokio::time::sleep(self.config.shutdown_timeout) => {
-                // Timeout - abort the read task
-                read_task.abort();
+            () = &mut read => invoke.await,
+            () = &mut invoke => {
+                let _ = tokio::time::timeout(self.config.shutdown_timeout, read).await;
             }
         }
 
@@ -252,73 +213,59 @@ impl<I: Invoker + 'static> Server<I> {
             .await
     }
 
-    /// Accepts and handles connections in a loop.
-    ///
-    /// This is a convenience method that accepts connections from a listener
-    /// and spawns a task to handle each one.
-    ///
-    /// Errors from individual connections are reported via the error handler
-    /// (if configured) but don't stop the server.
+    /// Accepts and handles connections concurrently, draining accepted calls when the listener ends.
+    /// Dropping this future drops every active call and cancels its transport context.
+    /// Individual connection failures reach the error handler without stopping other calls.
     pub async fn serve<L, T>(&self, mut listener: L) -> Result<()>
     where
         L: futures::Stream<Item = std::io::Result<T>> + Unpin,
         T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
     {
-        while let Some(result) = listener.next().await {
-            match result {
-                Ok(stream) => {
-                    let server = self.clone_for_spawn();
-                    tokio::spawn(async move {
-                        if let Err(e) = server.handle_stream(stream).await {
-                            server.report_error(e);
-                        }
-                    });
-                }
-                Err(e) => {
-                    self.report_error(Error::Io(e));
+        // Concurrent call futures remain owned by this invocation, without detached tasks.
+        let mut calls = FuturesUnordered::new();
+        let mut accepting = true;
+        loop {
+            tokio::select! {
+                connection = listener.next(), if accepting => {
+                    match connection {
+                        Some(Ok(stream)) => calls.push(async move {
+                            if let Err(error) = self.handle_stream(stream).await {
+                                self.report_error(error);
+                            }
+                        }),
+                        Some(Err(error)) => self.report_error(Error::Io(error)),
+                        None => accepting = false,
+                    }
                 }
+                Some(()) = calls.next(), if !calls.is_empty() => {}
+                else => return Ok(()),
             }
         }
-
-        Ok(())
-    }
-
-    /// Creates a clone of the server for spawning tasks.
-    pub(crate) fn clone_for_spawn(&self) -> Server<I> {
-        Server {
-            invoker: self.invoker.clone(),
-            config: self.config.clone(),
-            error_handler: self.error_handler.clone(),
-        }
     }
 }
 
-/// Wrapper to provide Stream interface for ServerRpc.
-struct ServerStream {
-    rpc: Arc<ServerRpc>,
-}
-
-#[async_trait::async_trait]
-impl Stream for ServerStream {
-    fn context(&self) -> &Context {
-        self.rpc.context()
-    }
-
-    async fn send_bytes(&self, data: Bytes) -> Result<()> {
-        self.rpc.send_bytes(data).await
-    }
-
-    async fn recv_bytes(&self) -> Result<Bytes> {
-        self.rpc.recv_bytes().await
-    }
-
-    async fn close_send(&self) -> Result<()> {
-        self.rpc.close_send().await
-    }
-
-    async fn close(&self) -> Result<()> {
-        self.rpc.close().await
-    }
+/// Invokes the method a server RPC names and settles the call with the outcome.
+///
+/// An unknown method or a handler error reaches the peer as an error, and a
+/// handler that returns cleanly closes the send side. The RPC is closed last,
+/// which releases its writer.
+pub(crate) async fn serve_rpc<I: Invoker + ?Sized>(invoker: &I, rpc: &Arc<ServerRpc>) {
+    // Invoke the method with the RPC as its stream.
+    let (found, result) = invoker
+        .invoke_method(rpc.service(), rpc.method(), Box::new(rpc.clone()))
+        .await;
+
+    // Report the outcome to the peer; a failed write leaves nothing to settle.
+    let _ = if !found {
+        rpc.send_error("method not implemented".to_string()).await
+    } else if let Err(err) = result {
+        rpc.send_error(err.to_string()).await
+    } else {
+        rpc.close_send().await
+    };
+
+    // Close the RPC so the connection does not stay open after the terminal response.
+    let _ = rpc.close().await;
 }
 
 #[cfg(test)]
@@ -330,10 +277,9 @@ mod tests {
     #[tokio::test]
     async fn test_server_config() {
         let mux = Mux::new();
-        let server = Server::new(mux)
-            .with_config(ServerConfig {
-                shutdown_timeout: Duration::from_secs(1),
-            });
+        let server = Server::new(mux).with_config(ServerConfig {
+            shutdown_timeout: Duration::from_secs(1),
+        });
 
         assert_eq!(server.config.shutdown_timeout, Duration::from_secs(1));
     }
@@ -346,10 +292,9 @@ mod tests {
         let errors_clone = errors.clone();
 
         let mux = Mux::new();
-        let server = Server::new(mux)
-            .with_error_handler(move |e| {
-                errors_clone.lock().unwrap().push(e.to_string());
-            });
+        let server = Server::new(mux).with_error_handler(move |e| {
+            errors_clone.lock().unwrap().push(e.to_string());
+        });
 
         // Report an error
         server.report_error(Error::StreamClosed);

package/srpc/server.test.ts

@@ -1,5 +1,7 @@
 import { describe, it, beforeEach, expect, vi } from 'vitest'
 import { pipe } from 'it-pipe'
+import { pushable } from 'it-pushable'
+import type { Source } from 'it-stream-types'
 import {
   createHandler,
   createMux,
@@ -175,6 +177,8 @@ describe('srpc server', () => {
     const server = new Server(mux.lookupMethod)
     const firstResponse = new Promise<Packet>((resolve, reject) => {
       server.handlePacketStream({
+        close: async () => {},
+        abort: () => {},
         source: (async function* () {
           yield Packet.toBinary({
             body: {
@@ -342,6 +346,52 @@ describe('srpc server', () => {
     await Promise.resolve()
   })
 
+  it('closes the packet stream when the server pipeline completes', async () => {
+    const server = new Server(createMux().lookupMethod)
+    const close = vi.fn(async () => {})
+    const source = pushable<Uint8Array>({ objectMode: true })
+    const stream = {
+      close,
+      abort: vi.fn(),
+      source,
+      sink: async (output: Source<Uint8Array>) => {
+        for await (const _packet of output) {
+          // Drain the response pipeline.
+        }
+      },
+    }
+    const rpc = server.handlePacketStream(stream)
+
+    await rpc.close()
+    source.end()
+
+    await vi.waitFor(() => expect(close).toHaveBeenCalledOnce())
+    expect(stream.abort).not.toHaveBeenCalled()
+  })
+
+  it('aborts the packet stream when the server pipeline fails', async () => {
+    const server = new Server(createMux().lookupMethod)
+    const error = new Error('input failed')
+    const abort = vi.fn()
+    const source = pushable<Uint8Array>({ objectMode: true })
+    const stream = {
+      close: vi.fn(async () => {}),
+      abort,
+      source,
+      sink: async (output: Source<Uint8Array>) => {
+        for await (const _packet of output) {
+          // Drain the response pipeline.
+        }
+      },
+    }
+
+    server.handlePacketStream(stream)
+    source.end(error)
+
+    await vi.waitFor(() => expect(abort).toHaveBeenCalledWith(error))
+    expect(stream.close).not.toHaveBeenCalled()
+  })
+
   it('tears down passive channel close state', async () => {
     const { port1, port2 } = new MessageChannel()
     const opts: ChannelStreamOpts = { idleTimeoutMs: 1000, keepAliveMs: 1000 }

package/srpc/server.ts

@@ -20,9 +20,7 @@ export class Server implements StreamHandler {
   public get rpcStreamHandler(): HandleStreamFunc {
     return async (stream: PacketStream) => {
       const rpc = this.startRpc()
-      return pipe(stream, decodePacketSource, rpc, encodePacketSource, stream)
-        .catch((err: Error) => rpc.close(err))
-        .then(() => rpc.close())
+      return runPacketStream(stream, rpc)
     }
   }
 
@@ -36,9 +34,27 @@ export class Server implements StreamHandler {
   // the stream has one Uint8Array per packet w/o length prefix.
   public handlePacketStream(stream: PacketStream): ServerRPC {
     const rpc = this.startRpc()
-    pipe(stream, decodePacketSource, rpc, encodePacketSource, stream)
-      .catch((err: Error) => rpc.close(err))
-      .then(() => rpc.close())
+    void runPacketStream(stream, rpc).catch(() => undefined)
     return rpc
   }
 }
+
+async function runPacketStream(
+  stream: PacketStream,
+  rpc: ServerRPC,
+): Promise<void> {
+  try {
+    await pipe(stream, decodePacketSource, rpc, encodePacketSource, stream)
+    if (rpc.isClosed instanceof Error) {
+      stream.abort(rpc.isClosed)
+      throw rpc.isClosed
+    }
+    await stream.close()
+    await rpc.close()
+  } catch (err) {
+    const error = err instanceof Error ? err : new Error(String(err))
+    stream.abort(error)
+    await rpc.close(error)
+    throw error
+  }
+}

package/srpc/starpc.hpp

@@ -1,7 +1,9 @@
 #pragma once
 
-// Main header for starpc C++ library
-// Include this to get all starpc functionality
+/*
+ * starpc.hpp is the library's umbrella header: include it for the whole
+ * hand-written C++ API.
+ */
 
 #include "client-rpc.hpp"
 #include "client.hpp"

package/srpc/stream.go

@@ -17,6 +17,10 @@ type Stream interface {
 	MsgRecv(msg Message) error
 
 	// CloseSend signals to the remote that we will no longer send any messages.
+	//
+	// An error means the signal was not written, usually because the call
+	// already ended. MsgRecv still returns every message received before the
+	// end, then the call's outcome.
 	CloseSend() error
 
 	// Close closes the stream for reading and writing.