Auto merge of #34866 - cynicaldevil:panic-counter, r=alexcrichton

Refactored code to access TLS only in case of panic (II)

Fixes #34787
r? @alexcrichton
Do it **very** carefully this time!

4 files changed, 40 insertions, 23 deletions

diff --git a/src/libstd/lib.rs b/src/libstd/lib.rs
index c6272012d66..c05e0c3ca68 100644
--- a/src/libstd/lib.rs
+++ b/src/libstd/lib.rs
@@ -278,6 +278,7 @@
 #![feature(unwind_attributes)]
 #![feature(vec_push_all)]
 #![feature(zero_one)]
+#![cfg_attr(test, feature(update_panic_count))]
 
 // Issue# 30592: Systematically use alloc_system during stage0 since jemalloc
 // might be unavailable or disabled
diff --git a/src/libstd/panic.rs b/src/libstd/panic.rs
index ba18d15f5c4..2f67081e0d7 100644
--- a/src/libstd/panic.rs
+++ b/src/libstd/panic.rs
@@ -340,5 +340,5 @@ pub fn catch_unwind<F: FnOnce() -> R + UnwindSafe, R>(f: F) -> Result<R> {
 /// ```
 #[stable(feature = "resume_unwind", since = "1.9.0")]
 pub fn resume_unwind(payload: Box<Any + Send>) -> ! {
-    panicking::rust_panic(payload)
+    panicking::update_count_then_panic(payload)
 }
diff --git a/src/libstd/panicking.rs b/src/libstd/panicking.rs
index 57a4c3df70a..8c1567939fb 100644
--- a/src/libstd/panicking.rs
+++ b/src/libstd/panicking.rs
@@ -21,7 +21,6 @@ use prelude::v1::*;
 use io::prelude::*;
 
 use any::Any;
-use cell::Cell;
 use cell::RefCell;
 use fmt;
 use intrinsics;
@@ -39,8 +38,6 @@ thread_local! {
     }
 }
 
-thread_local! { pub static PANIC_COUNT: Cell<usize> = Cell::new(0) }
-
 // Binary interface to the panic runtime that the standard library depends on.
 //
 // The standard library is tagged with `#![needs_panic_runtime]` (introduced in
@@ -187,7 +184,7 @@ fn default_hook(info: &PanicInfo) {
     // for this panic. Otherwise only print it if logging is enabled.
     #[cfg(any(not(cargobuild), feature = "backtrace"))]
     let log_backtrace = {
-        let panics = PANIC_COUNT.with(|c| c.get());
+        let panics = update_panic_count(0);
 
         panics >= 2 || backtrace::log_enabled()
     };
@@ -238,14 +235,31 @@ fn default_hook(info: &PanicInfo) {
     }
 }
 
+
+#[cfg(not(test))]
+#[doc(hidden)]
+#[unstable(feature = "update_panic_count", issue = "0")]
+pub fn update_panic_count(amt: isize) -> usize {
+    use cell::Cell;
+    thread_local! { static PANIC_COUNT: Cell<usize> = Cell::new(0) }
+
+    PANIC_COUNT.with(|c| {
+        let next = (c.get() as isize + amt) as usize;
+        c.set(next);
+        return next
+    })
+}
+
+#[cfg(test)]
+pub use realstd::rt::update_panic_count;
+
 /// Invoke a closure, capturing the cause of an unwinding panic if one occurs.
 pub unsafe fn try<R, F: FnOnce() -> R>(f: F) -> Result<R, Box<Any + Send>> {
     let mut slot = None;
     let mut f = Some(f);
-    let ret = PANIC_COUNT.with(|s| {
-        let prev = s.get();
-        s.set(0);
+    let ret;
 
+    {
         let mut to_run = || {
             slot = Some(f.take().unwrap()());
         };
@@ -258,18 +272,18 @@ pub unsafe fn try<R, F: FnOnce() -> R>(f: F) -> Result<R, Box<Any + Send>> {
                                          dataptr,
                                          &mut any_data,
                                          &mut any_vtable);
-        s.set(prev);
-
         if r == 0 {
-            Ok(())
+            ret = Ok(());
         } else {
-            Err(mem::transmute(raw::TraitObject {
+            update_panic_count(-1);
+            ret = Err(mem::transmute(raw::TraitObject {
                 data: any_data as *mut _,
                 vtable: any_vtable as *mut _,
-            }))
+            }));
         }
-    });
+    }
 
+    debug_assert!(update_panic_count(0) == 0);
     return ret.map(|()| {
         slot.take().unwrap()
     });
@@ -285,7 +299,7 @@ pub unsafe fn try<R, F: FnOnce() -> R>(f: F) -> Result<R, Box<Any + Send>> {
 
 /// Determines whether the current thread is unwinding because of panic.
 pub fn panicking() -> bool {
-    PANIC_COUNT.with(|c| c.get() != 0)
+    update_panic_count(0) != 0
 }
 
 /// Entry point of panic from the libcore crate.
@@ -350,18 +364,14 @@ fn rust_panic_with_hook(msg: Box<Any + Send>,
                         file_line: &(&'static str, u32)) -> ! {
     let (file, line) = *file_line;
 
-    let panics = PANIC_COUNT.with(|c| {
-        let prev = c.get();
-        c.set(prev + 1);
-        prev
-    });
+    let panics = update_panic_count(1);
 
     // If this is the third nested call (e.g. panics == 2, this is 0-indexed),
     // the panic hook probably triggered the last panic, otherwise the
     // double-panic check would have aborted the process. In this case abort the
     // process real quickly as we don't want to try calling it again as it'll
     // probably just panic again.
-    if panics > 1 {
+    if panics > 2 {
         util::dumb_print(format_args!("thread panicked while processing \
                                        panic. aborting.\n"));
         unsafe { intrinsics::abort() }
@@ -383,7 +393,7 @@ fn rust_panic_with_hook(msg: Box<Any + Send>,
         HOOK_LOCK.read_unlock();
     }
 
-    if panics > 0 {
+    if panics > 1 {
         // If a thread panics while it's already unwinding then we
         // have limited options. Currently our preference is to
         // just abort. In the future we may consider resuming
@@ -396,6 +406,12 @@ fn rust_panic_with_hook(msg: Box<Any + Send>,
     rust_panic(msg)
 }
 
+/// Shim around rust_panic. Called by resume_unwind.
+pub fn update_count_then_panic(msg: Box<Any + Send>) -> ! {
+    update_panic_count(1);
+    rust_panic(msg)
+}
+
 /// A private no-mangle function on which to slap yer breakpoints.
 #[no_mangle]
 #[allow(private_no_mangle_fns)] // yes we get it, but we like breakpoints
diff --git a/src/libstd/rt.rs b/src/libstd/rt.rs
index 5a7c0fe4816..a3d9e4db7d1 100644
--- a/src/libstd/rt.rs
+++ b/src/libstd/rt.rs
@@ -25,7 +25,7 @@
 
 
 // Reexport some of our utilities which are expected by other crates.
-pub use panicking::{begin_panic, begin_panic_fmt};
+pub use panicking::{begin_panic, begin_panic_fmt, update_panic_count};
 
 #[cfg(not(test))]
 #[lang = "start"]