polars_error/
abort.rs

1use std::any::Any;
2use std::panic::{UnwindSafe, catch_unwind};
3use std::sync::atomic::{AtomicU64, Ordering};
4
5/// Python hooks SIGINT to instead generate a KeyboardInterrupt exception.
6/// So we do the same to try and abort long-running computations and return to
7/// Python so that the Python exception can be generated.
8///
9/// We also use this mechanic to abort queries that run out of disk space while
10/// spilling, which can happen from anywhere in the out-of-core code, meaning
11/// there might not be a suitable PolarsResult return path.
12pub enum QueryAborted {
13    KeyboardInterrupt,
14    OocOutOfDisk,
15}
16
17// We use a unique string so we can detect it in backtraces.
18static POLARS_ABORT_PREFIX: &str = "__POLARS_ABORT_";
19static POLARS_ABORT_KEYBOARD_INTERRUPT_STR: &str = "__POLARS_ABORT_KEYBOARD_INTERRUPT";
20static POLARS_ABORT_OOC_OUT_OF_DISK_STR: &str = "__POLARS_ABORT_OOC_OUT_OF_DISK";
21
22// Bottom two bits: abort flags.
23// Top 62 bits: number of alive abort catchers.
24const ABORT_KEYBOARD_INTERRUPT_BIT: u64 = 1;
25const ABORT_OOC_OUT_OF_DISK_BIT: u64 = 2;
26const ABORT_CATCHERS_UNIT: u64 = 4;
27static ABORT_STATE: AtomicU64 = AtomicU64::new(0);
28
29fn decode_polars_abort(p: &dyn Any) -> Option<QueryAborted> {
30    let s = if let Some(s) = p.downcast_ref::<&str>() {
31        s
32    } else if let Some(s) = p.downcast_ref::<String>() {
33        s.as_str()
34    } else {
35        return None;
36    };
37
38    if !s.contains(POLARS_ABORT_PREFIX) {
39        return None;
40    }
41
42    if s.contains(POLARS_ABORT_KEYBOARD_INTERRUPT_STR) {
43        Some(QueryAborted::KeyboardInterrupt)
44    } else if s.contains(POLARS_ABORT_OOC_OUT_OF_DISK_STR) {
45        Some(QueryAborted::OocOutOfDisk)
46    } else {
47        unreachable!()
48    }
49}
50
51pub fn register_polars_abort_mechanism() {
52    let default_hook = std::panic::take_hook();
53    std::panic::set_hook(Box::new(move |p| {
54        // Suppress output if there is an active catcher and the panic message
55        // contains the abort string.
56        let num_catchers =
57            ABORT_STATE.load(Ordering::Relaxed) >> ABORT_CATCHERS_UNIT.trailing_zeros();
58        let suppress = num_catchers > 0 && decode_polars_abort(p.payload()).is_some();
59        if !suppress {
60            default_hook(p);
61        }
62    }));
63
64    // WASM doesn't support signals, so we just skip installing the hook there.
65    #[cfg(not(target_family = "wasm"))]
66    unsafe {
67        // SAFETY: we only do an atomic op in the signal handler, which is allowed.
68        signal_hook::low_level::register(signal_hook::consts::signal::SIGINT, move || {
69            // Set the keyboard interrupt flag, but only if there are active catchers.
70            ABORT_STATE
71                .fetch_update(Ordering::Release, Ordering::Relaxed, |state| {
72                    let num_catchers = state >> ABORT_CATCHERS_UNIT.trailing_zeros();
73                    if num_catchers > 0 {
74                        Some(state | ABORT_KEYBOARD_INTERRUPT_BIT)
75                    } else {
76                        None
77                    }
78                })
79                .ok();
80        })
81        .unwrap();
82    }
83}
84
85pub fn polars_abort_ooc_out_of_disk() -> ! {
86    ABORT_STATE
87        .fetch_update(Ordering::Release, Ordering::Relaxed, |state| {
88            let num_catchers = state >> ABORT_CATCHERS_UNIT.trailing_zeros();
89            if num_catchers > 0 {
90                Some(state | ABORT_OOC_OUT_OF_DISK_BIT)
91            } else {
92                None
93            }
94        })
95        .ok();
96
97    std::panic::panic_any(POLARS_ABORT_OOC_OUT_OF_DISK_STR);
98}
99
100/// Checks if the abort flag is set, and if yes panics. This function is very cheap.
101#[inline(always)]
102pub fn try_raise_polars_abort() {
103    if ABORT_STATE.load(Ordering::Acquire) & (ABORT_CATCHERS_UNIT - 1) != 0 {
104        try_raise_polars_abort_slow()
105    }
106}
107
108#[inline(never)]
109#[cold]
110fn try_raise_polars_abort_slow() {
111    let state = ABORT_STATE.load(Ordering::Acquire);
112    if state & ABORT_KEYBOARD_INTERRUPT_BIT != 0 {
113        std::panic::panic_any(POLARS_ABORT_KEYBOARD_INTERRUPT_STR);
114    } else if state & ABORT_OOC_OUT_OF_DISK_BIT != 0 {
115        std::panic::panic_any(POLARS_ABORT_OOC_OUT_OF_DISK_STR);
116    } else {
117        unreachable!()
118    }
119}
120
121/// Runs the passed function, catching any query abortions if they occur
122/// while running the function.
123pub fn catch_polars_abort<R, F: FnOnce() -> R + UnwindSafe>(try_fn: F) -> Result<R, QueryAborted> {
124    // Try to register this catcher (or immediately return if there is an
125    // uncaught interrupt).
126    try_register_catcher()?;
127    let ret = catch_unwind(try_fn);
128    unregister_catcher();
129    ret.map_err(|p| {
130        if let Some(reason) = decode_polars_abort(&*p) {
131            reason
132        } else {
133            std::panic::resume_unwind(p)
134        }
135    })
136}
137
138fn try_register_catcher() -> Result<(), QueryAborted> {
139    let old_state = ABORT_STATE.fetch_add(ABORT_CATCHERS_UNIT, Ordering::Relaxed);
140    if old_state & (ABORT_CATCHERS_UNIT - 1) != 0 {
141        unregister_catcher();
142
143        return if old_state & ABORT_KEYBOARD_INTERRUPT_BIT != 0 {
144            Err(QueryAborted::KeyboardInterrupt)
145        } else if old_state & ABORT_OOC_OUT_OF_DISK_BIT != 0 {
146            Err(QueryAborted::OocOutOfDisk)
147        } else {
148            unreachable!()
149        };
150    }
151    Ok(())
152}
153
154fn unregister_catcher() {
155    ABORT_STATE
156        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |state| {
157            let num_catchers = state >> ABORT_CATCHERS_UNIT.trailing_zeros();
158            if num_catchers > 1 {
159                Some(state - ABORT_CATCHERS_UNIT)
160            } else {
161                // Last catcher, clear abort flags.
162                Some(0)
163            }
164        })
165        .ok();
166}
polars_error/abort.rs

polars_error/
abort.rs
