From 49900d7a35a56f1c88c6f4d47d2ba5ac38dfa4ae Mon Sep 17 00:00:00 2001
From: Tim Kuehn <tikue@google.com>
Date: Fri, 8 Oct 2021 20:10:16 -0700
Subject: [PATCH] Close TcpStream when client disconnects

---
 tarpc/src/client.rs | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/tarpc/src/client.rs b/tarpc/src/client.rs
index f71f351..9f82cac 100644
--- a/tarpc/src/client.rs
+++ b/tarpc/src/client.rs
@@ -291,6 +291,9 @@ where
     /// Could not flush the transport.
     #[error("could not flush the transport")]
     Flush(#[source] E),
+    /// Could not close the write end of the transport.
+    #[error("could not close the write end of the transport")]
+    Close(#[source] E),
     /// Could not poll expired requests.
     #[error("could not poll expired requests")]
     Timer(#[source] tokio::time::error::Error),
@@ -335,6 +338,15 @@ where
             .map_err(ChannelError::Flush)
     }
 
+    fn poll_close<'a>(
+        self: &'a mut Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), ChannelError<C::Error>>> {
+        self.transport_pin_mut()
+            .poll_close(cx)
+            .map_err(ChannelError::Close)
+    }
+
     fn canceled_requests_mut<'a>(self: &'a mut Pin<&mut Self>) -> &'a mut CanceledRequests {
         self.as_mut().project().canceled_requests
     }
@@ -394,7 +406,7 @@ where
 
         match (pending_requests_status, canceled_requests_status) {
             (ReceiverStatus::Closed, ReceiverStatus::Closed) => {
-                ready!(self.poll_flush(cx)?);
+                ready!(self.poll_close(cx)?);
                 Poll::Ready(None)
             }
             (ReceiverStatus::Pending, _) | (_, ReceiverStatus::Pending) => {