proptest/test_runner/
scoped_panic_hook.rs

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

proptest/test_runner/
scoped_panic_hook.rs
