Support arbitrary stdout/stderr/logger handles

This will allow capturing of common things like logging messages, stdout prints
(using stdio println), and failure messages (printed to stderr).  Any new prints
added to libstd should be funneled through these task handles to allow capture
as well.

Additionally, this commit redirects logging back through a `Logger` trait so the
log level can be usefully consumed by an arbitrary logger.

This commit also introduces methods to set the task-local stdout handles:

* std::io::stdio::set_stdout
* std::io::stdio::set_stderr
* std::io::logging::set_logger

These methods all return the previous logger just in case it needs to be used
for inspection.

I plan on using this infrastructure for extra::test soon, but we don't quite
have the primitives that I'd like to use for it, so it doesn't migrate
extra::test at this time.

Closes #6369

Author: alexcrichton

src/libgreen/task.rs

@@ -118,13 +118,17 @@ impl GreenTask {
                      f: proc()) -> ~GreenTask {
         let TaskOpts {
             watched: _watched,
-            notify_chan, name, stack_size
+            notify_chan, name, stack_size,
+            stderr, stdout, logger,
         } = opts;
 
         let mut green = GreenTask::new(pool, stack_size, f);
         {
             let task = green.task.get_mut_ref();
             task.name = name;
+            task.logger = logger;
+            task.stderr = stderr;
+            task.stdout = stdout;
             match notify_chan {
                 Some(chan) => {
                     let on_exit = proc(task_result) { chan.send(task_result) };

src/libnative/lib.rs

@@ -34,7 +34,7 @@ pub mod io;
 pub mod task;
 
 // XXX: this should not exist here
-#[cfg(stage0)]
+#[cfg(stage0, nativestart)]
 #[lang = "start"]
 pub fn lang_start(main: *u8, argc: int, argv: **u8) -> int {
     use std::cast;

src/libnative/task.rs

@@ -55,11 +55,15 @@ pub fn spawn(f: proc()) {
 pub fn spawn_opts(opts: TaskOpts, f: proc()) {
     let TaskOpts {
         watched: _watched,
-        notify_chan, name, stack_size
+        notify_chan, name, stack_size,
+        logger, stderr, stdout,
     } = opts;
 
     let mut task = ~Task::new();
     task.name = name;
+    task.logger = logger;
+    task.stderr = stderr;
+    task.stdout = stdout;
     match notify_chan {
         Some(chan) => {
             let on_exit = proc(task_result) { chan.send(task_result) };

src/libstd/io/stdio.rs

@@ -32,8 +32,12 @@ use io::{Reader, Writer, io_error, IoError, OtherIoError,
          standard_error, EndOfFile};
 use libc;
 use option::{Option, Some, None};
+use prelude::drop;
 use result::{Ok, Err};
+use rt::local::Local;
 use rt::rtio::{DontClose, IoFactory, LocalIo, RtioFileStream, RtioTTY};
+use rt::task::Task;
+use util;
 
 // And so begins the tale of acquiring a uv handle to a stdio stream on all
 // platforms in all situations. Our story begins by splitting the world into two
@@ -101,6 +105,44 @@ pub fn stderr() -> StdWriter {
     src(libc::STDERR_FILENO, false, |src| StdWriter { inner: src })
 }
 
+fn reset_helper(w: ~Writer,
+                f: |&mut Task, ~Writer| -> Option<~Writer>) -> Option<~Writer> {
+    let mut t = Local::borrow(None::<Task>);
+    // Be sure to flush any pending output from the writer
+    match f(t.get(), w) {
+        Some(mut w) => {
+            drop(t);
+            w.flush();
+            Some(w)
+        }
+        None => None
+    }
+}
+
+/// Resets the task-local stdout handle to the specified writer
+///
+/// This will replace the current task's stdout handle, returning the old
+/// handle. All future calls to `print` and friends will emit their output to
+/// this specified handle.
+///
+/// Note that this does not need to be called for all new tasks; the default
+/// output handle is to the process's stdout stream.
+pub fn set_stdout(stdout: ~Writer) -> Option<~Writer> {
+    reset_helper(stdout, |t, w| util::replace(&mut t.stdout, Some(w)))
+}
+
+/// Resets the task-local stderr handle to the specified writer
+///
+/// This will replace the current task's stderr handle, returning the old
+/// handle. Currently, the stderr handle is used for printing failure messages
+/// during task failure.
+///
+/// Note that this does not need to be called for all new tasks; the default
+/// output handle is to the process's stderr stream.
+pub fn set_stderr(stderr: ~Writer) -> Option<~Writer> {
+    reset_helper(stderr, |t, w| util::replace(&mut t.stderr, Some(w)))
+}
+
 // Helper to access the local task's stdout handle
 //
 // Note that this is not a safe function to expose because you can create an
@@ -112,38 +154,49 @@ pub fn stderr() -> StdWriter {
 //      })
 //  })
 fn with_task_stdout(f: |&mut Writer|) {
-    use rt::local::Local;
-    use rt::task::Task;
-
-    unsafe {
-        let task: Option<*mut Task> = Local::try_unsafe_borrow();
-        match task {
-            Some(task) => {
-                match (*task).stdout_handle {
-                    Some(ref mut handle) => f(*handle),
-                    None => {
-                        let handle = ~LineBufferedWriter::new(stdout());
-                        let mut handle = handle as ~Writer;
-                        f(handle);
-                        (*task).stdout_handle = Some(handle);
-                    }
-                }
+    let task: Option<~Task> = Local::try_take();
+    match task {
+        Some(mut task) => {
+            // Printing may run arbitrary code, so ensure that the task is in
+            // TLS to allow all std services. Note that this means a print while
+            // printing won't use the task's normal stdout handle, but this is
+            // necessary to ensure safety (no aliasing).
+            let mut my_stdout = task.stdout.take();
+            Local::put(task);
+
+            if my_stdout.is_none() {
+                my_stdout = Some(~LineBufferedWriter::new(stdout()) as ~Writer);
             }
+            f(*my_stdout.get_mut_ref());
+
+            // Note that we need to be careful when putting the stdout handle
+            // back into the task. If the handle was set to `Some` while
+            // printing, then we can run aribitrary code when destroying the
+            // previous handle. This means that the local task needs to be in
+            // TLS while we do this.
+            //
+            // To protect against this, we do a little dance in which we
+            // temporarily take the task, swap the handles, put the task in TLS,
+            // and only then drop the previous handle.
+            let mut t = Local::borrow(None::<Task>);
+            let prev = util::replace(&mut t.get().stdout, my_stdout);
+            drop(t);
+            drop(prev);
+        }
 
-            None => {
-                struct Stdout;
-                impl Writer for Stdout {
-                    fn write(&mut self, data: &[u8]) {
-                        unsafe {
-                            libc::write(libc::STDOUT_FILENO,
-                                        data.as_ptr() as *libc::c_void,
-                                        data.len() as libc::size_t);
-                        }
+        None => {
+            struct Stdout;
+            impl Writer for Stdout {
+                fn write(&mut self, data: &[u8]) {
+                    unsafe {
+                        libc::write(libc::STDOUT_FILENO,
+                                    data.as_ptr() as *libc::c_void,
+                                    data.len() as libc::size_t);
                     }
                 }
-                let mut io = Stdout;
-                f(&mut io as &mut Writer);
             }
+            let mut io = Stdout;
+            f(&mut io as &mut Writer);
         }
     }
 }
@@ -313,4 +366,29 @@ mod tests {
         stdout();
         stderr();
     })
+
+    iotest!(fn capture_stdout() {
+        use io::comm_adapters::{PortReader, ChanWriter};
+
+        let (p, c) = Chan::new();
+        let (mut r, w) = (PortReader::new(p), ChanWriter::new(c));
+        do spawn {
+            set_stdout(~w as ~Writer);
+            println!("hello!");
+        }
+        assert_eq!(r.read_to_str(), ~"hello!\n");
+    })
+
+    iotest!(fn capture_stderr() {
+        use io::comm_adapters::{PortReader, ChanWriter};
+
+        let (p, c) = Chan::new();
+        let (mut r, w) = (PortReader::new(p), ChanWriter::new(c));
+        do spawn {
+            set_stderr(~w as ~Writer);
+            fail!("my special message");
+        }
+        let s = r.read_to_str();
+        assert!(s.contains("my special message"));
+    })
 }

src/libstd/logging.rs

@@ -96,10 +96,15 @@ start, print out all modules registered for logging, and then exit.
 */
 
 use fmt;
-use option::*;
+use io::buffered::LineBufferedWriter;
+use io;
+use io::Writer;
+use ops::Drop;
+use option::{Some, None, Option};
+use prelude::drop;
 use rt::local::Local;
-use rt::logging::{Logger, StdErrLogger};
 use rt::task::Task;
+use util;
 
 /// Debug log level
 pub static DEBUG: u32 = 4;
@@ -110,24 +115,65 @@ pub static WARN: u32 = 2;
 /// Error log level
 pub static ERROR: u32 = 1;
 
+/// A trait used to represent an interface to a task-local logger. Each task
+/// can have its own custom logger which can respond to logging messages
+/// however it likes.
+pub trait Logger {
+    /// Logs a single message described by the `args` structure. The level is
+    /// provided in case you want to do things like color the message, etc.
+    fn log(&mut self, level: u32, args: &fmt::Arguments);
+}
+
+struct DefaultLogger {
+    handle: LineBufferedWriter<io::stdio::StdWriter>,
+}
+
+impl Logger for DefaultLogger {
+    // by default, just ignore the level
+    fn log(&mut self, _level: u32, args: &fmt::Arguments) {
+        fmt::writeln(&mut self.handle, args);
+    }
+}
+
+impl Drop for DefaultLogger {
+    fn drop(&mut self) {
+        self.handle.flush();
+    }
+}
+
 /// This function is called directly by the compiler when using the logging
 /// macros. This function does not take into account whether the log level
 /// specified is active or not, it will always log something if this method is
 /// called.
 ///
 /// It is not recommended to call this function directly, rather it should be
 /// invoked through the logging family of macros.
-pub fn log(_level: u32, args: &fmt::Arguments) {
+pub fn log(level: u32, args: &fmt::Arguments) {
+    // See io::stdio::with_task_stdout for why there's a few dances here. The
+    // gist of it is that arbitrary code can run during logging (and set an
+    // arbitrary logging handle into the task) so we need to be careful that the
+    // local task is in TLS while we're running arbitrary code.
     let mut logger = {
         let mut task = Local::borrow(None::<Task>);
         task.get().logger.take()
     };
 
     if logger.is_none() {
-        logger = Some(StdErrLogger::new());
+        logger = Some(~DefaultLogger {
+            handle: LineBufferedWriter::new(io::stderr()),
+        } as ~Logger);
     }
-    logger.get_mut_ref().log(args);
+    logger.get_mut_ref().log(level, args);
+
+    let mut task = Local::borrow(None::<Task>);
+    let prev = util::replace(&mut task.get().logger, logger);
+    drop(task);
+    drop(prev);
+}
 
+/// Replaces the task-local logger with the specified logger, returning the old
+/// logger.
+pub fn set_logger(logger: ~Logger) -> Option<~Logger> {
     let mut task = Local::borrow(None::<Task>);
-    task.get().logger = logger;
+    util::replace(&mut task.get().logger, Some(logger))
 }

src/libstd/rt/logging.rs

@@ -8,13 +8,9 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-use fmt;
 use from_str::from_str;
 use libc::exit;
 use option::{Some, None, Option};
-use io;
-use io::stdio::StdWriter;
-use io::buffered::LineBufferedWriter;
 use rt::crate_map::{ModEntry, CrateMap, iter_crate_map, get_crate_map};
 use str::StrSlice;
 use vec::{ImmutableVector, MutableTotalOrdVector};
@@ -168,28 +164,6 @@ fn update_log_settings(crate_map: &CrateMap, settings: ~str) {
     }
 }
 
-pub trait Logger {
-    fn log(&mut self, args: &fmt::Arguments);
-}
-
-/// 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: LineBufferedWriter<StdWriter>,
-}
-
-impl StdErrLogger {
-    pub fn new() -> StdErrLogger {
-        StdErrLogger { handle: LineBufferedWriter::new(io::stderr()) }
-    }
-}
-
-impl Logger for StdErrLogger {
-    fn log(&mut self, args: &fmt::Arguments) {
-        fmt::writeln(&mut self.handle as &mut io::Writer, args);
-    }
-}
-
 /// Configure logging by traversing the crate map and setting the
 /// per-module global logging flags based on the logging spec
 pub fn init() {