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
// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License in the LICENSE file at the
// root of this repository, or online at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Panic utilities.

use std::any::Any;
use std::cell::RefCell;
use std::panic::{self, UnwindSafe};
use std::process;

#[cfg(feature = "async")]
use tokio::task_local;

thread_local! {
    static CATCHING_UNWIND: RefCell<bool> = RefCell::new(false);
}

#[cfg(feature = "async")]
task_local! {
    pub(crate) static CATCHING_UNWIND_ASYNC: bool;
}

/// Instructs the entire process to abort if any thread panics.
///
/// By default, when a thread panics in Rust, only that thread is affected, and
/// other threads continue running unaffected. This is a bad default. In almost
/// all programs, thread panics are unexpected, unrecoverable, and leave the
/// overall program in an invalid state. It is therefore typically less
/// confusing to abort the entire program.
///
/// For example, consider a simple program with two threads communicating
/// through a channel, where the first thread is waiting for the second thread
/// to send a value over the channel. If the second thread panics, the first
/// thread will block forever for a value that will never be produced. Blocking
/// forever will be more confusing to the end user than aborting the program
/// entirely.
///
/// Computations in which a panic is expected can use the special
/// [`catch_unwind`] function in this module to recover. Note that the
/// `catch_unwind` function in the standard library is **not** compatible with
/// this function.
pub fn set_abort_on_panic() {
    let old_hook = panic::take_hook();
    panic::set_hook(Box::new(move |panic_info| {
        old_hook(panic_info);
        let catching_unwind = CATCHING_UNWIND.with(|v| *v.borrow());
        #[cfg(feature = "async")]
        let catching_unwind_async = CATCHING_UNWIND_ASYNC.try_with(|v| *v).unwrap_or(false);
        #[cfg(not(feature = "async"))]
        let catching_unwind_async = false;
        if !catching_unwind && !catching_unwind_async {
            process::abort();
        }
    }))
}

/// Like [`std::panic::catch_unwind`], but can unwind panics even if
/// [`set_abort_on_panic`] has been called.
pub fn catch_unwind<F, R>(f: F) -> Result<R, Box<dyn Any + Send + 'static>>
where
    F: FnOnce() -> R + UnwindSafe,
{
    CATCHING_UNWIND.with(|catching_unwind| {
        *catching_unwind.borrow_mut() = true;
        #[allow(clippy::disallowed_methods)]
        let res = panic::catch_unwind(f);
        *catching_unwind.borrow_mut() = false;
        res
    })
}