Implement panic::update_hook

7 files changed, 102 insertions, 16 deletions

diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs
index eec24a5c3f7..7b8eeb90b5a 100644
--- a/library/alloc/tests/lib.rs
+++ b/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};
diff --git a/library/alloc/tests/slice.rs b/library/alloc/tests/slice.rs
index 18ea6a21413..a02f7b1f277 100644
--- a/library/alloc/tests/slice.rs
+++ b/library/alloc/tests/slice.rs
@@ -1783,12 +1783,13 @@ 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| {
-        if !SILENCE_PANIC.with(|s| s.get()) {
-            prev(info);
-        }
-    }));
+    panic::update_hook(|prev| {
+        Box::new(move |info| {
+            if !SILENCE_PANIC.with(|s| s.get()) {
+                prev(info);
+            }
+        })
+    });
 
     let mut rng = thread_rng();
 
diff --git a/library/proc_macro/src/bridge/client.rs b/library/proc_macro/src/bridge/client.rs
index 83a2ac6f0d4..5ff7bbf6f96 100644
--- a/library/proc_macro/src/bridge/client.rs
+++ b/library/proc_macro/src/bridge/client.rs
@@ -310,16 +310,17 @@ 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| {
-                let show = BridgeState::with(|state| match state {
-                    BridgeState::NotConnected => true,
-                    BridgeState::Connected(_) | BridgeState::InUse => force_show_panics,
-                });
-                if show {
-                    prev(info)
-                }
-            }));
+            panic::update_hook(|prev| {
+                Box::new(move |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))
diff --git a/library/proc_macro/src/lib.rs b/library/proc_macro/src/lib.rs
index 69af598f91e..c5afca6d56a 100644
--- a/library/proc_macro/src/lib.rs
+++ b/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")]
diff --git a/library/std/src/panic.rs b/library/std/src/panic.rs
index c0605b2f412..02ecf2e3e82 100644
--- a/library/std/src/panic.rs
+++ b/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};
 
diff --git a/library/std/src/panicking.rs b/library/std/src/panicking.rs
index 87854fe4f29..cf970dccfc9 100644
--- a/library/std/src/panicking.rs
+++ b/library/std/src/panicking.rs
@@ -180,6 +180,69 @@ pub fn take_hook() -> Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send> {
     }
 }
 
+/// Atomic combination of [`take_hook`] + [`set_hook`].
+///
+/// [`take_hook`]: ./fn.take_hook.html
+/// [`set_hook`]: ./fn.set_hook.html
+///
+/// # Panics
+///
+/// Panics if called from a panicking thread.
+///
+/// Panics if the provided closure calls any of the functions [`panic::take_hook`],
+/// [`panic::set_hook`], or [`panic::update_hook`].
+///
+/// Note: if the provided closure panics, the panic will not be able to be handled, resulting in a
+/// double panic that aborts the process with a generic error message.
+///
+/// [`panic::take_hook`]: ./fn.take_hook.html
+/// [`panic::set_hook`]: ./fn.set_hook.html
+/// [`panic::update_hook`]: ./fn.update_hook.html
+///
+/// # Examples
+///
+/// The following will print the custom message, and then the normal output of panic.
+///
+/// ```should_panic
+/// #![feature(panic_update_hook)]
+/// use std::panic;
+///
+/// panic::update_hook(|prev| {
+///     Box::new(move |panic_info| {
+///         println!("Print custom message and execute panic handler as usual");
+///         prev(panic_info);
+///     })
+/// });
+///
+/// panic!("Custom and then normal");
+/// ```
+#[unstable(feature = "panic_update_hook", issue = "92649")]
+pub fn update_hook<F>(hook_fn: F)
+where
+    F: FnOnce(
+        Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send>,
+    ) -> Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send>,
+{
+    if thread::panicking() {
+        panic!("cannot modify the panic hook from a panicking thread");
+    }
+
+    unsafe {
+        let guard = HOOK_LOCK.write();
+        let old_hook = HOOK;
+        HOOK = Hook::Default;
+
+        let hook_for_fn = match old_hook {
+            Hook::Default => Box::new(default_hook),
+            Hook::Custom(ptr) => Box::from_raw(ptr),
+        };
+
+        let hook = hook_fn(hook_for_fn);
+        HOOK = Hook::Custom(Box::into_raw(hook));
+        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.
diff --git a/src/test/ui/panics/panic-while-updating-hook.rs b/src/test/ui/panics/panic-while-updating-hook.rs
new file mode 100644
index 00000000000..8c95f1b8b78
--- /dev/null
+++ b/src/test/ui/panics/panic-while-updating-hook.rs
@@ -0,0 +1,16 @@
+// run-fail
+// error-pattern: panicked while processing panic
+#![allow(stable_features)]
+
+// ignore-emscripten no threads support
+
+#![feature(std_panic)]
+#![feature(panic_update_hook)]
+
+use std::panic;
+
+fn main() {
+    panic::update_hook(|_prev| {
+        panic!("inside update_hook");
+    })
+}