Support arbitrary stdout/stderr/logger handles

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() {