std panicking: Provide panic::always_abort

We must change the atomic read on panic entry to `Acquire`, to pick up
a possible an `always_panic` on another thread.

We add `count` to the names of panic_count::get and ::is_zaero,
because now there is another reason why panic ought to maybe abort.
Renaming these ensures that we have checked every call site to ensure
that they don't need further adjustment.

Signed-off-by: Ian Jackson <ijackson@chiark.greenend.org.uk>
Co-authored-by: Mara Bos <m-ou.se@m-ou.se>

Author: ijackson

library/std/src/panic.rs

@@ -461,5 +461,41 @@ pub fn resume_unwind(payload: Box<dyn Any + Send>) -> ! {
     panicking::rust_panic_without_hook(payload)
 }
 
+/// Make all future panics abort directly without running the panic hook or unwinding.
+///
+/// There is no way to undo this; the effect lasts until the process exits or
+/// execs (or the equivalent).
+///
+/// # Use after fork
+///
+/// This function is particularly useful for calling after `libc::fork`.  After `fork`, in a
+/// multithreaded program it is (on many platforms) not safe to call the allocator.  It is also
+/// generally highly undesirable for an unwind to unwind past the `fork`, because that results in
+/// the unwind propagating to code that was only ever expecting to run in the parent.
+///
+/// `panic::always_abort()` helps avoid both of these.  It directly avoids any further unwinding,
+/// and if there is a panic, the abort will occur without allocating provided that the arguments to
+/// panic can be formatted without allocating.
+///
+/// Examples
+///
+/// ```no_run
+/// #![feature(panic_always_abort)]
+/// use std::panic;
+///
+/// panic::always_abort();
+///
+/// let _ = panic::catch_unwind(|| {
+///     panic!("inside the catch");
+/// });
+///
+/// // We will have aborted already, due to the panic.
+/// unreachable!();
+/// ```
+#[unstable(feature = "panic_always_abort", issue = "84438")]
+pub fn always_abort() {
+    crate::panicking::panic_count::set_always_abort();
+}
+
 #[cfg(test)]
 mod tests;

library/std/src/panicking.rs

@@ -180,7 +180,7 @@ pub fn take_hook() -> Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send> {
 fn default_hook(info: &PanicInfo<'_>) {
     // If this is a double panic, make sure that we print a backtrace
     // for this panic. Otherwise only print it if logging is enabled.
-    let backtrace_env = if panic_count::get() >= 2 {
+    let backtrace_env = if panic_count::get_count() >= 2 {
         RustBacktrace::Print(crate::backtrace_rs::PrintFmt::Full)
     } else {
         backtrace::rust_backtrace_env()
@@ -233,6 +233,8 @@ pub mod panic_count {
     use crate::cell::Cell;
     use crate::sync::atomic::{AtomicUsize, Ordering};
 
+    pub const ALWAYS_ABORT_FLAG: usize = 1 << (usize::BITS - 1);
+
     // Panic count for the current thread.
     thread_local! { static LOCAL_PANIC_COUNT: Cell<usize> = Cell::new(0) }
 
@@ -241,15 +243,29 @@ pub mod panic_count {
     // thread, if that thread currently views `GLOBAL_PANIC_COUNT` as being zero,
     // then `LOCAL_PANIC_COUNT` in that thread is zero. This invariant holds before
     // and after increase and decrease, but not necessarily during their execution.
+    //
+    // Additionally, the top bit of GLOBAL_PANIC_COUNT (GLOBAL_ALWAYS_ABORT_FLAG)
+    // records whether panic::always_abort() has been called.  This can only be
+    // set, never cleared.
+    //
+    // This could be viewed as a struct containing a single bit and an n-1-bit
+    // value, but if we wrote it like that it would be more than a single word,
+    // and even a newtype around usize would be clumsy because we need atomics.
+    // But we use such a tuple for the return type of increase().
+    //
+    // Stealing a bit is fine because it just amounts to assuming that each
+    // panicking thread consumes at least 2 bytes of address space.
     static GLOBAL_PANIC_COUNT: AtomicUsize = AtomicUsize::new(0);
 
-    pub fn increase() -> usize {
-        GLOBAL_PANIC_COUNT.fetch_add(1, Ordering::Relaxed);
-        LOCAL_PANIC_COUNT.with(|c| {
-            let next = c.get() + 1;
-            c.set(next);
-            next
-        })
+    pub fn increase() -> (bool, usize) {
+        (
+            GLOBAL_PANIC_COUNT.fetch_add(1, Ordering::Acquire) & ALWAYS_ABORT_FLAG != 0,
+            LOCAL_PANIC_COUNT.with(|c| {
+                let next = c.get() + 1;
+                c.set(next);
+                next
+            }),
+        )
     }
 
     pub fn decrease() {
@@ -261,13 +277,19 @@ pub mod panic_count {
         });
     }
 
-    pub fn get() -> usize {
+    pub fn set_always_abort() {
+        GLOBAL_PANIC_COUNT.fetch_or(ALWAYS_ABORT_FLAG, Ordering::Release);
+    }
+
+    // Disregards ALWAYS_ABORT_FLAG
+    pub fn get_count() -> usize {
         LOCAL_PANIC_COUNT.with(|c| c.get())
     }
 
+    // Disregards ALWAYS_ABORT_FLAG
     #[inline]
-    pub fn is_zero() -> bool {
-        if GLOBAL_PANIC_COUNT.load(Ordering::Relaxed) == 0 {
+    pub fn count_is_zero() -> bool {
+        if GLOBAL_PANIC_COUNT.load(Ordering::Relaxed) & !ALWAYS_ABORT_FLAG == 0 {
             // Fast path: if `GLOBAL_PANIC_COUNT` is zero, all threads
             // (including the current one) will have `LOCAL_PANIC_COUNT`
             // equal to zero, so TLS access can be avoided.
@@ -410,7 +432,7 @@ pub unsafe fn r#try<R, F: FnOnce() -> R>(f: F) -> Result<R, Box<dyn Any + Send>>
 /// Determines whether the current thread is unwinding because of panic.
 #[inline]
 pub fn panicking() -> bool {
-    !panic_count::is_zero()
+    !panic_count::count_is_zero()
 }
 
 /// The entry point for panicking with a formatted message.
@@ -563,15 +585,27 @@ fn rust_panic_with_hook(
     message: Option<&fmt::Arguments<'_>>,
     location: &Location<'_>,
 ) -> ! {
-    let panics = panic_count::increase();
+    let (must_abort, panics) = panic_count::increase();
 
     // 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 > 2 {
-        util::dumb_print(format_args!("thread panicked while processing panic. aborting.\n"));
+    if must_abort || panics > 2 {
+        if panics > 2 {
+            // Don't try to print the message in this case
+            // - perhaps that is causing the recursive panics.
+            util::dumb_print(format_args!("thread panicked while processing panic. aborting.\n"));
+        } else {
+            // Unfortunately, this does not print a backtrace, because creating
+            // a `Backtrace` will allocate, which we must to avoid here.
+            let panicinfo = PanicInfo::internal_constructor(message, location);
+            util::dumb_print(format_args!(
+                "{}\npanicked after panic::always_abort(), aborting.\n",
+                panicinfo
+            ));
+        }
         intrinsics::abort()
     }