feat(grpc): add server-side Listener and Transport traits

grpc/src/client/name_resolution/dns/test.rs

@@ -269,6 +269,21 @@ impl Runtime for FakeRuntime {
     ) -> Pin<Box<dyn Future<Output = Result<Box<dyn GrpcEndpoint>, String>> + Send>> {
         self.inner.tcp_stream(target, opts)
     }
+
+    fn tcp_listener(
+        &self,
+        addr: std::net::SocketAddr,
+    ) -> Pin<Box<dyn Future<Output = Result<Box<dyn rt::EndpointListener>, String>> + Send>> {
+        self.inner.tcp_listener(addr)
+    }
+
+    fn unix_listener(
+        &self,
+        path: std::path::PathBuf,
+        opts: rt::UnixSocketOptions,
+    ) -> Pin<Box<dyn Future<Output = Result<Box<dyn rt::EndpointListener>, String>> + Send>> {
+        self.inner.unix_listener(path, opts)
+    }
 }
 
 #[tokio::test]

grpc/src/client/name_resolution/proxy_resolver.rs

@@ -324,6 +324,21 @@ mod tests {
         ) -> Pin<Box<dyn Future<Output = Result<Box<dyn GrpcEndpoint>, String>> + Send>> {
             self.inner.tcp_stream(target, opts)
         }
+
+        fn tcp_listener(
+            &self,
+            addr: std::net::SocketAddr,
+        ) -> rt::BoxFuture<Result<Box<dyn rt::EndpointListener>, String>> {
+            self.inner.tcp_listener(addr)
+        }
+
+        fn unix_listener(
+            &self,
+            path: std::path::PathBuf,
+            opts: rt::UnixSocketOptions,
+        ) -> rt::BoxFuture<Result<Box<dyn rt::EndpointListener>, String>> {
+            self.inner.unix_listener(path, opts)
+        }
     }
 
     struct MockResolverBuilder {}

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<Mutex<HashMap<String, mpsc::Sender<InMemoryServerCall>>>> =
     LazyLock::new(|| Mutex::new(HashMap::new()));
 
 static NEXT_ID: AtomicU64 = AtomicU64::new(0);
 
-struct InMemoryServerCall {
+pub struct InMemoryServerCall {
     headers: RequestHeaders,
     req_rx: mpsc::UnboundedReceiver<InMemoryRequestStreamItem>,
     resp_tx: mpsc::UnboundedSender<InMemoryResponseStreamItem>,
@@ -98,6 +104,25 @@ enum InMemoryResponseStreamItem {
     Headers(ResponseHeaders),
     Message(Box<dyn Buf + Send + Sync>),
 }
+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<ServerCall<Self::SendStream, Self::RecvStream>> {
+    async fn accept(&self, _token: crate::private::Internal) -> Option<Self::Transport> {
         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<dyn crate::rt::address::ListenerAddress> {
+        Box::new(InMemoryListenerAddress {
+            id: self.inner.id.clone(),
+        })
+    }
+}
+
+/// A serving in-memory connection.
+pub struct InMemoryServingConnection {
+    inner: Pin<Box<dyn std::future::Future<Output = ()> + 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<dyn DynHandle>,
+        _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<AtomicBool>, Arc<AtomicBool>);
+        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));
+    }
 }