From a66595bfe3c146daaa437bddd5ce3db4542b1bf6 Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Wed, 8 Jan 2020 11:00:57 -0500
Subject: [PATCH] feat(transport): Add `serve_with_incoming_shutdown` (#220)

---
 tonic/src/transport/server/mod.rs | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/tonic/src/transport/server/mod.rs b/tonic/src/transport/server/mod.rs
index 7df3123ce..9158545cc 100644
--- a/tonic/src/transport/server/mod.rs
+++ b/tonic/src/transport/server/mod.rs
@@ -388,12 +388,12 @@ where
     pub async fn serve_with_shutdown<F: Future<Output = ()>>(
         self,
         addr: SocketAddr,
-        f: F,
+        signal: F,
     ) -> Result<(), super::Error> {
         let incoming = TcpIncoming::new(addr, self.server.tcp_nodelay, self.server.tcp_keepalive)
             .map_err(super::Error::from_source)?;
         self.server
-            .serve_with_shutdown(self.routes, incoming, Some(f))
+            .serve_with_shutdown(self.routes, incoming, Some(signal))
             .await
     }
 
@@ -411,6 +411,28 @@ where
             .serve_with_shutdown::<_, _, future::Ready<()>, _, _>(self.routes, incoming, None)
             .await
     }
+
+    /// Consume this [`Server`] creating a future that will execute the server on
+    /// the provided incoming stream of `AsyncRead + AsyncWrite`. Similar to
+    /// `serve_with_shutdown` this method will also take a signal future to
+    /// gracefully shutdown the server.
+    ///
+    /// [`Server`]: struct.Server.html
+    pub async fn serve_with_incoming_shutdown<I, IO, IE, F>(
+        self,
+        incoming: I,
+        signal: F,
+    ) -> Result<(), super::Error>
+    where
+        I: Stream<Item = Result<IO, IE>>,
+        IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
+        IE: Into<crate::Error>,
+        F: Future<Output = ()>,
+    {
+        self.server
+            .serve_with_shutdown(self.routes, incoming, Some(signal))
+            .await
+    }
 }
 
 impl fmt::Debug for Server {