Implement panic::update_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));
+}