std: Fix partial writes in LineWriter

src/libstd/io/buffered.rs

@@ -652,6 +652,7 @@ impl<W> fmt::Display for IntoInnerError<W> {
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct LineWriter<W: Write> {
     inner: BufWriter<W>,
+    need_flush: bool,
 }
 
 impl<W: Write> LineWriter<W> {
@@ -692,7 +693,10 @@ impl<W: Write> LineWriter<W> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn with_capacity(cap: usize, inner: W) -> LineWriter<W> {
-        LineWriter { inner: BufWriter::with_capacity(cap, inner) }
+        LineWriter {
+            inner: BufWriter::with_capacity(cap, inner),
+            need_flush: false,
+        }
     }
 
     /// Gets a reference to the underlying writer.
@@ -759,28 +763,57 @@ impl<W: Write> LineWriter<W> {
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn into_inner(self) -> Result<W, IntoInnerError<LineWriter<W>>> {
         self.inner.into_inner().map_err(|IntoInnerError(buf, e)| {
-            IntoInnerError(LineWriter { inner: buf }, e)
+            IntoInnerError(LineWriter {
+                inner: buf,
+                need_flush: false,
+            }, e)
         })
     }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<W: Write> Write for LineWriter<W> {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        match memchr::memrchr(b'\n', buf) {
-            Some(i) => {
-                let n = self.inner.write(&buf[..i + 1])?;
-                if n != i + 1 || self.inner.flush().is_err() {
-                    // Do not return errors on partial writes.
-                    return Ok(n);
-                }
-                self.inner.write(&buf[i + 1..]).map(|i| n + i)
-            }
-            None => self.inner.write(buf),
+        if self.need_flush {
+            self.flush()?;
+        }
+
+        // Find the last newline character in the buffer provided. If found then
+        // we're going to write all the data up to that point and then flush,
+        // otherewise we just write the whole block to the underlying writer.
+        let i = match memchr::memrchr(b'\n', buf) {
+            Some(i) => i,
+            None => return self.inner.write(buf),
+        };
+
+
+        // Ok, we're going to write a partial amount of the data given first
+        // followed by flushing the newline. After we've successfully written
+        // some data then we *must* report that we wrote that data, so future
+        // errors are ignored. We set our internal `need_flush` flag, though, in
+        // case flushing fails and we need to try it first next time.
+        let n = self.inner.write(&buf[..i + 1])?;
+        self.need_flush = true;
+        if self.flush().is_err() || n != i + 1 {
+            return Ok(n)
+        }
+
+        // At this point we successfully wrote `i + 1` bytes and flushed it out,
+        // meaning that the entire line is now flushed out on the screen. While
+        // we can attempt to finish writing the rest of the data provided.
+        // Remember though that we ignore errors here as we've successfully
+        // written data, so we need to report that.
+        match self.inner.write(&buf[i + 1..]) {
+            Ok(i) => Ok(n + i),
+            Err(_) => Ok(n),
         }
     }
 
-    fn flush(&mut self) -> io::Result<()> { self.inner.flush() }
+    fn flush(&mut self) -> io::Result<()> {
+        self.inner.flush()?;
+        self.need_flush = false;
+        Ok(())
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -1153,4 +1186,44 @@ mod tests {
             BufWriter::new(io::sink())
         });
     }
+
+    struct AcceptOneThenFail {
+        written: bool,
+        flushed: bool,
+    }
+
+    impl Write for AcceptOneThenFail {
+        fn write(&mut self, data: &[u8]) -> io::Result<usize> {
+            if !self.written {
+                assert_eq!(data, b"a\nb\n");
+                self.written = true;
+                Ok(data.len())
+            } else {
+                Err(io::Error::new(io::ErrorKind::NotFound, "test"))
+            }
+        }
+
+        fn flush(&mut self) -> io::Result<()> {
+            assert!(self.written);
+            assert!(!self.flushed);
+            self.flushed = true;
+            Err(io::Error::new(io::ErrorKind::Other, "test"))
+        }
+    }
+
+    #[test]
+    fn erroneous_flush_retried() {
+        let a = AcceptOneThenFail {
+            written: false,
+            flushed: false,
+        };
+
+        let mut l = LineWriter::new(a);
+        assert_eq!(l.write(b"a\nb\na").unwrap(), 4);
+        assert!(l.get_ref().written);
+        assert!(l.get_ref().flushed);
+        l.get_mut().flushed = false;
+
+        assert_eq!(l.write(b"a").unwrap_err().kind(), io::ErrorKind::Other)
+    }
 }