Change panic::update_hook to simplify usage

And to remove possibility of panics while changing the panic handler,
because that resulted in a double panic.

5 files changed, 71 insertions, 54 deletions

diff --git a/library/alloc/tests/slice.rs b/library/alloc/tests/slice.rs
index a02f7b1f277..b93d7938bc9 100644
--- a/library/alloc/tests/slice.rs
+++ b/library/alloc/tests/slice.rs
@@ -1783,12 +1783,10 @@ thread_local!(static SILENCE_PANIC: Cell<bool> = Cell::new(false));
 #[test]
 #[cfg_attr(target_os = "emscripten", ignore)] // no threads
 fn panic_safe() {
-    panic::update_hook(|prev| {
-        Box::new(move |info| {
-            if !SILENCE_PANIC.with(|s| s.get()) {
-                prev(info);
-            }
-        })
+    panic::update_hook(move |prev, 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 5ff7bbf6f96..9e9750eb8de 100644
--- a/library/proc_macro/src/bridge/client.rs
+++ b/library/proc_macro/src/bridge/client.rs
@@ -310,16 +310,14 @@ 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(|| {
-            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)
-                    }
-                })
+            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)
+                }
             });
         });
 
diff --git a/library/std/src/panicking.rs b/library/std/src/panicking.rs
index cf970dccfc9..19040cb12e0 100644
--- a/library/std/src/panicking.rs
+++ b/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;
 
@@ -180,7 +186,8 @@ pub fn take_hook() -> Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send> {
     }
 }
 
-/// Atomic combination of [`take_hook`] + [`set_hook`].
+/// 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
@@ -189,16 +196,6 @@ pub fn take_hook() -> Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send> {
 ///
 /// 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.
@@ -207,11 +204,15 @@ pub fn take_hook() -> Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send> {
 /// #![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);
-///     })
+/// // 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");
@@ -219,9 +220,10 @@ pub fn take_hook() -> Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send> {
 #[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>,
+    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");
@@ -232,13 +234,12 @@ where
         let old_hook = HOOK;
         HOOK = Hook::Default;
 
-        let hook_for_fn = match old_hook {
+        let prev = 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));
+        HOOK = Hook::custom(move |info| hook_fn(&prev, info));
         drop(guard);
     }
 }
diff --git a/src/test/ui/panics/panic-handler-chain-update-hook.rs b/src/test/ui/panics/panic-handler-chain-update-hook.rs
new file mode 100644
index 00000000000..4dd08ba4ad4
--- /dev/null
+++ b/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));
+}
diff --git a/src/test/ui/panics/panic-while-updating-hook.rs b/src/test/ui/panics/panic-while-updating-hook.rs
deleted file mode 100644
index 8c95f1b8b78..00000000000
--- a/src/test/ui/panics/panic-while-updating-hook.rs
+++ /dev/null
@@ -1,16 +0,0 @@
-// 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");
-    })
-}