From ad156b1fce82ed441079bfe7f590686805d229a3 Mon Sep 17 00:00:00 2001
From: k-nasa <htilcs1115@gmail.com>
Date: Sun, 13 Oct 2019 15:55:32 +0900
Subject: [PATCH] feat: Add BufWriter::into_inner flush

---
 src/io/buf_writer.rs | 28 ++++++++++++++++++++++++++--
 tests/buf_writer.rs  |  2 +-
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/src/io/buf_writer.rs b/src/io/buf_writer.rs
index 440e270..ba3bf74 100644
--- a/src/io/buf_writer.rs
+++ b/src/io/buf_writer.rs
@@ -4,6 +4,7 @@ use std::pin::Pin;
 use futures_core::ready;
 
 use crate::io::{self, Seek, SeekFrom, Write};
+use crate::io::write::WriteExt;
 use crate::task::{Context, Poll};
 
 const DEFAULT_CAPACITY: usize = 8 * 1024;
@@ -83,6 +84,9 @@ pub struct BufWriter<W> {
     written: usize,
 }
 
+#[derive(Debug)]
+pub struct IntoInnerError<W>(W, std::io::Error);
+
 impl<W: Write> BufWriter<W> {
     pin_utils::unsafe_pinned!(inner: W);
     pin_utils::unsafe_unpinned!(buf: Vec<u8>);
@@ -180,8 +184,28 @@ impl<W: Write> BufWriter<W> {
     /// For method that will attempt to write before returning the writer see [`poll_into_inner`]
     ///
     /// [`poll_into_inner`]: #method.poll_into_inner
-    pub fn into_inner(self) -> W {
-        self.inner
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # fn main() -> std::io::Result<()> { async_std::task::block_on(async {
+    /// use async_std::io::BufWriter;
+    /// use async_std::net::TcpStream;
+    ///
+    /// let buf_writer = BufWriter::new(TcpStream::connect("127.0.0.1:34251").await?);
+    ///
+    /// // unwrap the TcpStream and flush the buffer
+    /// let stream = buf_writer.into_inner().await.unwrap();
+    /// #
+    /// # Ok(()) }) }
+    /// ```
+    pub async fn into_inner(mut self) -> Result<W, IntoInnerError<BufWriter<W>>>
+    where
+        Self: Unpin
+    {
+        match self.flush().await {
+            Err(e) => Err(IntoInnerError(self, e)),
+            Ok(()) => Ok(self.inner),
+        }
     }
 
     /// Returns a reference to the internally buffered data.
diff --git a/tests/buf_writer.rs b/tests/buf_writer.rs
index 188cd8c..fa0e1ed 100644
--- a/tests/buf_writer.rs
+++ b/tests/buf_writer.rs
@@ -53,7 +53,7 @@ fn test_buffered_writer_inner_into_inner_does_not_flush() {
         let mut w = BufWriter::with_capacity(3, Vec::new());
         w.write(&[0, 1]).await.unwrap();
         assert_eq!(*w.get_ref(), []);
-        let w = w.into_inner();
+        let w = w.into_inner().await.unwrap();
         assert_eq!(w, []);
     })
 }