mononoke/repo_listener: ignore errors in poll_shutdown

Summary:
If the socket is already shutdown for writes, then not being able to shut it
down is fine.

Reviewed By: ahornby

Differential Revision: D26052499

fbshipit-source-id: 2da6c34f657317419df00a0b7ba615e0eb351e0d

eden/mononoke/server/repo_listener/Cargo.toml

@@ -62,6 +62,7 @@ lazy_static = "1.0"
 maplit = "1.0"
 openssl = "0.10"
 percent-encoding = "2.1"
+pin-project = "0.4"
 sha-1 = "0.8"
 slog = { version = "2.5", features = ["max_level_debug"] }
 slog-kvfilter = "0.7"

eden/mononoke/server/repo_listener/src/connection_acceptor.rs

@@ -51,6 +51,7 @@ use stats::prelude::*;
 use crate::errors::ErrorKind;
 use crate::repo_handlers::RepoHandler;
 use crate::request_handler::{create_conn_logger, request_handler};
+use crate::stream::QuietShutdownStream;
 
 define_stats! {
     prefix = "mononoke.connection_acceptor";
@@ -226,6 +227,8 @@ async fn handle_connection(conn: PendingConnection, sock: TcpStream) -> Result<(
 
     let is_hgcli = ssl_socket.ssl().selected_alpn_protocol() == Some(alpn::HGCLI_ALPN.as_bytes());
 
+    let ssl_socket = QuietShutdownStream::new(ssl_socket);
+
     if is_hgcli {
         handle_hgcli(conn, ssl_socket)
             .await

eden/mononoke/server/repo_listener/src/lib.rs

@@ -16,6 +16,7 @@ mod netspeedtest;
 mod repo_handlers;
 mod request_handler;
 mod security_checker;
+mod stream;
 
 pub use crate::connection_acceptor::wait_for_connections_closed;
 

eden/mononoke/server/repo_listener/src/stream.rs

@@ -0,0 +1,77 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This software may be used and distributed according to the terms of the
+ * GNU General Public License version 2.
+ */
+
+use futures::{
+    ready,
+    task::{Context, Poll},
+};
+use pin_project::pin_project;
+use std::io::{Error, ErrorKind};
+use std::pin::Pin;
+use tokio::io::{AsyncRead, AsyncWrite};
+
+#[pin_project]
+pub struct QuietShutdownStream<T> {
+    #[pin]
+    inner: T,
+}
+
+impl<T> QuietShutdownStream<T> {
+    pub fn new(inner: T) -> Self {
+        Self { inner }
+    }
+}
+
+impl<T> AsyncRead for QuietShutdownStream<T>
+where
+    T: AsyncRead,
+{
+    #[inline]
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context,
+        buf: &mut [u8],
+    ) -> Poll<Result<usize, Error>> {
+        let this = self.project();
+        this.inner.poll_read(cx, buf)
+    }
+}
+
+impl<T> AsyncWrite for QuietShutdownStream<T>
+where
+    T: AsyncWrite,
+{
+    #[inline]
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context,
+        buf: &[u8],
+    ) -> Poll<Result<usize, Error>> {
+        let this = self.project();
+        this.inner.poll_write(cx, buf)
+    }
+
+    #[inline]
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Error>> {
+        let this = self.project();
+        this.inner.poll_flush(cx)
+    }
+
+    #[inline]
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Error>> {
+        // This is useful to wrap a SslStream. See here for why:
+        // https://github.com/sfackler/tokio-openssl/issues/27
+        let this = self.project();
+        let res = ready!(this.inner.poll_shutdown(cx));
+        let res = match res {
+            Ok(r) => Ok(r),
+            Err(e) if e.kind() == ErrorKind::NotConnected => Ok(()),
+            Err(e) => Err(e),
+        };
+        Poll::Ready(res)
+    }
+}