Improve handling of truncated streams

DisregardBrokenPipe needs to implement write_all() as well as write()
because the former treats a 0-byte write as an error, whereas DPB should
no-op in that scenario.

Note that this is fix is in main.rs and therefore _only_ fixes the CLI.
Along with the existing TODO on DPB it may be preferable to stop using
write_all() in lib.rs, but for now I'm erring on the side of letting
library users handle failures themselves.

Fixes #44

src/main.rs

@@ -24,6 +24,21 @@ impl Write for DisregardBrokenPipe {
         }
     }
 
+    // Custom implementation of write_all() that treats Ok(0) as success rather than an error as the
+    // default implementation does.
+    // TODO perhaps this should be inlined into maybe_tee() instead of calling write_all()
+    fn write_all(&mut self, mut buf: &[u8]) -> io::Result<()> {
+        while !buf.is_empty() {
+            match self.write(buf) {
+                Ok(0) => return Ok(()),
+                Ok(n) => buf = &buf[n..],
+                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {},
+                Err(e) => return Err(e),
+            }
+        }
+        Ok(())
+    }
+
     fn flush(&mut self) -> io::Result<()> {
         match self.0.flush() {
             Err(e) if e.kind() == std::io::ErrorKind::BrokenPipe => Ok(()),

tests/cli.rs

@@ -463,6 +463,34 @@ mod cli {
         assert_eq!(result.status, Some(0));
     }
 
+    #[test]
+    fn truncated_output() {
+        let dir = TestDir::temp();
+        let bytes = 1024*100; // 100KB is larger than the standard OS process buffer
+        // Write a large amount of data to stdout and close the process' stream without reading it;
+        // this should be supported silently, see https://github.com/dimo414/bkt/issues/44.
+        let script = format!(r#"printf '.%.0s' {{1..{0}}}"#, bytes);
+        let args = ["--", "bash", "-c", &script, "arg0"];
+        let mut cmd = bkt(dir.path("cache"));
+        let cmd = cmd.args(args).stdout(Stdio::piped()).stderr(Stdio::piped());
+
+        let mut child = cmd.spawn().unwrap();
+        // Read the beginning of stdout
+        // It's not strictly necessary to do this, in fact closing the stream without reading
+        // anything causes the error even for small outputs, but this seems like the more
+        // "interesting" case and it covers the read-nothing behavior too.
+        let mut buf = [0; 10];
+        child.stdout.as_mut().unwrap().read_exact(&mut buf).unwrap();
+        assert_eq!(buf, ['.' as u8; 10]);
+
+        std::mem::drop(child.stdout.take().unwrap()); // close stdout without reading further
+
+        let result: CmdResult = child.wait_with_output().unwrap().into();
+        assert_eq!(result.out, "");
+        assert_eq!(result.err, ""); // Unexpected error messages will show up in stderr
+        assert_eq!(result.status, Some(0));
+    }
+
     #[test]
     #[cfg(not(feature="debug"))]
     fn no_debug_output() {