Add AsyncWriteExt::shutdown (#1382)

3 files changed, 47 insertions, 0 deletions

diff --git a/tokio-io/src/io/async_write_ext.rs b/tokio-io/src/io/async_write_ext.rs
index dc7b4a86..84671a03 100644
--- a/tokio-io/src/io/async_write_ext.rs
+++ b/tokio-io/src/io/async_write_ext.rs
@@ -1,4 +1,5 @@
 use crate::io::flush::{flush, Flush};
+use crate::io::shutdown::{shutdown, Shutdown};
 use crate::io::write::{write, Write};
 use crate::io::write_all::{write_all, WriteAll};
 use crate::AsyncWrite;
@@ -28,6 +29,14 @@ pub trait AsyncWriteExt: AsyncWrite {
     {
         flush(self)
     }
+
+    /// Shutdown this writer.
+    fn shutdown(&mut self) -> Shutdown<'_, Self>
+    where
+        Self: Unpin,
+    {
+        shutdown(self)
+    }
 }
 
 impl<W: AsyncWrite + ?Sized> AsyncWriteExt for W {}
diff --git a/tokio-io/src/io/mod.rs b/tokio-io/src/io/mod.rs
index 34117780..0b0f24fa 100644
--- a/tokio-io/src/io/mod.rs
+++ b/tokio-io/src/io/mod.rs
@@ -48,6 +48,7 @@ mod read_line;
 mod read_to_end;
 mod read_to_string;
 mod read_until;
+mod shutdown;
 mod write;
 mod write_all;
 
diff --git a/tokio-io/src/io/shutdown.rs b/tokio-io/src/io/shutdown.rs
new file mode 100644
index 00000000..4d01c46a
--- /dev/null
+++ b/tokio-io/src/io/shutdown.rs
@@ -0,0 +1,37 @@
+use crate::AsyncWrite;
+use std::future::Future;
+use std::io;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+/// A future used to shutdown an I/O object.
+///
+/// Created by the [`AsyncWriteExt::shutdown`] function.
+///
+/// [`shutdown`]: fn.shutdown.html
+#[derive(Debug)]
+pub struct Shutdown<'a, A: ?Sized> {
+    a: &'a mut A,
+}
+
+/// Creates a future which will shutdown an I/O object.
+pub(super) fn shutdown<A>(a: &mut A) -> Shutdown<'_, A>
+where
+    A: AsyncWrite + Unpin + ?Sized,
+{
+    Shutdown { a }
+}
+
+impl<'a, A> Unpin for Shutdown<'a, A> where A: Unpin + ?Sized {}
+
+impl<A> Future for Shutdown<'_, A>
+where
+    A: AsyncWrite + Unpin + ?Sized,
+{
+    type Output = io::Result<()>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let me = &mut *self;
+        Pin::new(&mut *me.a).poll_shutdown(cx)
+    }
+}