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.

2 files changed, 46 insertions, 35 deletions

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.