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/inmemory/mod.rs b/grpc/src/inmemory/mod.rs index 4bba7357f..f476ad6e8 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,69 @@ 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 { + let trailers = handler + .dyn_handle(self.headers, options, &mut send, recv) + .await; + let _ = trailers_tx.send(trailers); + }); + InMemoryServingConnection { inner } + } } pub struct InMemoryServerSendStream { @@ -491,4 +558,85 @@ 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()); + } + + #[tokio::test] + async fn inmemory_handler_cancelled_on_connection_drop() { + use crate::client::CallOptions; + use crate::core::{RequestHeaders, Trailers}; + use crate::server::{RecvStream, SendStream}; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + + let handler_started = Arc::new(AtomicBool::new(false)); + let handler_finished = Arc::new(AtomicBool::new(false)); + let hs = handler_started.clone(); + let hf = handler_finished.clone(); + + struct HangingHandler(Arc, Arc); + impl crate::server::Handle for HangingHandler { + async fn handle( + &self, + _: RequestHeaders, + _: CallOptions, + _: &mut impl SendStream, + _: impl RecvStream + 'static, + ) -> Trailers { + self.0.store(true, Ordering::SeqCst); + tokio::time::sleep(std::time::Duration::from_secs(60)).await; + self.1.store(true, Ordering::SeqCst); + Trailers::new(Ok(())) + } + } + + let (req_tx, req_rx) = mpsc::unbounded_channel(); + let (resp_tx, _resp_rx) = mpsc::unbounded_channel(); + let (trailer_tx, _trailer_rx) = oneshot::channel(); + + let transport = InMemoryServerCall { + headers: RequestHeaders::new(), + req_rx, + resp_tx, + trailer_tx, + }; + + let mut serving_conn = transport.serve( + Arc::new(HangingHandler(hs, hf)), + crate::rt::default_runtime(), + crate::private::Internal, + ); + + // Poll the serving connection briefly to start the handler. + tokio::select! { + _ = &mut serving_conn => {} + _ = tokio::time::sleep(std::time::Duration::from_millis(20)) => {} + } + assert!(handler_started.load(Ordering::SeqCst)); + + // Now DROP `serving_conn` (simulating what watch() does on `break`). + drop(serving_conn); + + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + // With the inline fix, dropping the connection future drops the handler future immediately -> handler never finishes. + assert!(!handler_finished.load(Ordering::SeqCst)); + } } 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..0bdd22375 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,125 @@ 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. +/// +/// # Contract +/// Implementors MUST drive all of the connection's work within this future +/// (or own the handles to any sub-tasks they spawn). Dropping the connection +/// future MUST cancel all associated work and release the socket. Detaching +/// work via a fire-and-forget spawn breaks force-shutdown and leaks tasks. +/// +/// Furthermore, `graceful_shutdown()` MUST be idempotent: calling it multiple +/// times on the same connection MUST be safe and act as a no-op after the +/// initial shutdown signal is sent. +#[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; } + res = rx.changed() => { + match res { + // Operator graceful signal: begin GOAWAY drain, + // keep polling the connection to completion. + Ok(()) => conn.as_mut().graceful_shutdown(crate::private::Internal), + // Sender dropped == serve future dropped (force + // shutdown): stop now. Dropping `conn` closes the + // socket and cancels in-flight work. + Err(_) => break, + } + } + } + } + } + } + + 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 +160,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 +243,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 +279,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 +300,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 +338,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 +408,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 +426,231 @@ 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(); + } + + #[tokio::test] + async fn dropped_sender_breaks_watch_loop() { + 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; + + // Drop the coordinator/Sender explicitly (simulating serve future being dropped). + drop(coordinator); + + // The watched task should exit promptly because `rx.changed()` returns `Err` -> `break`. + // If the old busy-loop bug existed, this `await` would time out / hang forever. + tokio::time::timeout(std::time::Duration::from_secs(1), handle) + .await + .expect("watched connection did not exit when coordinator was dropped") + .unwrap(); + + // `graceful_shutdown` should NOT have been called on the force-close path. + assert!(!shutdown_called.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn dropping_serve_future_force_closes_connection() { + let listener = crate::inmemory::InMemoryListener::new(); + let server = Server::new(); + + // A never-resolving signal future (we won't signal gracefully, we will drop the serve future directly). + let (_signal_tx, signal_rx) = tokio::sync::oneshot::channel::<()>(); + + let listener_clone = listener.clone(); + let (serve_started_tx, serve_started_rx) = tokio::sync::oneshot::channel::<()>(); + + // Run `serve_with_shutdown` inside a task we can abort/drop (or wrap with timeout). + let serve_handle = tokio::spawn(async move { + let serve_fut = server.serve_with_shutdown(&listener_clone, async { + let _ = signal_rx.await; + }); + let _ = serve_started_tx.send(()); + // Wrap in a short timeout so the future drops mid-flight while connections are active. + let _ = tokio::time::timeout(std::time::Duration::from_millis(50), serve_fut).await; + }); + + serve_started_rx.await.unwrap(); + + // When the 50ms timeout in `serve_handle` elapses, `serve_fut` is dropped -> `coordinator.tx` dropped -> connection self-terminates. + serve_handle.await.unwrap(); + } + + #[tokio::test] + async fn infinite_connection_graceful_shutdown_then_drop() { + let coordinator = GracefulCoordinator::new(); + // MockConnection never notifies finish on its own or when graceful_shutdown is called. + let (conn, shutdown_called, _finish) = MockConnection::new(); + + let watched = coordinator.watch(conn); + let handle = tokio::spawn(watched); + + tokio::task::yield_now().await; + + // 1. Initiate graceful shutdown. `tx.send(())` wakes `rx.changed()`, triggering `graceful_shutdown()`. + // But because `conn` never completes (`_finish` never notified), `coordinator.shutdown()` would hang if awaited. + let _ = coordinator.tx.send(()); + tokio::task::yield_now().await; + assert!(shutdown_called.load(Ordering::SeqCst)); + + // 2. Escalate to forceful shutdown by dropping the coordinator (simulating dropping the serve future). + drop(coordinator); + + // The watched task MUST self-terminate (`rx.changed()` returns `Err` -> `break`) despite being mid-drain. + tokio::time::timeout(std::time::Duration::from_secs(1), handle) + .await + .expect("infinite watched connection did not exit when coordinator was dropped after graceful shutdown") + .unwrap(); + } +}