Auto merge of #92598 - Badel2:panic-update-hook, r=yaahc

Implement `panic::update_hook`

Add a new function `panic::update_hook` to allow creating panic hooks that forward the call to the previously set panic hook, without race conditions. It works by taking a closure that transforms the old panic hook into a new one, while ensuring that during the execution of the closure no other thread can modify the panic hook. This is a small function so I hope it can be discussed here without a formal RFC, however if you prefer I can write one.

Consider the following example:

```rust
let prev = panic::take_hook();
panic::set_hook(Box::new(move |info| {
    println!("panic handler A");
    prev(info);
}));
```

This is a common pattern in libraries that need to do something in case of panic: log panic to a file, record code coverage, send panic message to a monitoring service, print custom message with link to github to open a new issue, etc. However it is impossible to avoid race conditions with the current API, because two threads can execute in this order:

* Thread A calls `panic::take_hook()`
* Thread B calls `panic::take_hook()`
* Thread A calls `panic::set_hook()`
* Thread B calls `panic::set_hook()`

And the result is that the original panic hook has been lost, as well as the panic hook set by thread A. The resulting panic hook will be the one set by thread B, which forwards to the default panic hook. This is not considered a big issue because the panic handler setup is usually run during initialization code, probably before spawning any other threads.

Using the new `panic::update_hook` function, this race condition is impossible, and the result will be either `A, B, original` or `B, A, original`.

```rust
panic::update_hook(|prev| {
    Box::new(move |info| {
        println!("panic handler A");
        prev(info);
    })
});
```

I found one real world use case here: https://github.com/dtolnay/proc-macro2/blob/988cf403e741aadfd5340bbf67e35e1062a526aa/src/detection.rs#L32 the workaround is to detect the race condition and panic in that case.

The pattern of `take_hook` + `set_hook` is very common, you can see some examples in this pull request, so I think it's natural to have a function that combines them both. Also using `update_hook` instead of `take_hook` + `set_hook` reduces the number of calls to `HOOK_LOCK.write()` from 2 to 1, but I don't expect this to make any difference in performance.

### Unresolved questions:

* `panic::update_hook` takes a closure, if that closure panics the error message is "panicked while processing panic" which is not nice. This is a consequence of holding the `HOOK_LOCK` while executing the closure. Could be avoided using `catch_unwind`?

* Reimplement `panic::set_hook` as `panic::update_hook(|_prev| hook)`?

library/alloc/tests/lib.rs

@@ -38,6 +38,7 @@
 #![feature(const_trait_impl)]
 #![feature(const_str_from_utf8)]
 #![feature(nonnull_slice_from_raw_parts)]
+#![feature(panic_update_hook)]
 
 use std::collections::hash_map::DefaultHasher;
 use std::hash::{Hash, Hasher};

library/alloc/tests/slice.rs

@@ -1783,12 +1783,11 @@ thread_local!(static SILENCE_PANIC: Cell<bool> = Cell::new(false));
 #[test]
 #[cfg_attr(target_os = "emscripten", ignore)] // no threads
 fn panic_safe() {
-    let prev = panic::take_hook();
-    panic::set_hook(Box::new(move |info| {
+    panic::update_hook(move |prev, info| {
         if !SILENCE_PANIC.with(|s| s.get()) {
             prev(info);
         }
-    }));
+    });
 
     let mut rng = thread_rng();
 

library/proc_macro/src/bridge/client.rs

@@ -310,16 +310,15 @@ impl Bridge<'_> {
         // NB. the server can't do this because it may use a different libstd.
         static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
         HIDE_PANICS_DURING_EXPANSION.call_once(|| {
-            let prev = panic::take_hook();
-            panic::set_hook(Box::new(move |info| {
+            panic::update_hook(move |prev, info| {
                 let show = BridgeState::with(|state| match state {
                     BridgeState::NotConnected => true,
                     BridgeState::Connected(_) | BridgeState::InUse => force_show_panics,
                 });
                 if show {
                     prev(info)
                 }
-            }));
+            });
         });
 
         BRIDGE_STATE.with(|state| state.set(BridgeState::Connected(self), f))

library/proc_macro/src/lib.rs

@@ -30,6 +30,7 @@
 #![feature(restricted_std)]
 #![feature(rustc_attrs)]
 #![feature(min_specialization)]
+#![feature(panic_update_hook)]
 #![recursion_limit = "256"]
 
 #[unstable(feature = "proc_macro_internals", issue = "27812")]

library/std/src/panic.rs

@@ -36,6 +36,9 @@ pub use core::panic::panic_2021;
 #[stable(feature = "panic_hooks", since = "1.10.0")]
 pub use crate::panicking::{set_hook, take_hook};
 
+#[unstable(feature = "panic_update_hook", issue = "92649")]
+pub use crate::panicking::update_hook;
+
 #[stable(feature = "panic_hooks", since = "1.10.0")]
 pub use core::panic::{Location, PanicInfo};
 

library/std/src/panicking.rs

@@ -76,6 +76,12 @@ enum Hook {
     Custom(*mut (dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send)),
 }
 
+impl Hook {
+    fn custom(f: impl Fn(&PanicInfo<'_>) + 'static + Sync + Send) -> Self {
+        Self::Custom(Box::into_raw(Box::new(f)))
+    }
+}
+
 static HOOK_LOCK: StaticRWLock = StaticRWLock::new();
 static mut HOOK: Hook = Hook::Default;
 
@@ -118,6 +124,11 @@ pub fn set_hook(hook: Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send>) {
         panic!("cannot modify the panic hook from a panicking thread");
     }
 
+    // SAFETY:
+    //
+    // - `HOOK` can only be modified while holding write access to `HOOK_LOCK`.
+    // - The argument of `Box::from_raw` is always a valid pointer that was created using
+    // `Box::into_raw`.
     unsafe {
         let guard = HOOK_LOCK.write();
         let old_hook = HOOK;
@@ -167,6 +178,11 @@ pub fn take_hook() -> Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send> {
         panic!("cannot modify the panic hook from a panicking thread");
     }
 
+    // SAFETY:
+    //
+    // - `HOOK` can only be modified while holding write access to `HOOK_LOCK`.
+    // - The argument of `Box::from_raw` is always a valid pointer that was created using
+    // `Box::into_raw`.
     unsafe {
         let guard = HOOK_LOCK.write();
         let hook = HOOK;
@@ -180,6 +196,69 @@ pub fn take_hook() -> Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send> {
     }
 }
 
+/// Atomic combination of [`take_hook`] and [`set_hook`]. Use this to replace the panic handler with
+/// a new panic handler that does something and then executes the old handler.
+///
+/// [`take_hook`]: ./fn.take_hook.html
+/// [`set_hook`]: ./fn.set_hook.html
+///
+/// # Panics
+///
+/// Panics if called from a panicking thread.
+///
+/// # Examples
+///
+/// The following will print the custom message, and then the normal output of panic.
+///
+/// ```should_panic
+/// #![feature(panic_update_hook)]
+/// use std::panic;
+///
+/// // Equivalent to
+/// // let prev = panic::take_hook();
+/// // panic::set_hook(move |info| {
+/// //     println!("...");
+/// //     prev(info);
+/// // );
+/// panic::update_hook(move |prev, info| {
+///     println!("Print custom message and execute panic handler as usual");
+///     prev(info);
+/// });
+///
+/// panic!("Custom and then normal");
+/// ```
+#[unstable(feature = "panic_update_hook", issue = "92649")]
+pub fn update_hook<F>(hook_fn: F)
+where
+    F: Fn(&(dyn Fn(&PanicInfo<'_>) + Send + Sync + 'static), &PanicInfo<'_>)
+        + Sync
+        + Send
+        + 'static,
+{
+    if thread::panicking() {
+        panic!("cannot modify the panic hook from a panicking thread");
+    }
+
+    // SAFETY:
+    //
+    // - `HOOK` can only be modified while holding write access to `HOOK_LOCK`.
+    // - The argument of `Box::from_raw` is always a valid pointer that was created using
+    // `Box::into_raw`.
+    unsafe {
+        let guard = HOOK_LOCK.write();
+        let old_hook = HOOK;
+        HOOK = Hook::Default;
+
+        let prev = match old_hook {
+            Hook::Default => Box::new(default_hook),
+            Hook::Custom(ptr) => Box::from_raw(ptr),
+        };
+
+        HOOK = Hook::custom(move |info| hook_fn(&prev, info));
+        drop(guard);
+    }
+}
+
 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.

src/test/ui/panics/panic-handler-chain-update-hook.rs

@@ -0,0 +1,36 @@
+// run-pass
+// needs-unwind
+#![allow(stable_features)]
+
+// ignore-emscripten no threads support
+
+#![feature(std_panic)]
+#![feature(panic_update_hook)]
+
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::panic;
+use std::thread;
+
+static A: AtomicUsize = AtomicUsize::new(0);
+static B: AtomicUsize = AtomicUsize::new(0);
+static C: AtomicUsize = AtomicUsize::new(0);
+
+fn main() {
+    panic::set_hook(Box::new(|_| { A.fetch_add(1, Ordering::SeqCst); }));
+    panic::update_hook(|prev, info| {
+        B.fetch_add(1, Ordering::SeqCst);
+        prev(info);
+    });
+    panic::update_hook(|prev, info| {
+        C.fetch_add(1, Ordering::SeqCst);
+        prev(info);
+    });
+
+    let _ = thread::spawn(|| {
+        panic!();
+    }).join();
+
+    assert_eq!(1, A.load(Ordering::SeqCst));
+    assert_eq!(1, B.load(Ordering::SeqCst));
+    assert_eq!(1, C.load(Ordering::SeqCst));
+}