oom_safe/
oom.rs

1//! Catch out-of-memory panic.
2
3use crate::AllocError;
4use std::alloc::Layout;
5use std::panic::{PanicInfo, UnwindSafe};
6use std::sync::atomic::{AtomicBool, Ordering};
7
8fn oom_hook(layout: Layout) {
9    std::panic::panic_any(AllocError(layout))
10}
11
12type PanicHook = Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send>;
13
14fn panic_hook(panic_info: &PanicInfo<'_>) {
15    // panic abort except alloc error
16    if !panic_info.payload().is::<AllocError>() {
17        std::process::abort();
18    }
19}
20
21/// Invokes a closure, capturing the out-of-memory panic if one occurs.
22///
23/// This function will return `Ok` with the closure's result if the closure
24/// does not panic, and will return `AllocError` if allocation error occurs. The
25/// process will abort if other panics occur.
26#[inline]
27pub fn catch_oom<F: FnOnce() -> R + UnwindSafe, R>(f: F) -> Result<R, AllocError> {
28    static SET_HOOK: AtomicBool = AtomicBool::new(false);
29    if !SET_HOOK.load(Ordering::Acquire) {
30        let hook: PanicHook =
31            Box::try_new(panic_hook).map_err(|_| AllocError::new(Layout::new::<PanicHook>()))?;
32        std::panic::set_hook(hook);
33        std::alloc::set_alloc_error_hook(oom_hook);
34        SET_HOOK.store(true, Ordering::Release);
35    }
36
37    let result = std::panic::catch_unwind(f);
38    match result {
39        Ok(r) => Ok(r),
40        Err(e) => match e.downcast_ref::<AllocError>() {
41            None => {
42                unreachable!()
43            }
44            Some(e) => Err(*e),
45        },
46    }
47}
oom_safe/oom.rs

oom_safe/
oom.rs
