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.

1 files changed, 16 insertions, 8 deletions

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
     }
 }