Shut down the socket on close()

Author: Stjepan Glavina

src/lib.rs

@@ -623,8 +623,8 @@ impl<T: Write> AsyncWrite for Async<T> {
         poll_once(cx, self.write_with_mut(|io| io.flush()))
     }
 
-    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
-        self.poll_flush(cx)
+    fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Poll::Ready(self.source.shutdown_write())
     }
 }
 
@@ -652,8 +652,8 @@ where
         poll_once(cx, self.write_with(|io| (&*io).flush()))
     }
 
-    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
-        self.poll_flush(cx)
+    fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Poll::Ready(self.source.shutdown_write())
     }
 }
 

src/parking.rs

@@ -7,11 +7,12 @@
 use std::collections::BTreeMap;
 use std::fmt;
 use std::io;
-use std::mem;
+use std::mem::{self, ManuallyDrop};
+use std::net::{Shutdown, TcpStream};
 #[cfg(unix)]
-use std::os::unix::io::RawFd;
+use std::os::unix::io::{FromRawFd, RawFd};
 #[cfg(windows)]
-use std::os::windows::io::RawSocket;
+use std::os::windows::io::{FromRawSocket, RawSocket};
 use std::panic;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Condvar, Mutex, MutexGuard};
@@ -739,4 +740,21 @@ impl Source {
         })
         .await
     }
+
+    /// Shuts down the write side of the socket.
+    ///
+    /// If this source is not a socket, the `shutdown` syscall error is ignored.
+    pub(crate) fn shutdown_write(&self) -> io::Result<()> {
+        // This may not be a TCP stream, but that's okay - all we do is call `shutdown()` on it.
+        #[cfg(unix)]
+        let stream = unsafe { ManuallyDrop::new(TcpStream::from_raw_fd(self.raw)) };
+        #[cfg(windows)]
+        let stream = unsafe { ManuallyDrop::new(TcpStream::from_raw_socket(self.raw)) };
+
+        // The only actual error may be ENOTCONN.
+        match stream.shutdown(Shutdown::Write) {
+            Err(err) if err.kind() == io::ErrorKind::NotConnected => Err(err),
+            _ => Ok(()),
+        }
+    }
 }

tests/async.rs

@@ -9,7 +9,7 @@ use std::time::Duration;
 
 use async_io::{Async, Timer};
 use blocking::block_on;
-use futures::{AsyncReadExt, AsyncWriteExt, StreamExt};
+use futures::{future, AsyncReadExt, AsyncWriteExt, StreamExt};
 #[cfg(unix)]
 use tempfile::tempdir;
 
@@ -337,3 +337,15 @@ fn tcp_duplex() -> io::Result<()> {
         Ok(())
     })
 }
+
+#[test]
+fn close() -> io::Result<()> {
+    block_on(async {
+        let (mut reader, mut writer) = Async::<UnixStream>::pair()?;
+        let mut buf = Vec::new();
+
+        // The writer must be closed in order for `read_to_end()` to finish.
+        future::try_join(reader.read_to_end(&mut buf), writer.close()).await?;
+        Ok(())
+    })
+}