Cache and buffer stdout per-task for printing

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.

2 files changed, 108 insertions, 20 deletions

diff --git a/src/libstd/rt/io/buffered.rs b/src/libstd/rt/io/buffered.rs
index 9dcb35c806f..47c8dbd35c6 100644
--- a/src/libstd/rt/io/buffered.rs
+++ b/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]);
+    }
 }
diff --git a/src/libstd/rt/io/stdio.rs b/src/libstd/rt/io/stdio.rs
index b922e6400cc..3f34d32b350 100644
--- a/src/libstd/rt/io/stdio.rs
+++ b/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