From 02dc9d9a3419618fc729542b45c96c32b0f178bb Mon Sep 17 00:00:00 2001
From: Siddharth Agarwal <sid0@fb.com>
Date: Fri, 7 Apr 2017 17:58:06 -0700
Subject: [PATCH] encoder: instead of using write_all, keep track of partial
 writes

Non-blocking IO can cause the write_all to partially succeed. To retry in
potential future attempts, keep track of how much was written.

This has no impact on blocking IO.
---
 src/stream/encoder.rs | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/stream/encoder.rs b/src/stream/encoder.rs
index bc12b3cd..4fd2c29d 100644
--- a/src/stream/encoder.rs
+++ b/src/stream/encoder.rs
@@ -205,9 +205,9 @@ impl<W: Write> Encoder<W> {
             // Need to flush?
             panic!("Need to flush, but I'm lazy.");
         }
+        self.offset = 0;
 
-        // Write the end out
-        self.writer.write_all(&self.buffer)?;
+        self.write_from_offset()?;
 
         // Return the writer, because why not
         Ok(self.writer)
@@ -217,6 +217,18 @@ impl<W: Write> Encoder<W> {
     pub fn recommended_input_size() -> usize {
         unsafe { zstd_sys::ZSTD_CStreamInSize() }
     }
+
+    /// write_all, except keep track of partial writes for non-blocking IO.
+    fn write_from_offset(&mut self) -> io::Result<()> {
+        while self.offset < self.buffer.len() {
+            match self.writer.write(&self.buffer[self.offset..]) {
+                Ok(n) => self.offset += n,
+                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {},
+                Err(e) => return Err(e),
+            }
+        }
+        Ok(())
+    }
 }
 
 impl<W: Write> Write for Encoder<W> {
@@ -284,8 +296,9 @@ impl<W: Write> Write for Encoder<W> {
             self.buffer.set_len(buffer.pos);
             let _ = parse_code(code)?;
         }
+        self.offset = 0;
 
-        self.writer.write_all(&self.buffer)?;
+        self.write_from_offset()?;
         Ok(())
     }
 }