auto merge of #10060 : alexcrichton/rust/cached-stdout, r=brson

Almost all languages provide some form of buffering of the stdout stream, and
this commit adds this feature for rust. A handle to stdout is lazily initialized
in the Task structure as a buffered owned Writer trait object. The buffer
behavior depends on where stdout is directed to. Like C, this line-buffers the
stream when the output goes to a terminal (flushes on newlines), and also like C
this uses a fixed-size buffer when output is not directed at a terminal.

We may decide the fixed-size buffering is overkill, but it certainly does reduce
write syscall counts when piping output elsewhere. This is a *huge* benefit to
any code using logging macros or the printing macros. Formatting emits calls to
`write` very frequently, and to have each of them backed by a write syscall was
very expensive.

In a local benchmark of printing 10000 lines of "what" to stdout, I got the
following timings:

  when |  terminal   |  redirected
----------|---------------|--------
before |  0.575s     |   0.525s
after  |  0.197s     |   0.013s
  C    |  0.019s     |   0.004s

I can also confirm that we're buffering the output appropriately in both
situtations. We're still far slower than C, but I believe much of that has to do
with the "homing" that all tasks due, we're still performing an order of
magnitude more write syscalls than C does.

Author: bors

src/libstd/rt/io/buffered.rs

@@ -221,17 +221,48 @@ impl<W: Writer> Writer for BufferedWriter<W> {
 }
 
 impl<W: Writer> Decorator<W> for BufferedWriter<W> {
-    fn inner(self) -> W {
-        self.inner
-    }
+    fn inner(self) -> W { self.inner }
+    fn inner_ref<'a>(&'a self) -> &'a W { &self.inner }
+    fn inner_mut_ref<'a>(&'a mut self) -> &'a mut W { &mut self.inner }
+}
 
-    fn inner_ref<'a>(&'a self) -> &'a W {
-        &self.inner
+/// Wraps a Writer and buffers output to it, flushing whenever a newline (0xa,
+/// '\n') is detected.
+///
+/// Note that this structure does NOT flush the output when dropped.
+pub struct LineBufferedWriter<W> {
+    priv inner: BufferedWriter<W>,
+}
+
+impl<W: Writer> LineBufferedWriter<W> {
+    /// Creates a new `LineBufferedWriter`
+    pub fn new(inner: W) -> LineBufferedWriter<W> {
+        // Lines typically aren't that long, don't use a giant buffer
+        LineBufferedWriter {
+            inner: BufferedWriter::with_capacity(1024, inner)
+        }
     }
+}
 
-    fn inner_mut_ref<'a>(&'a mut self) -> &'a mut W {
-        &mut self.inner
+impl<W: Writer> Writer for LineBufferedWriter<W> {
+    fn write(&mut self, buf: &[u8]) {
+        match buf.iter().position(|&b| b == '\n' as u8) {
+            Some(i) => {
+                self.inner.write(buf.slice_to(i + 1));
+                self.inner.flush();
+                self.inner.write(buf.slice_from(i + 1));
+            }
+            None => self.inner.write(buf),
+        }
     }
+
+    fn flush(&mut self) { self.inner.flush() }
+}
+
+impl<W: Writer> Decorator<W> for LineBufferedWriter<W> {
+    fn inner(self) -> W { self.inner.inner() }
+    fn inner_ref<'a>(&'a self) -> &'a W { self.inner.inner_ref() }
+    fn inner_mut_ref<'a>(&'a mut self) -> &'a mut W { self.inner.inner_mut_ref() }
 }
 
 struct InternalBufferedWriter<W>(BufferedWriter<W>);
@@ -413,4 +444,19 @@ mod test {
         assert_eq!(reader.read_until(8), Some(~[0]));
         assert_eq!(reader.read_until(9), None);
     }
+
+    #[test]
+    fn test_line_buffer() {
+        let mut writer = LineBufferedWriter::new(MemWriter::new());
+        writer.write([0]);
+        assert_eq!(*writer.inner_ref().inner_ref(), ~[]);
+        writer.write([1]);
+        assert_eq!(*writer.inner_ref().inner_ref(), ~[]);
+        writer.flush();
+        assert_eq!(*writer.inner_ref().inner_ref(), ~[0, 1]);
+        writer.write([0, '\n' as u8, 1]);
+        assert_eq!(*writer.inner_ref().inner_ref(), ~[0, 1, 0, '\n' as u8]);
+        writer.flush();
+        assert_eq!(*writer.inner_ref().inner_ref(), ~[0, 1, 0, '\n' as u8, 1]);
+    }
 }

src/libstd/rt/io/stdio.rs

@@ -30,6 +30,7 @@ use fmt;
 use libc;
 use option::{Option, Some, None};
 use result::{Ok, Err};
+use rt::io::buffered::{LineBufferedWriter, BufferedWriter};
 use rt::rtio::{IoFactory, RtioTTY, RtioFileStream, with_local_io,
                CloseAsynchronously};
 use super::{Reader, Writer, io_error, IoError, OtherIoError};
@@ -111,37 +112,78 @@ pub fn stderr() -> StdWriter {
     do src(libc::STDERR_FILENO, false) |src| { StdWriter { inner: src } }
 }
 
+/// Executes a closure with the local task's handle on stdout. By default, this
+/// stream is a buffering stream, so the handled yielded to the given closure
+/// can be used to flush the stdout stream (if necessary). The buffering used is
+/// line-buffering when stdout is attached to a terminal, and a fixed sized
+/// buffer if it is not attached to a terminal.
+///
+/// Note that handles generated via the `stdout()` function all output to the
+/// same stream, and output among all task may be interleaved as a result of
+/// this. This is provided to have access to the default stream for `print` and
+/// `println` (and the related macros) for this task.
+///
+/// Also note that logging macros do not use this stream. Using the logging
+/// macros will emit output to stderr.
+pub fn with_task_stdout(f: &fn(&mut Writer)) {
+    use rt::local::Local;
+    use rt::task::Task;
+
+    unsafe {
+        // Logging may require scheduling operations, so we can't remove the
+        // task from TLS right now, hence the unsafe borrow. Sad.
+        let task: *mut Task = Local::unsafe_borrow();
+
+        match (*task).stdout_handle {
+            Some(ref mut handle) => f(*handle),
+            None => {
+                let handle = stdout();
+                let mut handle = if handle.isatty() {
+                    ~LineBufferedWriter::new(handle) as ~Writer
+                } else {
+                    // The default capacity is very large, 64k, but this is just
+                    // a stdout stream, and possibly per task, so let's not make
+                    // this too expensive.
+                    ~BufferedWriter::with_capacity(4096, handle) as ~Writer
+                };
+                f(handle);
+                (*task).stdout_handle = Some(handle);
+            }
+        }
+    }
+}
+
 /// Prints a string to the stdout of the current process. No newline is emitted
 /// after the string is printed.
 pub fn print(s: &str) {
-    // XXX: need to see if not caching stdin() is the cause of performance
-    //      issues, it should be possible to cache a stdout handle in each Task
-    //      and then re-use that across calls to print/println. Note that the
-    //      resolution of this comment will affect all of the prints below as
-    //      well.
-    stdout().write(s.as_bytes());
+    do with_task_stdout |io| {
+        io.write(s.as_bytes());
+    }
 }
 
 /// Prints a string as a line. to the stdout of the current process. A literal
 /// `\n` character is printed to the console after the string.
 pub fn println(s: &str) {
-    let mut out = stdout();
-    out.write(s.as_bytes());
-    out.write(['\n' as u8]);
+    do with_task_stdout |io| {
+        io.write(s.as_bytes());
+        io.write(['\n' as u8]);
+    }
 }
 
 /// Similar to `print`, but takes a `fmt::Arguments` structure to be compatible
 /// with the `format_args!` macro.
 pub fn print_args(fmt: &fmt::Arguments) {
-    let mut out = stdout();
-    fmt::write(&mut out as &mut Writer, fmt);
+    do with_task_stdout |io| {
+        fmt::write(io, fmt);
+    }
 }
 
 /// Similar to `println`, but takes a `fmt::Arguments` structure to be
 /// compatible with the `format_args!` macro.
 pub fn println_args(fmt: &fmt::Arguments) {
-    let mut out = stdout();
-    fmt::writeln(&mut out as &mut Writer, fmt);
+    do with_task_stdout |io| {
+        fmt::writeln(io, fmt);
+    }
 }
 
 /// Representation of a reader of a standard input stream

src/libstd/rt/logging.rs

@@ -13,6 +13,8 @@ use from_str::from_str;
 use libc::exit;
 use option::{Some, None, Option};
 use rt::io;
+use rt::io::stdio::StdWriter;
+use rt::io::buffered::LineBufferedWriter;
 use rt::crate_map::{ModEntry, CrateMap, iter_crate_map, get_crate_map};
 use str::StrSlice;
 use u32;
@@ -170,7 +172,7 @@ pub trait Logger {
 /// This logger emits output to the stderr of the process, and contains a lazily
 /// initialized event-loop driven handle to the stream.
 pub struct StdErrLogger {
-    priv handle: Option<io::stdio::StdWriter>,
+    priv handle: Option<LineBufferedWriter<StdWriter>>,
 }
 
 impl StdErrLogger {
@@ -181,7 +183,7 @@ impl Logger for StdErrLogger {
     fn log(&mut self, args: &fmt::Arguments) {
         // First time logging? Get a handle to the stderr of this process.
         if self.handle.is_none() {
-            self.handle = Some(io::stderr());
+            self.handle = Some(LineBufferedWriter::new(io::stderr()));
         }
         fmt::writeln(self.handle.get_mut_ref() as &mut io::Writer, args);
     }

src/libstd/rt/task.rs

@@ -23,6 +23,7 @@ use option::{Option, Some, None};
 use rt::borrowck;
 use rt::borrowck::BorrowRecord;
 use rt::env;
+use rt::io::Writer;
 use rt::kill::Death;
 use rt::local::Local;
 use rt::logging::StdErrLogger;
@@ -56,7 +57,8 @@ pub struct Task {
     sched: Option<~Scheduler>,
     task_type: TaskType,
     // Dynamic borrowck debugging info
-    borrow_list: Option<~[BorrowRecord]>
+    borrow_list: Option<~[BorrowRecord]>,
+    stdout_handle: Option<~Writer>,
 }
 
 pub enum TaskType {
@@ -141,7 +143,8 @@ impl Task {
             name: None,
             sched: None,
             task_type: SchedTask,
-            borrow_list: None
+            borrow_list: None,
+            stdout_handle: None,
         }
     }
 
@@ -175,7 +178,8 @@ impl Task {
             coroutine: Some(Coroutine::new(stack_pool, stack_size, start)),
             sched: None,
             task_type: GreenTask(Some(home)),
-            borrow_list: None
+            borrow_list: None,
+            stdout_handle: None,
         }
     }
 
@@ -198,7 +202,8 @@ impl Task {
             coroutine: Some(Coroutine::new(stack_pool, stack_size, start)),
             sched: None,
             task_type: GreenTask(Some(home)),
-            borrow_list: None
+            borrow_list: None,
+            stdout_handle: None,
         }
     }
 
@@ -234,6 +239,7 @@ impl Task {
 
             // Run the task main function, then do some cleanup.
             do f.finally {
+
                 // First, destroy task-local storage. This may run user dtors.
                 //
                 // FIXME #8302: Dear diary. I'm so tired and confused.
@@ -257,6 +263,17 @@ impl Task {
 
                 // Destroy remaining boxes. Also may run user dtors.
                 unsafe { cleanup::annihilate(); }
+
+                // Finally flush and destroy any output handles which the task
+                // owns. There are no boxes here, and no user destructors should
+                // run after this any more.
+                match self.stdout_handle.take() {
+                    Some(handle) => {
+                        let mut handle = handle;
+                        handle.flush();
+                    }
+                    None => {}
+                }
             }
         }
 
@@ -331,7 +348,7 @@ impl Task {
 impl Drop for Task {
     fn drop(&mut self) {
         rtdebug!("called drop for a task: {}", borrow::to_uint(self));
-        rtassert!(self.destroyed)
+        rtassert!(self.destroyed);
     }
 }