diff --git a/grpc/src/client/name_resolution/dns/test.rs b/grpc/src/client/name_resolution/dns/test.rs index f085f695a..f50616673 100644 --- a/grpc/src/client/name_resolution/dns/test.rs +++ b/grpc/src/client/name_resolution/dns/test.rs @@ -269,6 +269,21 @@ impl Runtime for FakeRuntime { ) -> Pin, String>> + Send>> { self.inner.tcp_stream(target, opts) } + + fn tcp_listener( + &self, + addr: std::net::SocketAddr, + ) -> Pin, String>> + Send>> { + self.inner.tcp_listener(addr) + } + + fn unix_listener( + &self, + path: std::path::PathBuf, + opts: rt::UnixSocketOptions, + ) -> Pin, String>> + Send>> { + self.inner.unix_listener(path, opts) + } } #[tokio::test] diff --git a/grpc/src/client/name_resolution/proxy_resolver.rs b/grpc/src/client/name_resolution/proxy_resolver.rs index 6e6488794..11ead3967 100644 --- a/grpc/src/client/name_resolution/proxy_resolver.rs +++ b/grpc/src/client/name_resolution/proxy_resolver.rs @@ -324,6 +324,21 @@ mod tests { ) -> Pin, String>> + Send>> { self.inner.tcp_stream(target, opts) } + + fn tcp_listener( + &self, + addr: std::net::SocketAddr, + ) -> rt::BoxFuture, String>> { + self.inner.tcp_listener(addr) + } + + fn unix_listener( + &self, + path: std::path::PathBuf, + opts: rt::UnixSocketOptions, + ) -> rt::BoxFuture, String>> { + self.inner.unix_listener(path, opts) + } } struct MockResolverBuilder {} diff --git a/grpc/src/credentials/dyn_wrapper.rs b/grpc/src/credentials/dyn_wrapper.rs deleted file mode 100644 index f6067c32e..000000000 --- a/grpc/src/credentials/dyn_wrapper.rs +++ /dev/null @@ -1,132 +0,0 @@ -/* - * - * Copyright 2026 gRPC authors. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to - * deal in the Software without restriction, including without limitation the - * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or - * sell copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING - * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS - * IN THE SOFTWARE. - * - */ - -use tonic::async_trait; - -use crate::credentials::ProtocolInfo; -use crate::credentials::ServerCredentials; -use crate::credentials::server::HandshakeOutput as ServerHandshakeOutput; -use crate::private; -use crate::rt::BoxEndpoint; -use crate::rt::GrpcEndpoint; -use crate::rt::GrpcRuntime; -use crate::send_future::SendFuture; - -// Bridge trait for type erasure. -#[async_trait] -pub(crate) trait DynServerCredentials: Send + Sync { - async fn dyn_accept( - &self, - source: BoxEndpoint, - runtime: GrpcRuntime, - ) -> Result, String>; - - fn info(&self) -> &ProtocolInfo; -} - -#[async_trait] -impl DynServerCredentials for T -where - T: ServerCredentials, - T::Output: GrpcEndpoint, -{ - async fn dyn_accept( - &self, - source: BoxEndpoint, - runtime: GrpcRuntime, - ) -> Result, String> { - let output = SendFuture::make_send(self.accept(source, runtime, private::Internal)).await?; - Ok(ServerHandshakeOutput { - endpoint: Box::new(output.endpoint), - security: output.security, - }) - } - - fn info(&self) -> &ProtocolInfo { - self.info() - } -} - -#[cfg(test)] -mod tests { - use tokio::io::AsyncReadExt; - use tokio::io::AsyncWriteExt; - use tokio::net::TcpListener; - use tokio::net::TcpStream; - - use super::*; - use crate::credentials::LocalServerCredentials; - use crate::credentials::SecurityLevel; - use crate::rt; - use crate::rt::AsyncIoAdapter; - use crate::rt::tokio::TokioIoStream; - - #[tokio::test] - async fn test_dyn_server_credential_dispatch() { - let creds = LocalServerCredentials::new(); - let dyn_creds: Box = Box::new(creds); - - let info = dyn_creds.info(); - assert_eq!(info.security_protocol, "local"); - - let addr = "127.0.0.1:0"; - let runtime = rt::default_runtime(); - let listener = TcpListener::bind(addr).await.unwrap(); - let server_addr = listener.local_addr().unwrap(); - - let client_handle = tokio::spawn(async move { - let mut stream = TcpStream::connect(server_addr).await.unwrap(); - let data = b"hello dynamic grpc server"; - stream.write_all(data).await.unwrap(); - - // Keep the connection alive for a bit so server can read - let mut buf = vec![0u8; 1]; - let _ = stream.read(&mut buf).await; - }); - - let (stream, _) = listener.accept().await.unwrap(); - let server_stream = TokioIoStream::new_from_tcp(stream).unwrap(); - - let result = dyn_creds - .dyn_accept(Box::new(server_stream) as Box, runtime) - .await; - - assert!(result.is_ok()); - let output = result.unwrap(); - let endpoint = output.endpoint; - let security_info = output.security; - - assert_eq!(security_info.security_protocol(), "local"); - assert_eq!(security_info.security_level(), SecurityLevel::NoSecurity); - - let mut buf = vec![0u8; 25]; - AsyncIoAdapter::new(endpoint) - .read_exact(&mut buf) - .await - .unwrap(); - assert_eq!(&buf[..], b"hello dynamic grpc server"); - - client_handle.abort(); - } -} diff --git a/grpc/src/credentials/local.rs b/grpc/src/credentials/local.rs index 0a8905b5c..1ec785671 100644 --- a/grpc/src/credentials/local.rs +++ b/grpc/src/credentials/local.rs @@ -45,7 +45,6 @@ use crate::credentials::server; use crate::credentials::server::ServerConnectionSecurityInfo; use crate::private; use crate::rt::BoxEndpoint; -use crate::rt::GrpcEndpoint; use crate::rt::GrpcRuntime; pub const PROTOCOL_NAME: &str = "local"; @@ -162,15 +161,14 @@ impl LocalServerCredentials { } } +#[async_trait] impl ServerCredentials for LocalServerCredentials { - type Output = I; - - async fn accept( + async fn accept( &self, - source: Input, + source: BoxEndpoint, _runtime: GrpcRuntime, _token: private::Internal, - ) -> Result>, String> { + ) -> Result { let security_level = security_level_for_endpoint(source.get_peer_address(), source.get_network_type())?; Ok(server::HandshakeOutput { @@ -204,7 +202,6 @@ mod test { use crate::credentials::common::Authority; use crate::rt; use crate::rt::AsyncIoAdapter; - use crate::rt::GrpcEndpoint; use crate::rt::TcpOptions; use crate::rt::tokio::TokioIoStream; @@ -321,7 +318,7 @@ mod test { let server_stream = TokioIoStream::new_from_tcp(stream).unwrap(); let output = creds - .accept(server_stream, runtime, private::Internal) + .accept(Box::new(server_stream), runtime, private::Internal) .await .unwrap(); let endpoint = output.endpoint; @@ -339,4 +336,56 @@ mod test { client_handle.abort(); } + + #[tokio::test] + async fn test_dyn_server_credential_dispatch() { + let creds = LocalServerCredentials::new(); + let dyn_creds: Arc = Arc::new(creds); + + let info = dyn_creds.info(); + assert_eq!(info.security_protocol, "local"); + + let addr = "127.0.0.1:0"; + let runtime = rt::default_runtime(); + let listener = TcpListener::bind(addr).await.unwrap(); + let server_addr = listener.local_addr().unwrap(); + + let client_handle = tokio::spawn(async move { + let mut stream = TcpStream::connect(server_addr).await.unwrap(); + let data = b"hello dynamic grpc server"; + stream.write_all(data).await.unwrap(); + + // Keep the connection alive for a bit so server can read + let mut buf = vec![0u8; 1]; + let _ = stream.read(&mut buf).await; + }); + + let (stream, _) = listener.accept().await.unwrap(); + let server_stream = TokioIoStream::new_from_tcp(stream).unwrap(); + + let result = dyn_creds + .accept( + Box::new(server_stream) as Box, + runtime, + private::Internal, + ) + .await; + + assert!(result.is_ok()); + let output = result.unwrap(); + let endpoint = output.endpoint; + let security_info = output.security; + + assert_eq!(security_info.security_protocol(), "local"); + assert_eq!(security_info.security_level(), SecurityLevel::NoSecurity); + + let mut buf = vec![0u8; 25]; + AsyncIoAdapter::new(endpoint) + .read_exact(&mut buf) + .await + .unwrap(); + assert_eq!(&buf[..], b"hello dynamic grpc server"); + + client_handle.abort(); + } } diff --git a/grpc/src/credentials/mod.rs b/grpc/src/credentials/mod.rs index cb5b48ee0..a01e690c6 100644 --- a/grpc/src/credentials/mod.rs +++ b/grpc/src/credentials/mod.rs @@ -37,7 +37,6 @@ pub mod call; pub(crate) mod client; -pub(crate) mod dyn_wrapper; mod local; #[cfg(feature = "tls-rustls")] pub mod rustls; @@ -56,7 +55,6 @@ use crate::credentials::client::HandshakeOutput; use crate::credentials::common::Authority; use crate::private; use crate::rt::BoxEndpoint; -use crate::rt::GrpcEndpoint; use crate::rt::GrpcRuntime; /// Client-side trait for all live gRPC wire protocols and supported transport @@ -100,11 +98,8 @@ pub trait ChannelCredentials: Send + Sync + 'static { /// Server-side trait for all live gRPC wire protocols and supported /// transport security protocols (e.g., TLS, ALTS). -#[trait_variant::make(Send)] -pub trait ServerCredentials: Sync + 'static { - #[doc(hidden)] - type Output; - +#[async_trait] +pub trait ServerCredentials: Send + Sync + 'static { /// Provides the ProtocolInfo of these credentials. fn info(&self) -> &ProtocolInfo; @@ -113,12 +108,12 @@ pub trait ServerCredentials: Sync + 'static { /// This method wraps the incoming raw `source` connection with the configured /// security protocol (e.g., TLS). #[doc(hidden)] - async fn accept( + async fn accept( &self, - source: Input, + source: BoxEndpoint, runtime: GrpcRuntime, token: private::Internal, - ) -> Result>, String>; + ) -> Result; } /// Defines the level of protection provided by an established connection. diff --git a/grpc/src/credentials/rustls/server/mod.rs b/grpc/src/credentials/rustls/server/mod.rs index d571545df..45fefa8f7 100644 --- a/grpc/src/credentials/rustls/server/mod.rs +++ b/grpc/src/credentials/rustls/server/mod.rs @@ -56,7 +56,7 @@ use crate::credentials::server::HandshakeOutput; use crate::credentials::server::ServerConnectionSecurityInfo; use crate::private; use crate::rt::AsyncIoAdapter; -use crate::rt::GrpcEndpoint; +use crate::rt::BoxEndpoint; use crate::rt::GrpcRuntime; #[cfg(test)] @@ -338,15 +338,14 @@ impl ProducesTickets for NoTicketer { } } +#[tonic::async_trait] impl ServerCredentials for RustlsServerCredendials { - type Output = TlsStream; - - async fn accept( + async fn accept( &self, - source: Input, + source: BoxEndpoint, _runtime: GrpcRuntime, _token: private::Internal, - ) -> Result>, String> { + ) -> Result { let input_io = AsyncIoAdapter::new(source); let tls_stream = self .acceptor @@ -366,7 +365,7 @@ impl ServerCredentials for RustlsServerCredendials { ); let endpoint = TlsStream::new(RustlsStream::Server(tls_stream)); Ok(HandshakeOutput { - endpoint, + endpoint: Box::new(endpoint), security: auth_info, }) } diff --git a/grpc/src/credentials/rustls/server/test.rs b/grpc/src/credentials/rustls/server/test.rs index 988845565..9c6dc2af9 100644 --- a/grpc/src/credentials/rustls/server/test.rs +++ b/grpc/src/credentials/rustls/server/test.rs @@ -74,7 +74,9 @@ async fn test_tls_server_handshake() { let server_task = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); let stream = TokioIoStream::new_from_tcp(stream).unwrap(); - let result = creds.accept(stream, runtime, private::Internal).await; + let result = creds + .accept(Box::new(stream), runtime, private::Internal) + .await; assert!( result.is_ok(), "Server handshake failed: {:?}", @@ -130,7 +132,9 @@ async fn test_tls_server_handshake_no_alpn() { let server_task = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); let stream = TokioIoStream::new_from_tcp(stream).unwrap(); - let result = creds.accept(stream, runtime, private::Internal).await; + let result = creds + .accept(Box::new(stream), runtime, private::Internal) + .await; assert!(result.is_err(), "Server handshake should have failed"); }); @@ -174,7 +178,9 @@ async fn test_tls_server_handshake_bad_alpn() { let (stream, _) = listener.accept().await.unwrap(); let stream = TokioIoStream::new_from_tcp(stream).unwrap(); let runtime = rt::default_runtime(); - let result = creds.accept(stream, runtime, private::Internal).await; + let result = creds + .accept(Box::new(stream), runtime, private::Internal) + .await; assert!(result.is_err(), "Server handshake should have failed"); }); @@ -212,7 +218,7 @@ async fn test_tls_handshake_alpn_h1_and_h2() { let stream = TokioIoStream::new_from_tcp(stream).unwrap(); let runtime = rt::default_runtime(); creds - .accept(stream, runtime, private::Internal) + .accept(Box::new(stream), runtime, private::Internal) .await .unwrap(); }); @@ -257,7 +263,9 @@ async fn test_tls_server_mtls_require_fail() { let server_task = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); let stream = TokioIoStream::new_from_tcp(stream).unwrap(); - let result = creds.accept(stream, runtime, private::Internal).await; + let result = creds + .accept(Box::new(stream), runtime, private::Internal) + .await; assert!(result.is_err(), "Handshake should fail without client cert"); }); @@ -309,7 +317,7 @@ async fn test_tls_server_mtls_success() { let (stream, _) = listener.accept().await.unwrap(); let stream = TokioIoStream::new_from_tcp(stream).unwrap(); let result = creds - .accept(stream, runtime, private::Internal) + .accept(Box::new(stream), runtime, private::Internal) .await .expect("Server handshake failed"); let mut stream = AsyncIoAdapter::new(result.endpoint); @@ -369,7 +377,7 @@ async fn test_tls_server_mtls_optional() { let (stream, _) = listener.accept().await.unwrap(); let stream = TokioIoStream::new_from_tcp(stream).unwrap(); let result = creds - .accept(stream, runtime, private::Internal) + .accept(Box::new(stream), runtime, private::Internal) .await .expect("Server handshake failed"); let mut stream = AsyncIoAdapter::new(result.endpoint); @@ -419,7 +427,7 @@ async fn test_tls_server_key_log() { let (stream, _) = listener.accept().await.unwrap(); let stream = TokioIoStream::new_from_tcp(stream).unwrap(); let result = creds - .accept(stream, runtime, private::Internal) + .accept(Box::new(stream), runtime, private::Internal) .await .expect("Server handshake failed"); let mut stream = AsyncIoAdapter::new(result.endpoint); @@ -473,7 +481,9 @@ async fn check_resumption_disabled(versions: Vec<&'static rustls::SupportedProto let (stream, _) = listener.accept().await.unwrap(); let stream = TokioIoStream::new_from_tcp(stream).unwrap(); let runtime = rt::default_runtime(); - let result = creds.accept(stream, runtime, private::Internal).await; + let result = creds + .accept(Box::new(stream), runtime, private::Internal) + .await; assert!(result.is_ok()); let stream = result.unwrap().endpoint; AsyncIoAdapter::new(stream) @@ -549,7 +559,9 @@ async fn test_tls_server_sni() { let (stream, _) = listener.accept().await.unwrap(); let stream = TokioIoStream::new_from_tcp(stream).unwrap(); let runtime = rt::default_runtime(); - let result = creds.accept(stream, runtime, private::Internal).await; + let result = creds + .accept(Box::new(stream), runtime, private::Internal) + .await; assert!( result.is_ok(), "Server handshake failed: {:?}", diff --git a/grpc/src/credentials/server.rs b/grpc/src/credentials/server.rs index 7eeb7858a..9c9591616 100644 --- a/grpc/src/credentials/server.rs +++ b/grpc/src/credentials/server.rs @@ -24,9 +24,10 @@ use crate::attributes::Attributes; use crate::credentials::SecurityLevel; +use crate::rt::BoxEndpoint; -pub struct HandshakeOutput { - pub endpoint: T, +pub struct HandshakeOutput { + pub endpoint: BoxEndpoint, pub security: ServerConnectionSecurityInfo, } diff --git a/grpc/src/inmemory/mod.rs b/grpc/src/inmemory/mod.rs index 4bba7357f..f9bc1ebc4 100644 --- a/grpc/src/inmemory/mod.rs +++ b/grpc/src/inmemory/mod.rs @@ -70,19 +70,25 @@ use crate::credentials::client::ChannelSecurityContext; use crate::credentials::client::ChannelSecurityInfo; use crate::credentials::common::Authority; use crate::rt::GrpcRuntime; -use crate::server::Call as ServerCall; -use crate::server::Listener as ServerListener; +use crate::server::BoxedRecvStream; +use crate::server::DynHandle; +use crate::server::GracefulConnection; +use crate::server::Listener; use crate::server::RecvStream as ServerRecvStream; use crate::server::ResponseStreamItem as ServerResponseStreamItem; use crate::server::SendOptions as ServerSendOptions; use crate::server::SendStream as ServerSendStream; +use crate::server::Transport as ServerTransport; + +use std::pin::Pin; +use std::task::{Context, Poll}; static LISTENERS: LazyLock>>> = LazyLock::new(|| Mutex::new(HashMap::new())); static NEXT_ID: AtomicU64 = AtomicU64::new(0); -struct InMemoryServerCall { +pub struct InMemoryServerCall { headers: RequestHeaders, req_rx: mpsc::UnboundedReceiver, resp_tx: mpsc::UnboundedSender, @@ -98,6 +104,25 @@ enum InMemoryResponseStreamItem { Headers(ResponseHeaders), Message(Box), } +struct InMemoryListenerAddress { + id: String, +} + +impl crate::rt::address::ListenerAddress for InMemoryListenerAddress { + fn network(&self) -> &str { + "inmemory" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +impl std::fmt::Display for InMemoryListenerAddress { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "inmemory:{}", self.id) + } +} #[derive(Clone)] pub struct InMemoryListener { @@ -166,27 +191,71 @@ impl InMemoryListener { pub async fn await_connection(&self) {} } -impl ServerListener for InMemoryListener { - type SendStream = InMemoryServerSendStream; - type RecvStream = InMemoryServerRecvStream; +impl Listener for InMemoryListener { + type Transport = InMemoryServerCall; - async fn accept(&self) -> Option> { + async fn accept(&self, _token: crate::private::Internal) -> Option { let mut r = self.inner.r.lock().await; tokio::select! { call = r.recv() => { - let call = call?; - Some(ServerCall { - headers: call.headers, - send: InMemoryServerSendStream { tx: call.resp_tx }, - recv: InMemoryServerRecvStream { rx: call.req_rx }, - trailers_tx: call.trailer_tx, - }) + call } _ = self.inner.close_notify.notified() => { None } } } + + fn local_addr(&self) -> Box { + Box::new(InMemoryListenerAddress { + id: self.inner.id.clone(), + }) + } +} + +/// A serving in-memory connection. +pub struct InMemoryServingConnection { + inner: Pin + Send>>, +} + +impl std::future::Future for InMemoryServingConnection { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + self.inner.as_mut().poll(cx) + } +} + +impl GracefulConnection for InMemoryServingConnection { + fn graceful_shutdown(self: Pin<&mut Self>, _token: crate::private::Internal) { + // No-op: each in-memory "connection" is a single RPC. + } +} + +impl ServerTransport for InMemoryServerCall { + type Connection = InMemoryServingConnection; + + fn serve( + self, + handler: Arc, + runtime: GrpcRuntime, + _token: crate::private::Internal, + ) -> InMemoryServingConnection { + let mut send = InMemoryServerSendStream { tx: self.resp_tx }; + let recv = BoxedRecvStream(Box::new(InMemoryServerRecvStream { rx: self.req_rx })); + let options = crate::client::CallOptions::default(); + let trailers_tx = self.trailer_tx; + + let inner = Box::pin(async move { + runtime.spawn(Box::pin(async move { + let trailers = handler + .dyn_handle(self.headers, options, &mut send, recv) + .await; + let _ = trailers_tx.send(trailers); + })); + }); + InMemoryServingConnection { inner } + } } pub struct InMemoryServerSendStream { @@ -491,4 +560,24 @@ mod tests { _ => panic!("expected trailers with error, got {:?}", item), } } + + #[test] + fn in_memory_listener_local_addr() { + use crate::server::Listener; + let listener = InMemoryListener::new(); + let addr = listener.local_addr(); + assert_eq!(addr.network(), "inmemory"); + let display = addr.to_string(); + assert!( + display.starts_with("inmemory:"), + "expected 'inmemory:' prefix, got: {display}" + ); + } + + #[test] + fn in_memory_listener_id_is_unique() { + let a = InMemoryListener::new(); + let b = InMemoryListener::new(); + assert_ne!(a.id(), b.id()); + } } diff --git a/grpc/src/rt/address.rs b/grpc/src/rt/address.rs new file mode 100644 index 000000000..54b3c9b20 --- /dev/null +++ b/grpc/src/rt/address.rs @@ -0,0 +1,153 @@ +/* + * + * Copyright 2025 gRPC authors. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + * + */ + +use std::any::Any; +use std::fmt; + +/// The address a [`Listener`](super::Listener) is bound to. +/// +/// All implementations provide a human-readable string via [`Display`]. +/// Callers needing structured data (e.g. `std::net::SocketAddr`) can +/// downcast via [`as_any`](ListenerAddress::as_any) and +/// [`downcast_ref`](Any::downcast_ref). +/// +/// This is analogous to Go's `net.Addr` interface. +pub trait ListenerAddress: fmt::Display + Send + Sync + 'static { + /// Returns the network type (e.g. `"tcp"`, `"unix"`, `"inmemory"`). + fn network(&self) -> &str; + + /// Enables downcasting to a concrete address type. + fn as_any(&self) -> &dyn Any; +} + +// --------------------------------------------------------------------------- +// impl ListenerAddress for std::net::SocketAddr +// --------------------------------------------------------------------------- + +impl ListenerAddress for std::net::SocketAddr { + fn network(&self) -> &str { + "tcp" + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +// --------------------------------------------------------------------------- +// UnixListenerAddress +// --------------------------------------------------------------------------- + +/// Address for a Unix socket listener. +pub(crate) struct UnixListenerAddress { + path: String, +} + +impl UnixListenerAddress { + /// Creates a new Unix listener address. + pub(crate) fn new(path: String) -> Self { + Self { path } + } + + /// Returns the socket path. + pub(crate) fn path(&self) -> &str { + &self.path + } +} + +impl ListenerAddress for UnixListenerAddress { + fn network(&self) -> &str { + "unix" + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +impl fmt::Display for UnixListenerAddress { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.path) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::SocketAddr; + + // -- SocketAddr as ListenerAddress -- + + #[test] + fn socket_addr_network_is_tcp() { + let addr: SocketAddr = "127.0.0.1:8080".parse().unwrap(); + assert_eq!(addr.network(), "tcp"); + } + + #[test] + fn socket_addr_display() { + let addr: SocketAddr = "127.0.0.1:8080".parse().unwrap(); + assert_eq!(addr.to_string(), "127.0.0.1:8080"); + } + + #[test] + fn socket_addr_downcast_via_as_any() { + let addr: SocketAddr = "[::1]:443".parse().unwrap(); + let any_ref = addr.as_any(); + let downcasted = any_ref + .downcast_ref::() + .expect("downcast failed"); + assert_eq!(*downcasted, addr); + } + + // -- UnixListenerAddress -- + + #[test] + fn unix_address_path() { + let addr = UnixListenerAddress::new("/tmp/grpc.sock".to_string()); + assert_eq!(addr.path(), "/tmp/grpc.sock"); + } + + #[test] + fn unix_address_network_is_unix() { + let addr = UnixListenerAddress::new("/tmp/grpc.sock".to_string()); + assert_eq!(addr.network(), "unix"); + } + + #[test] + fn unix_address_display() { + let addr = UnixListenerAddress::new("/var/run/server.sock".to_string()); + assert_eq!(addr.to_string(), "/var/run/server.sock"); + } + + #[test] + fn unix_address_downcast_via_as_any() { + let addr = UnixListenerAddress::new("/tmp/test.sock".to_string()); + let any_ref = addr.as_any(); + let downcasted = any_ref + .downcast_ref::() + .expect("downcast failed"); + assert_eq!(downcasted.path(), "/tmp/test.sock"); + } +} diff --git a/grpc/src/rt/mod.rs b/grpc/src/rt/mod.rs index fb1b3166f..6cc9a674b 100644 --- a/grpc/src/rt/mod.rs +++ b/grpc/src/rt/mod.rs @@ -40,6 +40,7 @@ use ::tokio::io::ReadBuf; use crate::private; +pub mod address; pub(crate) mod hyper_wrapper; #[cfg(feature = "_runtime-tokio")] pub(crate) mod tokio; @@ -48,6 +49,16 @@ pub type BoxFuture = Pin + Send>>; pub type BoxedTaskHandle = Box; pub type BoxEndpoint = Box; +/// A server-side listening socket that yields incoming connections. +#[tonic::async_trait] +pub(crate) trait EndpointListener: Send + Sync + 'static { + /// Accepts the next incoming connection. + async fn accept(&self) -> Result, String>; + + /// Returns the local address this listener is bound to. + fn local_addr(&self) -> Box; +} + /// An abstraction over an asynchronous runtime. /// /// The `Runtime` trait defines the core functionality required for @@ -86,6 +97,23 @@ pub trait Runtime: Send + Sync + Debug { Err("Unix sockets are not supported by this runtime on this platform".to_string()) }) } + + /// Binds a TCP listener to the given address. + #[doc(hidden)] + #[allow(private_interfaces)] + fn tcp_listener( + &self, + addr: SocketAddr, + ) -> BoxFuture, String>>; + + /// Binds a Unix listener to the given path. + #[doc(hidden)] + #[allow(private_interfaces)] + fn unix_listener( + &self, + path: PathBuf, + opts: UnixSocketOptions, + ) -> BoxFuture, String>>; } /// A future that resolves after a specified duration. @@ -325,6 +353,21 @@ impl Runtime for NoOpRuntime { ) -> Pin, String>> + Send>> { unimplemented!() } + + fn tcp_listener( + &self, + _addr: SocketAddr, + ) -> BoxFuture, String>> { + unimplemented!() + } + + fn unix_listener( + &self, + _path: PathBuf, + _opts: UnixSocketOptions, + ) -> BoxFuture, String>> { + unimplemented!() + } } pub(crate) fn default_runtime() -> GrpcRuntime { @@ -377,4 +420,19 @@ impl GrpcRuntime { ) -> BoxFuture, String>> { self.inner.unix_stream(path, opts) } + + pub(crate) fn tcp_listener( + &self, + addr: SocketAddr, + ) -> BoxFuture, String>> { + self.inner.tcp_listener(addr) + } + + pub(crate) fn unix_listener( + &self, + path: PathBuf, + opts: UnixSocketOptions, + ) -> BoxFuture, String>> { + self.inner.unix_listener(path, opts) + } } diff --git a/grpc/src/rt/tokio/mod.rs b/grpc/src/rt/tokio/mod.rs index ebde593d3..fa2e8e98f 100644 --- a/grpc/src/rt/tokio/mod.rs +++ b/grpc/src/rt/tokio/mod.rs @@ -156,6 +156,41 @@ impl Runtime for TokioRuntime { Ok(stream) }) } + + fn tcp_listener( + &self, + addr: SocketAddr, + ) -> BoxFuture, String>> { + Box::pin(async move { + let listener = tokio::net::TcpListener::bind(addr) + .await + .map_err(|e| e.to_string())?; + Ok(Box::new(TokioTcpListener { listener }) as Box) + }) + } + + #[cfg(unix)] + fn unix_listener( + &self, + path: std::path::PathBuf, + _opts: super::UnixSocketOptions, + ) -> BoxFuture, String>> { + Box::pin(async move { + let listener = tokio::net::UnixListener::bind(&path).map_err(|e| e.to_string())?; + Ok(Box::new(TokioUnixListener { listener }) as Box) + }) + } + + #[cfg(not(unix))] + fn unix_listener( + &self, + _path: std::path::PathBuf, + _opts: super::UnixSocketOptions, + ) -> BoxFuture, String>> { + Box::pin( + async move { Err("Unix listeners are not supported on this platform".to_string()) }, + ) + } } impl TokioDefaultDnsResolver { @@ -254,6 +289,67 @@ impl super::GrpcEndpoint for } } +// --------------------------------------------------------------------------- +// TokioTcpListener — EndpointListener for TCP +// --------------------------------------------------------------------------- + +/// Wraps `tokio::net::TcpListener` as an [`EndpointListener`](super::EndpointListener). +struct TokioTcpListener { + listener: tokio::net::TcpListener, +} + +#[tonic::async_trait] +impl super::EndpointListener for TokioTcpListener { + async fn accept(&self) -> Result, String> { + let (stream, _addr) = self.listener.accept().await.map_err(|e| e.to_string())?; + let io = TokioIoStream::new_from_tcp(stream)?; + Ok(Box::new(io)) + } + + fn local_addr(&self) -> Box { + let addr = self.listener.local_addr().expect("TCP listener has addr"); + Box::new(addr) + } +} + +// --------------------------------------------------------------------------- +// TokioUnixListener — EndpointListener for Unix sockets +// --------------------------------------------------------------------------- + +#[cfg(unix)] +struct TokioUnixListener { + listener: tokio::net::UnixListener, +} + +#[cfg(unix)] +#[tonic::async_trait] +impl super::EndpointListener for TokioUnixListener { + async fn accept(&self) -> Result, String> { + use crate::client::name_resolution::UNIX_NETWORK_TYPE; + + let (stream, _addr) = self.listener.accept().await.map_err(|e| e.to_string())?; + let peer_addr = stream.peer_addr().map_err(|e| e.to_string())?; + let local_addr = stream.local_addr().map_err(|e| e.to_string())?; + + let io: Box = Box::new(TokioIoStream { + peer_addr: format!("{peer_addr:?}").into_boxed_str(), + local_addr: format!("{local_addr:?}").into_boxed_str(), + network_type: UNIX_NETWORK_TYPE, + inner: stream, + }); + Ok(io) + } + + fn local_addr(&self) -> Box { + let path = self + .listener + .local_addr() + .map(|a| format!("{a:?}")) + .unwrap_or_default(); + Box::new(crate::rt::address::UnixListenerAddress::new(path)) + } +} + #[cfg(test)] mod tests { use super::DnsResolver; diff --git a/grpc/src/server/mod.rs b/grpc/src/server/mod.rs index c71eaa64c..79f987aa7 100644 --- a/grpc/src/server/mod.rs +++ b/grpc/src/server/mod.rs @@ -22,9 +22,10 @@ * */ +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; -use tokio::sync::oneshot; use tonic::async_trait; use crate::client::CallOptions; @@ -33,32 +34,107 @@ use crate::core::RequestHeaders; use crate::core::ResponseHeaders; use crate::core::SendMessage; use crate::core::Trailers; +use crate::rt::GrpcRuntime; pub(crate) mod interceptor; +/// A serving connection that supports graceful shutdown. +/// +/// Implements `Future` to drive the connection, and exposes +/// `graceful_shutdown()` to initiate a graceful drain. +#[doc(hidden)] +pub trait GracefulConnection: Future + Send + 'static { + /// Initiate graceful shutdown. + fn graceful_shutdown(self: Pin<&mut Self>, token: crate::private::Internal); +} + pub struct Server { handler: Option>, + runtime: GrpcRuntime, } -pub struct Call { - pub headers: RequestHeaders, - pub send: SS, - pub recv: RS, - pub trailers_tx: oneshot::Sender, +mod sealed { + pub trait Sealed {} + impl Sealed for crate::inmemory::InMemoryListener {} } +/// A bound listening socket that yields incoming connections. +/// +/// **Sealed** — external crates cannot implement this trait. #[trait_variant::make(Send)] -pub trait Listener { - type SendStream: SendStream + 'static; - type RecvStream: RecvStream + 'static; - async fn accept(&self) -> Option>; +pub trait Listener: sealed::Sealed { + /// The concrete transport type yielded by this listener. + type Transport: Transport + 'static; + + /// Accepts the next incoming connection. + #[doc(hidden)] + async fn accept(&self, token: crate::private::Internal) -> Option; + + /// Returns the address this listener is bound to. + fn local_addr(&self) -> Box; +} + +/// A connection accepted by a [`Listener`] that can serve RPCs. +#[doc(hidden)] +pub trait Transport: Send + 'static { + /// The serving connection type. + type Connection: GracefulConnection; + + /// Bind the handler and start serving. + /// Returns a [`GracefulConnection`] that drives the connection when polled. + fn serve( + self, + handler: Arc, + runtime: GrpcRuntime, + token: crate::private::Internal, + ) -> Self::Connection; +} + +/// Coordinates graceful shutdown across all spawned connections. +/// +/// Each connection is watched via `watch()`. When `shutdown()` is called, +/// all watched connections receive a `graceful_shutdown()` signal. +struct GracefulCoordinator { + tx: tokio::sync::watch::Sender<()>, +} + +impl GracefulCoordinator { + fn new() -> Self { + let (tx, _) = tokio::sync::watch::channel(()); + Self { tx } + } + + fn watch(&self, conn: C) -> impl Future + Send + 'static { + let mut rx = self.tx.subscribe(); + async move { + let mut conn = std::pin::pin!(conn); + loop { + tokio::select! { + _ = &mut conn => { break; } + _ = rx.changed() => { + conn.as_mut().graceful_shutdown(crate::private::Internal); + } + } + } + } + } + + async fn shutdown(self) { + let _ = self.tx.send(()); + self.tx.closed().await; + } } impl Server { + /// Creates a new server with no handler. pub fn new() -> Self { - Self { handler: None } + Self { + handler: None, + runtime: crate::rt::default_runtime(), + } } + /// Sets the RPC handler for this server. pub fn set_handler(&mut self, h: H) where H: Handle + Send + Sync + 'static, @@ -66,20 +142,64 @@ impl Server { self.handler = Some(Arc::new(h)) } - pub async fn serve(&self, l: &impl Listener) { - while let Some(call) = l.accept().await { - let mut send: Box = Box::new(call.send); - let recv = BoxedRecvStream(Box::new(call.recv)); - let options = CallOptions::default(); - let trailers_tx = call.trailers_tx; - let trailers = self - .handler - .as_ref() - .unwrap() - .dyn_handle(call.headers, options, &mut *send, recv) - .await; - let _ = trailers_tx.send(trailers); + /// Serves on the given listener until it stops producing connections. + /// + /// After the accept loop ends, waits for all in-flight connections to + /// drain before returning. Connections are not notified to shut down + /// gracefully; they run until completion or client disconnect. + // TODO: Consider returning a `ServeOutcome` enum (e.g. + // ShutdownSignal / ListenerClosed) so callers can distinguish + // why the server stopped. + pub async fn serve(&self, listener: &impl Listener) { + self.serve_with_shutdown(listener, std::future::pending()) + .await; + } + + /// Serves on the given listener until `signal` resolves, then drains + /// all in-flight connections before returning. + /// + /// When `signal` completes: + /// 1. The server stops accepting new connections. + /// 2. All open HTTP/2 connections receive a GOAWAY frame. + /// 3. In-flight RPCs continue to completion. + /// 4. This method returns once all connections are closed. + // TODO(sauravz): The shutdown signal is currently per-listener like tonic + // Other implementations have signal on the server level. + // We need to evaluate what's the correct behavior here. + // Per Listener is more flexible and easy to implement, but + // doesn't allign with other implementations. + pub(crate) async fn serve_with_shutdown( + &self, + listener: &impl Listener, + signal: impl Future, + ) { + let graceful = GracefulCoordinator::new(); + + tokio::pin!(signal); + + loop { + tokio::select! { + conn = listener.accept(crate::private::Internal) => { + let Some(connection) = conn else { break }; + let handler = match self.handler.as_ref() { + Some(h) => h.clone(), + None => continue, + }; + let serving = connection.serve( + handler, + self.runtime.clone(), + crate::private::Internal, + ); + let fut = graceful.watch(serving); + self.runtime.spawn(Box::pin(fut)); + } + _ = &mut signal => { + break; + } + } } + + graceful.shutdown().await; } } @@ -105,8 +225,9 @@ pub trait Handle: Send + Sync { ) -> Trailers; } +#[doc(hidden)] #[async_trait] -trait DynHandle: Send + Sync { +pub trait DynHandle: Send + Sync { async fn dyn_handle( &self, headers: RequestHeaders, @@ -140,7 +261,8 @@ impl DynHandle for T { // | // = note: `Box<(dyn server::DynRecvStream + '0)>` must implement `server::RecvStream`, for any lifetime `'0`... // = note: ...but `server::RecvStream` is actually implemented for the type `Box<(dyn server::DynRecvStream + 'static)>` -struct BoxedRecvStream(Box); +#[doc(hidden)] +pub struct BoxedRecvStream(pub Box); // Implement RecvStream for the wrapper instead of the Box directly impl RecvStream for BoxedRecvStream { @@ -160,6 +282,23 @@ pub enum ResponseStreamItem<'a> { Message(&'a dyn SendMessage), } +/// Bridges [`DynHandle`] (object-safe) back to [`Handle`] (generic), +/// enabling interceptor composition on top of a dynamic handler. +pub(crate) struct DynHandleWrapper(pub Arc); + +impl Handle for DynHandleWrapper { + async fn handle( + &self, + headers: RequestHeaders, + options: CallOptions, + tx: &mut impl SendStream, + rx: impl RecvStream + 'static, + ) -> Trailers { + self.0 + .dyn_handle(headers, options, tx, BoxedRecvStream(Box::new(rx))) + .await + } +} /// Represents the sending side of a server stream. See `ResponseStream` /// documentation for information about the different types of items and the /// order in which they must be sent. @@ -181,8 +320,9 @@ pub trait SendStream { ) -> Result<(), ()>; } +#[doc(hidden)] #[async_trait] -trait DynSendStream: Send { +pub trait DynSendStream: Send { async fn dyn_send<'a>( &mut self, item: ResponseStreamItem<'a>, @@ -250,8 +390,9 @@ pub trait RecvStream { async fn next(&mut self, msg: &mut dyn RecvMessage) -> Option>; } +#[doc(hidden)] #[async_trait] -trait DynRecvStream: Send { +pub trait DynRecvStream: Send { async fn dyn_next(&mut self, msg: &mut dyn RecvMessage) -> Option>; } @@ -267,3 +408,153 @@ impl<'a> RecvStream for Box { (**self).dyn_next(msg).await } } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::task::{Context, Poll}; + use tokio::sync::Notify; + + /// A mock connection whose completion is controlled by a [`Notify`], + /// and which records whether [`graceful_shutdown`] was called. + struct MockConnection { + shutdown_called: Arc, + finish: Arc, + } + + impl MockConnection { + fn new() -> (Self, Arc, Arc) { + let shutdown_called = Arc::new(AtomicBool::new(false)); + let finish = Arc::new(Notify::new()); + ( + Self { + shutdown_called: shutdown_called.clone(), + finish: finish.clone(), + }, + shutdown_called, + finish, + ) + } + } + + impl Future for MockConnection { + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + let notified = self.finish.notified(); + tokio::pin!(notified); + notified.poll(cx) + } + } + + impl GracefulConnection for MockConnection { + fn graceful_shutdown(self: Pin<&mut Self>, _token: crate::private::Internal) { + self.shutdown_called.store(true, Ordering::SeqCst); + self.finish.notify_one(); + } + } + + #[tokio::test] + async fn shutdown_signals_watched_connection() { + let coordinator = GracefulCoordinator::new(); + let (conn, shutdown_called, _finish) = MockConnection::new(); + + let watched = coordinator.watch(conn); + let handle = tokio::spawn(watched); + + tokio::task::yield_now().await; + + coordinator.shutdown().await; + + handle.await.unwrap(); + assert!(shutdown_called.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn connection_completes_before_shutdown() { + let coordinator = GracefulCoordinator::new(); + let (conn, shutdown_called, finish) = MockConnection::new(); + + let watched = coordinator.watch(conn); + let handle = tokio::spawn(watched); + + // Connection finishes on its own (e.g. client disconnected). + finish.notify_one(); + handle.await.unwrap(); + + // graceful_shutdown was never called. + assert!(!shutdown_called.load(Ordering::SeqCst)); + + // Shutdown still returns promptly (no connections to drain). + coordinator.shutdown().await; + } + + #[tokio::test] + async fn two_connections_both_drained() { + let coordinator = GracefulCoordinator::new(); + + let (conn_a, shutdown_a, _finish_a) = MockConnection::new(); + let (conn_b, shutdown_b, _finish_b) = MockConnection::new(); + + let handle_a = tokio::spawn(coordinator.watch(conn_a)); + let handle_b = tokio::spawn(coordinator.watch(conn_b)); + + tokio::task::yield_now().await; + coordinator.shutdown().await; + + handle_a.await.unwrap(); + handle_b.await.unwrap(); + assert!(shutdown_a.load(Ordering::SeqCst)); + assert!(shutdown_b.load(Ordering::SeqCst)); + } + + // -- Server + InMemoryListener integration tests -- + + #[tokio::test] + async fn server_stops_on_shutdown_signal() { + let listener = crate::inmemory::InMemoryListener::new(); + let server = Server::new(); + + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>(); + + let listener_clone = listener.clone(); + let server_handle = tokio::spawn(async move { + server + .serve_with_shutdown(&listener_clone, async { + let _ = shutdown_rx.await; + }) + .await; + }); + + // Fire shutdown immediately. + let _ = shutdown_tx.send(()); + + // Server should exit promptly. + tokio::time::timeout(std::time::Duration::from_secs(2), server_handle) + .await + .expect("server did not shut down within 2s") + .unwrap(); + } + + #[tokio::test] + async fn server_stops_when_listener_closes() { + let listener = crate::inmemory::InMemoryListener::new(); + let server = Server::new(); + + let listener_for_serve = listener.clone(); + let server_handle = tokio::spawn(async move { + // serve() runs until the listener stops producing connections. + server.serve(&listener_for_serve).await; + }); + + // Closing the listener makes accept() return None. + listener.close().await; + + tokio::time::timeout(std::time::Duration::from_secs(2), server_handle) + .await + .expect("server did not exit after listener close") + .unwrap(); + } +}