changeset 45382:eb55274d3650

rhg: add buffered stdout writing possibility

Improve batch stdout writing performance.
At some point line buffered output should be introduced.

Differential Revision: https://phab.mercurial-scm.org/D8866

--- a/rust/rhg/src/ui.rs	Mon Jul 20 18:14:52 2020 +0200
+++ b/rust/rhg/src/ui.rs	Wed Jul 29 10:26:17 2020 +0200
@@ -1,7 +1,11 @@
 use std::io;
-use std::io::Write;
+use std::io::{ErrorKind, Write};
 
-pub struct Ui {}
+#[derive(Debug)]
+pub struct Ui {
+    stdout: std::io::Stdout,
+    stderr: std::io::Stderr,
+}
 
 /// The kind of user interface error
 pub enum UiError {
@@ -14,20 +18,31 @@
 /// The commandline user interface
 impl Ui {
     pub fn new() -> Self {
-        Ui {}
+        Ui {
+            stdout: std::io::stdout(),
+            stderr: std::io::stderr(),
+        }
+    }
+
+    /// Returns a buffered handle on stdout for faster batch printing
+    /// operations.
+    pub fn stdout_buffer(&self) -> StdoutBuffer<std::io::StdoutLock> {
+        StdoutBuffer::new(self.stdout.lock())
     }
 
     /// Write bytes to stdout
     pub fn write_stdout(&self, bytes: &[u8]) -> Result<(), UiError> {
-        let mut stdout = io::stdout();
+        let mut stdout = self.stdout.lock();
 
         self.write_stream(&mut stdout, bytes)
-            .or_else(|e| self.into_stdout_error(e))?;
+            .or_else(|e| self.handle_stdout_error(e))?;
 
-        stdout.flush().or_else(|e| self.into_stdout_error(e))
+        stdout.flush().or_else(|e| self.handle_stdout_error(e))
     }
 
-    fn into_stdout_error(&self, error: io::Error) -> Result<(), UiError> {
+    /// Sometimes writing to stdout is not possible, try writing to stderr to
+    /// signal that failure, otherwise just bail.
+    fn handle_stdout_error(&self, error: io::Error) -> Result<(), UiError> {
         self.write_stderr(
             &[b"abort: ", error.to_string().as_bytes(), b"\n"].concat(),
         )?;
@@ -36,7 +51,7 @@
 
     /// Write bytes to stderr
     pub fn write_stderr(&self, bytes: &[u8]) -> Result<(), UiError> {
-        let mut stderr = io::stderr();
+        let mut stderr = self.stderr.lock();
 
         self.write_stream(&mut stderr, bytes)
             .or_else(|e| Err(UiError::StderrError(e)))?;
@@ -52,3 +67,43 @@
         stream.write_all(bytes)
     }
 }
+
+/// A buffered stdout writer for faster batch printing operations.
+pub struct StdoutBuffer<W: Write> {
+    buf: io::BufWriter<W>,
+}
+
+impl<W: Write> StdoutBuffer<W> {
+    pub fn new(writer: W) -> Self {
+        let buf = io::BufWriter::new(writer);
+        Self { buf }
+    }
+
+    /// Write bytes to stdout buffer
+    pub fn write_all(&mut self, bytes: &[u8]) -> Result<(), UiError> {
+        self.buf.write_all(bytes).or_else(|e| self.io_err(e))
+    }
+
+    /// Flush bytes to stdout
+    pub fn flush(&mut self) -> Result<(), UiError> {
+        self.buf.flush().or_else(|e| self.io_err(e))
+    }
+
+    fn io_err(&self, error: io::Error) -> Result<(), UiError> {
+        if let ErrorKind::BrokenPipe = error.kind() {
+            // This makes `| head` work for example
+            return Ok(());
+        }
+        let mut stderr = io::stderr();
+
+        stderr
+            .write_all(
+                &[b"abort: ", error.to_string().as_bytes(), b"\n"].concat(),
+            )
+            .map_err(|e| UiError::StderrError(e))?;
+
+        stderr.flush().map_err(|e| UiError::StderrError(e))?;
+
+        Err(UiError::StdoutError(error))
+    }
+}