1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
//-
// Copyright 2024, 2026 The proptest developers
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
#[cfg(feature = "handle-panics")]
mod internal {
//! Implementation of scoped panic hooks
//!
//! 1. `with_hook` serves as entry point, it executes body closure with panic hook closure
//! installed as scoped panic hook
//! 2. Upon first execution, current panic hook is replaced with `scoped_hook_dispatcher`
//! in a thread-safe manner, and original hook is stored for later use
//! 3. When panic occurs, `scoped_hook_dispatcher` either delegates execution to scoped
//! panic hook, if one is installed, or back to original hook stored earlier.
//! This preserves original behavior when scoped hook isn't used
//! 4. When `with_hook` is used, it replaces stored scoped hook pointer with pointer to
//! hook closure passed as parameter. Old hook pointer is set to be restored unconditionally
//! via drop guard. Then, normal body closure is executed.
use std::boxed::Box;
use std::cell::Cell;
use std::panic::{set_hook, take_hook, PanicInfo};
use std::sync::{Once, OnceLock};
use std::{mem, ptr};
thread_local! {
/// Pointer to currently installed scoped panic hook, if any
///
/// NB: pointers to arbitrary fn's are fat, and Rust doesn't allow crafting null pointers
/// to fat objects. So we just store const pointer to tuple with whatever data we need
static SCOPED_HOOK_PTR: Cell<*const (*mut dyn FnMut(&PanicInfo<'_>),)> = Cell::new(ptr::null());
}
static INIT_ONCE: Once = Once::new();
/// Default panic hook, the one which was present before installing scoped one.
static DEFAULT_HOOK: OnceLock<Box<dyn Fn(&PanicInfo<'_>) + Send + Sync>> =
OnceLock::new();
/// Replaces currently installed panic hook with `scoped_hook_dispatcher` once,
/// in a thread-safe manner
fn init() {
INIT_ONCE.call_once(|| {
let old_handler = take_hook();
// set DEFAULT_HOOK before installing the global hook. once set_hook
// is invoked, any other thread that panics will invoke the hook. if
// DEFAULT_HOOK is set after the call to set_hook, a race condition
// between the thread assigning DEFAULT_HOOK and the thread panicking,
// reading DEFAULT_HOOK, exists.
assert!(
DEFAULT_HOOK.set(old_handler).is_ok(),
"DEFAULT_HOOK is only ever assigned within a Once"
);
set_hook(Box::new(scoped_hook_dispatcher));
});
}
/// Panic hook which delegates execution to scoped hook,
/// if one installed, or to default hook
fn scoped_hook_dispatcher(info: &PanicInfo<'_>) {
let handler = SCOPED_HOOK_PTR.get();
if !handler.is_null() {
// It's assumed that if container's ptr is not null, ptr to `FnMut` is non-null too.
// Correctness **must** be ensured by hook switch code in `with_hook`
let hook = unsafe { &mut *(*handler).0 };
(hook)(info);
return;
}
if let Some(hook) = DEFAULT_HOOK.get() {
(hook)(info);
}
}
/// Executes stored closure when dropped
struct Finally<F: FnOnce()>(Option<F>);
impl<F: FnOnce()> Finally<F> {
fn new(body: F) -> Self {
Self(Some(body))
}
}
impl<F: FnOnce()> Drop for Finally<F> {
fn drop(&mut self) {
if let Some(body) = self.0.take() {
body();
}
}
}
/// Executes main closure `body` while installing `guard` as scoped panic hook,
/// for execution duration.
///
/// Any panics which happen during execution of `body` are passed to `guard` hook
/// to collect any info necessary, although unwind process is **NOT** interrupted.
/// See module documentation for details
///
/// # Parameters
/// * `panic_hook` - scoped panic hook, functions for the duration of `body` execution
/// * `body` - actual logic covered by `panic_hook`
///
/// # Returns
/// `body`'s return value
pub fn with_hook<R>(
mut panic_hook: impl FnMut(&PanicInfo<'_>),
body: impl FnOnce() -> R,
) -> R {
init();
// Construct scoped hook pointer
let guard_tuple = (unsafe {
// `mem::transmute` is needed due to borrow checker restrictions to erase all lifetimes
mem::transmute(&mut panic_hook as *mut dyn FnMut(&PanicInfo<'_>))
},);
let old_tuple = SCOPED_HOOK_PTR.replace(&guard_tuple);
// Old scoped hook **must** be restored before leaving function scope to keep it sound
let _undo = Finally::new(|| {
SCOPED_HOOK_PTR.set(old_tuple);
});
body()
}
#[cfg(test)]
mod tests {
use super::with_hook;
use std::cell::Cell;
use std::panic::{catch_unwind, AssertUnwindSafe};
#[test]
fn restores_scoped_hook_after_unwind() {
let outer_panics = Cell::new(0);
let inner_panics = Cell::new(0);
with_hook(
|_| outer_panics.set(outer_panics.get() + 1),
|| {
assert!(catch_unwind(AssertUnwindSafe(|| {
with_hook(
|_| inner_panics.set(inner_panics.get() + 1),
|| panic!("inner panic"),
);
}))
.is_err());
assert!(catch_unwind(|| panic!("outer panic")).is_err());
},
);
assert_eq!(inner_panics.get(), 1);
assert_eq!(outer_panics.get(), 1);
}
}
}
#[cfg(not(feature = "handle-panics"))]
mod internal {
use core::panic::PanicInfo;
/// Simply executes `body` and returns its execution result.
/// Hook parameter is ignored
pub fn with_hook<R>(
_: impl FnMut(&PanicInfo<'_>),
body: impl FnOnce() -> R,
) -> R {
body()
}
}
pub use internal::with_hook;