std: Tweak rt::at_exit behavior

There have been some recent panics on the bots and this commit is an attempt to
appease them. Previously it was considered invalid to run `rt::at_exit` after
the handlers had already started running. Due to the multithreaded nature of
applications, however, it is not always possible to guarantee this. For example
[this program][ex] will show off the abort.

[ex]: https://gist.github.com/alexcrichton/56300b87af6fa554e52d

The semantics of the `rt::at_exit` function have been modified as such:

* It is now legal to call `rt::at_exit` at any time. The return value now
  indicates whether the closure was successfully registered or not. Callers must
  now decide what to do with this information.
* The `rt::at_exit` handlers will now be run for a fixed number of iterations.
  Common cases (such as the example shown) may end up registering a new handler
  while others are running perhaps once or twice, so this common condition is
  covered by re-running the handlers a fixed number of times, after which new
  registrations are forbidden.

Some usage of `rt::at_exit` was updated to handle these new semantics, but
deprecated or unstable libraries calling `rt::at_exit` were not updated.

6 files changed, 65 insertions, 46 deletions

diff --git a/src/liblog/lib.rs b/src/liblog/lib.rs
index 4537fc763c9..df94a6ac80f 100644
--- a/src/liblog/lib.rs
+++ b/src/liblog/lib.rs
@@ -443,7 +443,7 @@ fn init() {
         DIRECTIVES = boxed::into_raw(box directives);
 
         // Schedule the cleanup for the globals for when the runtime exits.
-        rt::at_exit(move || {
+        let _ = rt::at_exit(move || {
             let _g = LOCK.lock();
             assert!(!DIRECTIVES.is_null());
             let _directives = Box::from_raw(DIRECTIVES);
diff --git a/src/libstd/io/lazy.rs b/src/libstd/io/lazy.rs
index c9b105f72a5..df280dab37d 100644
--- a/src/libstd/io/lazy.rs
+++ b/src/libstd/io/lazy.rs
@@ -35,25 +35,33 @@ impl<T: Send + Sync + 'static> Lazy<T> {
     pub fn get(&'static self) -> Option<Arc<T>> {
         let _g = self.lock.lock();
         unsafe {
-            let mut ptr = *self.ptr.get();
+            let ptr = *self.ptr.get();
             if ptr.is_null() {
-                ptr = boxed::into_raw(self.init());
-                *self.ptr.get() = ptr;
+                Some(self.init())
             } else if ptr as usize == 1 {
-                return None
+                None
+            } else {
+                Some((*ptr).clone())
             }
-            Some((*ptr).clone())
         }
     }
 
-    fn init(&'static self) -> Box<Arc<T>> {
-        rt::at_exit(move || unsafe {
+    unsafe fn init(&'static self) -> Arc<T> {
+        // If we successfully register an at exit handler, then we cache the
+        // `Arc` allocation in our own internal box (it will get deallocated by
+        // the at exit handler). Otherwise we just return the freshly allocated
+        // `Arc`.
+        let registered = rt::at_exit(move || {
             let g = self.lock.lock();
             let ptr = *self.ptr.get();
             *self.ptr.get() = 1 as *mut _;
             drop(g);
             drop(Box::from_raw(ptr))
         });
-        Box::new((self.init)())
+        let ret = (self.init)();
+        if registered.is_ok() {
+            *self.ptr.get() = boxed::into_raw(Box::new(ret.clone()));
+        }
+        return ret
     }
 }
diff --git a/src/libstd/old_io/stdio.rs b/src/libstd/old_io/stdio.rs
index a1c8630e0ec..a48758366f3 100644
--- a/src/libstd/old_io/stdio.rs
+++ b/src/libstd/old_io/stdio.rs
@@ -238,7 +238,7 @@ pub fn stdin() -> StdinReader {
             STDIN = boxed::into_raw(box stdin);
 
             // Make sure to free it at exit
-            rt::at_exit(|| {
+            let _ = rt::at_exit(|| {
                 Box::from_raw(STDIN);
                 STDIN = ptr::null_mut();
             });
diff --git a/src/libstd/rt/at_exit_imp.rs b/src/libstd/rt/at_exit_imp.rs
index 9fa12b1ef30..9079c0aaffb 100644
--- a/src/libstd/rt/at_exit_imp.rs
+++ b/src/libstd/rt/at_exit_imp.rs
@@ -12,6 +12,10 @@
 //!
 //! Documentation can be found on the `rt::at_exit` function.
 
+// FIXME: switch this to use atexit. Currently this
+// segfaults (the queue's memory is mysteriously gone), so
+// instead the cleanup is tied to the `std::rt` entry point.
+
 use boxed;
 use boxed::Box;
 use vec::Vec;
@@ -27,47 +31,56 @@ type Queue = Vec<Thunk<'static>>;
 static LOCK: Mutex = MUTEX_INIT;
 static mut QUEUE: *mut Queue = 0 as *mut Queue;
 
-unsafe fn init() {
+// The maximum number of times the cleanup routines will be run. While running
+// the at_exit closures new ones may be registered, and this count is the number
+// of times the new closures will be allowed to register successfully. After
+// this number of iterations all new registrations will return `false`.
+const ITERS: usize = 10;
+
+unsafe fn init() -> bool {
     if QUEUE.is_null() {
         let state: Box<Queue> = box Vec::new();
         QUEUE = boxed::into_raw(state);
-    } else {
+    } else if QUEUE as usize == 1 {
         // can't re-init after a cleanup
-        rtassert!(QUEUE as uint != 1);
+        return false
     }
 
-    // FIXME: switch this to use atexit as below. Currently this
-    // segfaults (the queue's memory is mysteriously gone), so
-    // instead the cleanup is tied to the `std::rt` entry point.
-    //
-    // ::libc::atexit(cleanup);
+    return true
 }
 
 pub fn cleanup() {
-    unsafe {
-        LOCK.lock();
-        let queue = QUEUE;
-        QUEUE = 1 as *mut _;
-        LOCK.unlock();
+    for i in 0..ITERS {
+        unsafe {
+            LOCK.lock();
+            let queue = QUEUE;
+            QUEUE = if i == ITERS - 1 {1} else {0} as *mut _;
+            LOCK.unlock();
 
-        // make sure we're not recursively cleaning up
-        rtassert!(queue as uint != 1);
+            // make sure we're not recursively cleaning up
+            rtassert!(queue as usize != 1);
 
-        // If we never called init, not need to cleanup!
-        if queue as uint != 0 {
-            let queue: Box<Queue> = Box::from_raw(queue);
-            for to_run in *queue {
-                to_run.invoke(());
+            // If we never called init, not need to cleanup!
+            if queue as usize != 0 {
+                let queue: Box<Queue> = Box::from_raw(queue);
+                for to_run in *queue {
+                    to_run.invoke(());
+                }
             }
         }
     }
 }
 
-pub fn push(f: Thunk<'static>) {
+pub fn push(f: Thunk<'static>) -> bool {
+    let mut ret = true;
     unsafe {
         LOCK.lock();
-        init();
-        (*QUEUE).push(f);
+        if init() {
+            (*QUEUE).push(f);
+        } else {
+            ret = false;
+        }
         LOCK.unlock();
     }
+    return ret
 }
diff --git a/src/libstd/rt/mod.rs b/src/libstd/rt/mod.rs
index e52e68dad23..43aa4155629 100644
--- a/src/libstd/rt/mod.rs
+++ b/src/libstd/rt/mod.rs
@@ -17,14 +17,9 @@
 //! time being.
 
 #![unstable(feature = "std_misc")]
-
-// FIXME: this should not be here.
 #![allow(missing_docs)]
 
-#![allow(dead_code)]
-
-use marker::Send;
-use ops::FnOnce;
+use prelude::v1::*;
 use sys;
 use thunk::Thunk;
 use usize;
@@ -149,13 +144,16 @@ fn lang_start(main: *const u8, argc: int, argv: *const *const u8) -> int {
 
 /// Enqueues a procedure to run when the main thread exits.
 ///
-/// It is forbidden for procedures to register more `at_exit` handlers when they
-/// are running, and doing so will lead to a process abort.
+/// Currently these closures are only run once the main *Rust* thread exits.
+/// Once the `at_exit` handlers begin running, more may be enqueued, but not
+/// infinitely so. Eventually a handler registration will be forced to fail.
 ///
-/// Note that other threads may still be running when `at_exit` routines start
-/// running.
-pub fn at_exit<F: FnOnce() + Send + 'static>(f: F) {
-    at_exit_imp::push(Thunk::new(f));
+/// Returns `Ok` if the handler was successfully registered, meaning that the
+/// closure will be run once the main thread exits. Returns `Err` to indicate
+/// that the closure could not be registered, meaning that it is not scheduled
+/// to be rune.
+pub fn at_exit<F: FnOnce() + Send + 'static>(f: F) -> Result<(), ()> {
+    if at_exit_imp::push(Thunk::new(f)) {Ok(())} else {Err(())}
 }
 
 /// One-time runtime cleanup.
diff --git a/src/libstd/sys/common/helper_thread.rs b/src/libstd/sys/common/helper_thread.rs
index 2a852fbcd57..53f18a57325 100644
--- a/src/libstd/sys/common/helper_thread.rs
+++ b/src/libstd/sys/common/helper_thread.rs
@@ -112,7 +112,7 @@ impl<M: Send> Helper<M> {
                     self.cond.notify_one()
                 });
 
-                rt::at_exit(move || { self.shutdown() });
+                let _ = rt::at_exit(move || { self.shutdown() });
                 *self.initialized.get() = true;
             } else if *self.chan.get() as uint == 1 {
                 panic!("cannot continue usage after shutdown");