starpc 0.52.1 → 0.52.2

package/srpc/accept.go

@@ -7,14 +7,20 @@ import (
 	"github.com/libp2p/go-yamux/v5"
 )
 
-// AcceptMuxedListener accepts incoming connections from a net.Listener.
-//
-// Uses the default yamux muxer.
-// If yamux conf is nil, uses the defaults.
+// AcceptMuxedListener accepts connections until ctx is canceled or the listener
+// fails. Cancellation closes lis to interrupt Accept. A nil yamuxConf uses defaults.
 func AcceptMuxedListener(ctx context.Context, lis net.Listener, srv *Server, yamuxConf *yamux.Config) error {
+	// Interrupt a blocked Accept when its serving context ends.
+	stop := context.AfterFunc(ctx, func() { _ = lis.Close() })
+	defer stop()
+
+	// Transfer accepted connections to the server and release rejected ones.
 	for {
 		nc, err := lis.Accept()
 		if err != nil {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
 			return err
 		}
 

package/srpc/accept_test.go

@@ -0,0 +1,25 @@
+package srpc
+
+import (
+	"context"
+	"errors"
+	"net"
+	"testing"
+)
+
+func TestAcceptMuxedListenerCancellation(t *testing.T) {
+	// An idle listener must not retain a blocked Accept after cancellation.
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = listener.Close() })
+	ctx, cancel := context.WithCancel(t.Context())
+	t.Cleanup(cancel)
+	done := make(chan error, 1)
+	go func() { done <- AcceptMuxedListener(ctx, listener, nil, nil) }()
+	cancel()
+	if err := <-done; !errors.Is(err, context.Canceled) {
+		t.Fatalf("listener cancellation returned %v", err)
+	}
+}

package/srpc/build.rs

@@ -3,12 +3,21 @@
 //! Enable the `build` feature and call [`configure`] from a downstream
 //! `build.rs` to install the Starpc service generator into `prost-build`.
 
+/// Precedes each generated `close_send` whose error is ignored.
+///
+/// A failed half-close means the call already ended; returning that error would
+/// discard responses the stream already received, and `msg_recv` reports the
+/// call's outcome after them.
+const HALF_CLOSE_COMMENT: &str =
+    "        // A failed half-close means the call ended; msg_recv reports its outcome.";
+
 /// Configures `prost-build` to generate Starpc service bindings.
 ///
 /// The returned config generates normal Prost message types and appends Starpc
 /// client/server traits, client implementations, stream wrappers, and handler
 /// glue for every protobuf service.
 pub fn configure() -> prost_build::Config {
+    // Use the packaged compiler and its standard protobuf imports when available.
     let mut config = prost_build::Config::new();
     if let Ok(protoc) = protoc_bin_vendored::protoc_bin_path() {
         config.protoc_executable(protoc);
@@ -16,39 +25,49 @@ pub fn configure() -> prost_build::Config {
     if let Ok(include) = protoc_bin_vendored::include_path() {
         config.protoc_arg(format!("--proto_path={}", include.display()));
     }
+
+    // Share runtime message types and generate services beside Prost messages.
     config.extern_path(".rpcstream", "::starpc::rpcstream");
-    config.service_generator(Box::new(StarpcServiceGenerator::default()));
+    config.service_generator(Box::new(StarpcServiceGenerator));
     config
 }
 
-/// StarpcServiceGenerator emits Starpc client and server glue for Prost services.
+/// Emits StarPC client and server bindings for Prost services.
 #[derive(Debug, Default)]
 pub struct StarpcServiceGenerator;
 
 impl prost_build::ServiceGenerator for StarpcServiceGenerator {
+    /// Appends one service's clients, streams, and dispatch implementation.
     fn generate(&mut self, service: prost_build::Service, buf: &mut String) {
         let mut gen = Generator::new(service, buf);
         gen.generate();
     }
 }
 
+/// Emits bindings for one service into its protobuf package's output module.
 struct Generator<'a> {
+    /// Service identity, message types, and streaming modes supplied by Prost.
     service: prost_build::Service,
+    /// Package output shared by every service in the descriptor graph.
     buf: &'a mut String,
 }
 
 impl<'a> Generator<'a> {
+    /// Borrows the package output for this service's generation pass.
     fn new(service: prost_build::Service, buf: &'a mut String) -> Self {
         Self { service, buf }
     }
 
+    /// Emits the complete client and server contract for this service.
     fn generate(&mut self) {
+        // Resolve stable identifiers before emitting any declarations.
         let service_id = self.service_id();
         let service_name = self.service.name.clone();
 
+        // Keep extension methods available without adding a shared import name.
         self.line("");
         self.line("#[allow(unused_imports)]");
-        self.line("use starpc::StreamExt;");
+        self.line("use starpc::StreamExt as _;");
         self.line("");
         self.line(&format!("/// Service ID for {}.", service_name));
         self.line(&format!(
@@ -58,6 +77,7 @@ impl<'a> Generator<'a> {
         ));
         self.line("");
 
+        // Compose the client and server surfaces around the same wire identity.
         self.generate_stream_traits();
         self.generate_client_trait();
         self.generate_client_impl();
@@ -66,15 +86,17 @@ impl<'a> Generator<'a> {
         self.generate_handler();
     }
 
+    /// Declares the send, receive, and close operations for each streaming call.
     fn generate_stream_traits(&mut self) {
+        // Emit the client stream contract only for streaming methods.
         let service_name = self.service.name.clone();
         for method in self.service.methods.clone() {
             if !method.client_streaming && !method.server_streaming {
                 continue;
             }
 
+            // Name the typed stream and retain its common transport contract.
             let stream_trait = stream_trait_name(&service_name, &method);
-
             self.line(&format!(
                 "/// Stream trait for {}.{}.",
                 service_name, method.proto_name
@@ -84,6 +106,7 @@ impl<'a> Generator<'a> {
             self.line("    /// Returns the context for this stream.");
             self.line("    fn context(&self) -> &starpc::Context;");
 
+            // Provide typed sends when the caller owns the request stream.
             if method.client_streaming {
                 self.line("    /// Sends a message on the stream.");
                 self.line(&format!(
@@ -92,6 +115,7 @@ impl<'a> Generator<'a> {
                 ));
             }
 
+            // Provide typed receives when the server owns the response stream.
             if method.server_streaming {
                 self.line("    /// Receives a message from the stream.");
                 self.line(&format!(
@@ -100,6 +124,7 @@ impl<'a> Generator<'a> {
                 ));
             }
 
+            // Preserve the final response after half-closing a request stream.
             if method.client_streaming && !method.server_streaming {
                 self.line("    /// Closes the send side and receives the response.");
                 self.line(&format!(
@@ -110,18 +135,20 @@ impl<'a> Generator<'a> {
                 self.line("    /// Closes the stream.");
                 self.line("    async fn close(&self) -> starpc::Result<()>;");
             }
-
             self.line("}");
             self.line("");
         }
     }
 
+    /// Declares the caller-facing methods for all four RPC shapes.
     fn generate_client_trait(&mut self) {
+        // Declare the caller contract before its typed methods.
         let service_name = self.service.name.clone();
         self.line(&format!("/// Client trait for {}.", service_name));
         self.line("#[starpc::async_trait]");
         self.line(&format!("pub trait {}Client: Send + Sync {{", service_name));
 
+        // Select each method signature from the descriptor streaming modes.
         for method in self.service.methods.clone() {
             self.line(&format!("    /// {}.", method.proto_name));
             self.line(&format!(
@@ -129,18 +156,25 @@ impl<'a> Generator<'a> {
                 client_trait_method(&service_name, &method)
             ));
         }
-
         self.line("}");
         self.line("");
     }
 
+    /// Implements calls through a retained transport client.
     fn generate_client_impl(&mut self) {
+        // Retain even an empty service's client for its declared lifetime.
         let service_id = self.service_id();
         let service_name = self.service.name.clone();
+        let client_field = if self.service.methods.is_empty() {
+            "_client"
+        } else {
+            "client"
+        };
 
+        // Emit the client handle and its constructor.
         self.line(&format!("/// Client implementation for {}.", service_name));
         self.line(&format!("pub struct {}ClientImpl<C> {{", service_name));
-        self.line("    client: C,");
+        self.line(&format!("    {client_field}: C,"));
         self.line("}");
         self.line("");
         self.line(&format!(
@@ -149,7 +183,11 @@ impl<'a> Generator<'a> {
         ));
         self.line("    /// Creates a new client.");
         self.line("    pub fn new(client: C) -> Self {");
-        self.line("        Self { client }");
+        if self.service.methods.is_empty() {
+            self.line("        Self { _client: client }");
+        } else {
+            self.line("        Self { client }");
+        }
         self.line("    }");
         self.line("}");
         self.line("");
@@ -159,6 +197,7 @@ impl<'a> Generator<'a> {
             service_name, service_name
         ));
 
+        // Implement each call using its wire identity and typed stream adapter.
         for method in self.service.methods.clone() {
             let method_name = &method.name;
             if method.client_streaming && method.server_streaming {
@@ -189,7 +228,8 @@ impl<'a> Generator<'a> {
                     "        let stream = self.client.new_stream({:?}, {:?}, Some(&data)).await?;",
                     service_id, method.proto_name
                 ));
-                self.line("        stream.close_send().await?;");
+                self.line(HALF_CLOSE_COMMENT);
+                self.line("        let _ = stream.close_send().await;");
                 self.line(&format!(
                     "        Ok(Box::new({}Impl {{ stream }}))",
                     stream_trait_name(&service_name, &method)
@@ -222,20 +262,21 @@ impl<'a> Generator<'a> {
                 self.line("    }");
             }
         }
-
         self.line("}");
         self.line("");
     }
 
+    /// Adapts each transport stream to its generated typed client contract.
     fn generate_stream_impls(&mut self) {
+        // Adapt each streaming method through its retained transport stream.
         let service_name = self.service.name.clone();
         for method in self.service.methods.clone() {
             if !method.client_streaming && !method.server_streaming {
                 continue;
             }
 
+            // Name the typed stream and retain its common transport contract.
             let stream_trait = stream_trait_name(&service_name, &method);
-
             self.line(&format!("struct {}Impl {{", stream_trait));
             self.line("    stream: Box<dyn starpc::Stream>,");
             self.line("}");
@@ -249,6 +290,7 @@ impl<'a> Generator<'a> {
             self.line("        self.stream.context()");
             self.line("    }");
 
+            // Provide typed sends when the caller owns the request stream.
             if method.client_streaming {
                 self.line(&format!(
                     "    async fn send(&self, msg: &{}) -> starpc::Result<()> {{",
@@ -258,6 +300,7 @@ impl<'a> Generator<'a> {
                 self.line("    }");
             }
 
+            // Provide typed receives when the server owns the response stream.
             if method.server_streaming {
                 self.line(&format!(
                     "    async fn recv(&self) -> starpc::Result<{}> {{",
@@ -267,12 +310,14 @@ impl<'a> Generator<'a> {
                 self.line("    }");
             }
 
+            // Preserve the final response after half-closing a request stream.
             if method.client_streaming && !method.server_streaming {
                 self.line(&format!(
                     "    async fn close_and_recv(&self) -> starpc::Result<{}> {{",
                     method.output_type
                 ));
-                self.line("        self.stream.close_send().await?;");
+                self.line(HALF_CLOSE_COMMENT);
+                self.line("        let _ = self.stream.close_send().await;");
                 self.line("        self.stream.msg_recv().await");
                 self.line("    }");
             } else {
@@ -280,19 +325,20 @@ impl<'a> Generator<'a> {
                 self.line("        self.stream.close().await");
                 self.line("    }");
             }
-
             self.line("}");
             self.line("");
         }
     }
 
+    /// Declares the implementation contract consumed by handler dispatch.
     fn generate_server_trait(&mut self) {
+        // Declare the implementation contract consumed by dispatch.
         let service_name = self.service.name.clone();
-
         self.line(&format!("/// Server trait for {}.", service_name));
         self.line("#[starpc::async_trait]");
         self.line(&format!("pub trait {}Server: Send + Sync {{", service_name));
 
+        // Select each method signature from the descriptor streaming modes.
         for method in self.service.methods.clone() {
             self.line(&format!("    /// {}.", method.proto_name));
             if method.client_streaming && method.server_streaming {
@@ -312,34 +358,43 @@ impl<'a> Generator<'a> {
                 ));
             } else {
                 self.line(&format!(
-                    "    async fn {}(&self, request: {}) -> starpc::Result<{}>;",
+                    "    async fn {}(&self, context: &starpc::Context, request: {}) -> starpc::Result<{}>;",
                     method.name, method.input_type, method.output_type
                 ));
             }
         }
-
         self.line("}");
         self.line("");
     }
 
+    /// Routes wire method names through the exact server trait.
     fn generate_handler(&mut self) {
+        // Retain empty services without emitting unused field or parameter names.
         let service_id = self.service_id();
         let service_name = self.service.name.clone();
         let methods = self.service.methods.clone();
         let method_ids = method_ids_const(&service_name);
+        let (server_field, stream_parameter, method_parameter) = if methods.is_empty() {
+            ("_server", "_stream", "_method_id")
+        } else {
+            ("server", "stream", "method_id")
+        };
 
+        // Emit the supported wire method identities.
         self.line(&format!("const {}: &[&str] = &[", method_ids));
         for method in &methods {
             self.line(&format!("    {:?},", method.proto_name));
         }
         self.line("];");
         self.line("");
+
+        // Retain the server through both owned and shared construction.
         self.line(&format!("/// Handler for {}.", service_name));
         self.line(&format!(
             "pub struct {}Handler<S: {}Server> {{",
             service_name, service_name
         ));
-        self.line("    server: std::sync::Arc<S>,");
+        self.line(&format!("    {server_field}: std::sync::Arc<S>,"));
         self.line("}");
         self.line("");
         self.line(&format!(
@@ -348,15 +403,23 @@ impl<'a> Generator<'a> {
         ));
         self.line("    /// Creates a new handler wrapping the server implementation.");
         self.line("    pub fn new(server: S) -> Self {");
-        self.line("        Self { server: std::sync::Arc::new(server) }");
+        self.line(&format!(
+            "        Self {{ {server_field}: std::sync::Arc::new(server) }}"
+        ));
         self.line("    }");
         self.line("");
         self.line("    /// Creates a new handler with a shared server.");
         self.line("    pub fn with_arc(server: std::sync::Arc<S>) -> Self {");
-        self.line("        Self { server }");
+        if methods.is_empty() {
+            self.line("        Self { _server: server }");
+        } else {
+            self.line("        Self { server }");
+        }
         self.line("    }");
         self.line("}");
         self.line("");
+
+        // Dispatch through the named trait so Arc methods cannot shadow RPCs.
         self.line("#[starpc::async_trait]");
         self.line(&format!(
             "impl<S: {}Server + 'static> starpc::Invoker for {}Handler<S> {{",
@@ -365,17 +428,24 @@ impl<'a> Generator<'a> {
         self.line("    async fn invoke_method(");
         self.line("        &self,");
         self.line("        _service_id: &str,");
-        self.line("        method_id: &str,");
-        self.line("        stream: Box<dyn starpc::Stream>,");
+        self.line(&format!("        {method_parameter}: &str,"));
+        self.line(&format!(
+            "        {stream_parameter}: Box<dyn starpc::Stream>,"
+        ));
         self.line("    ) -> (bool, starpc::Result<()>) {");
-        self.line("        match method_id {");
+        if methods.is_empty() {
+            self.line("        (false, Err(starpc::Error::Unimplemented))");
+        } else {
+            self.line("        match method_id {");
+        }
 
+        // Receive and return the messages required by each RPC shape.
         for method in &methods {
             self.line(&format!("            {:?} => {{", method.proto_name));
             if method.client_streaming && method.server_streaming {
                 self.line(&format!(
-                    "                (true, self.server.{}(stream).await)",
-                    method.name
+                    "                (true, <S as {}Server>::{}(self.server.as_ref(), stream).await)",
+                    service_name, method.name
                 ));
             } else if method.server_streaming {
                 self.line(&format!(
@@ -386,13 +456,13 @@ impl<'a> Generator<'a> {
                 self.line("                    Err(e) => return (true, Err(e)),");
                 self.line("                };");
                 self.line(&format!(
-                    "                (true, self.server.{}(request, stream).await)",
-                    method.name
+                    "                (true, <S as {}Server>::{}(self.server.as_ref(), request, stream).await)",
+                    service_name, method.name
                 ));
             } else if method.client_streaming {
                 self.line(&format!(
-                    "                match self.server.{}(stream.as_ref()).await {{",
-                    method.name
+                    "                match <S as {}Server>::{}(self.server.as_ref(), stream.as_ref()).await {{",
+                    service_name, method.name
                 ));
                 self.line("                    Ok(response) => {");
                 self.line(
@@ -413,8 +483,8 @@ impl<'a> Generator<'a> {
                 self.line("                    Err(e) => return (true, Err(e)),");
                 self.line("                };");
                 self.line(&format!(
-                    "                match self.server.{}(request).await {{",
-                    method.name
+                    "                match <S as {}Server>::{}(self.server.as_ref(), stream.context(), request).await {{",
+                    service_name, method.name
                 ));
                 self.line("                    Ok(response) => {");
                 self.line(
@@ -430,11 +500,16 @@ impl<'a> Generator<'a> {
             self.line("            }");
         }
 
-        self.line("            _ => (false, Err(starpc::Error::Unimplemented)),");
-        self.line("        }");
+        // Leave unknown wire methods available for another handler.
+        if !methods.is_empty() {
+            self.line("            _ => (false, Err(starpc::Error::Unimplemented)),");
+            self.line("        }");
+        }
         self.line("    }");
         self.line("}");
         self.line("");
+
+        // Advertise the service and method identities used for routing.
         self.line(&format!(
             "impl<S: {}Server + 'static> starpc::Handler for {}Handler<S> {{",
             service_name, service_name
@@ -450,6 +525,7 @@ impl<'a> Generator<'a> {
         self.line("");
     }
 
+    /// Returns the fully qualified protobuf service identity used on the wire.
     fn service_id(&self) -> String {
         if self.service.package.is_empty() {
             self.service.proto_name.clone()
@@ -458,12 +534,14 @@ impl<'a> Generator<'a> {
         }
     }
 
+    /// Appends one complete generated line to the package output.
     fn line(&mut self, line: &str) {
         self.buf.push_str(line);
         self.buf.push('\n');
     }
 }
 
+/// Formats the caller signature selected by a method's streaming modes.
 fn client_trait_method(service_name: &str, method: &prost_build::Method) -> String {
     let stream_trait = stream_trait_name(service_name, method);
     if method.client_streaming && method.server_streaming {
@@ -489,6 +567,7 @@ fn client_trait_method(service_name: &str, method: &prost_build::Method) -> Stri
     }
 }
 
+/// Names a stream within its service, avoiding collisions with other services.
 fn stream_trait_name(service_name: &str, method: &prost_build::Method) -> String {
     format!(
         "{}{}Stream",
@@ -497,14 +576,17 @@ fn stream_trait_name(service_name: &str, method: &prost_build::Method) -> String
     )
 }
 
+/// Names the public wire-service identity constant.
 fn service_id_const(service_name: &str) -> String {
     format!("{}_SERVICE_ID", screaming_snake(service_name))
 }
 
+/// Names the handler's supported-method list.
 fn method_ids_const(service_name: &str) -> String {
     format!("{}_METHOD_IDS", screaming_snake(service_name))
 }
 
+/// Separates uppercase transitions using the existing generated naming contract.
 fn screaming_snake(name: &str) -> String {
     let mut out = String::new();
     for (idx, ch) in name.chars().enumerate() {
@@ -516,6 +598,7 @@ fn screaming_snake(name: &str) -> String {
     out
 }
 
+/// Converts a Prost method identifier into the stream type's name segment.
 fn upper_camel_from_snake(name: &str) -> String {
     let mut out = String::new();
     let mut upper = true;
@@ -534,12 +617,15 @@ fn upper_camel_from_snake(name: &str) -> String {
     out
 }
 
+/// Checks generated API shape and half-close failure handling.
 #[cfg(test)]
 mod tests {
     use std::fs;
 
+    /// Keeps every RPC shape and preserves responses after a failed half-close.
     #[test]
     fn configure_generates_starpc_service_glue() {
+        // Isolate generated sources from other generator runs.
         let root = std::env::temp_dir().join(format!(
             "starpc-codegen-test-{}-{}",
             std::process::id(),
@@ -550,10 +636,10 @@ mod tests {
         ));
         let proto_dir = root.join("proto");
         let out_dir = root.join("out");
-        let _ = fs::remove_dir_all(&root);
         fs::create_dir_all(&proto_dir).unwrap();
         fs::create_dir_all(&out_dir).unwrap();
 
+        // Exercise all call shapes in one service declaration.
         let proto = proto_dir.join("test.proto");
         fs::write(
             &proto,
@@ -574,10 +660,12 @@ message TestMsg {
         )
         .unwrap();
 
+        // Generate the real bindings and check their public surfaces.
         let mut config = super::configure();
         config.out_dir(&out_dir);
         config.compile_protos(&[proto], &[proto_dir]).unwrap();
 
+        // Verify the externally consumed client, server, and stream names.
         let generated = fs::read_to_string(out_dir.join("fixture.rs")).unwrap();
         assert!(generated.contains("pub const TEST_SERVICE_SERVICE_ID"));
         assert!(generated.contains("pub trait TestServiceClient"));
@@ -588,6 +676,13 @@ message TestMsg {
         assert!(generated.contains("TestServiceClientStreamStream"));
         assert!(generated.contains("TestServiceBidiStream"));
 
-        let _ = fs::remove_dir_all(&root);
+        // A failed half-close must not discard the stream; msg_recv reports
+        // the call's outcome.
+        assert!(!generated.contains("close_send().await?"));
+        assert!(generated.contains("let _ = stream.close_send().await;"));
+        assert!(generated.contains("let _ = self.stream.close_send().await;"));
+
+        // Release the disposable generated sources after checking their contract.
+        fs::remove_dir_all(&root).unwrap();
     }
 }