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
pub use atexit::_run_exitfuncs;
pub(crate) use atexit::module_def;
#[pymodule]
mod atexit {
use crate::{
AsObject, PyObjectRef, PyResult, VirtualMachine, common::rc::PyRc, function::FuncArgs,
};
#[pyfunction]
fn register(func: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyObjectRef {
// Callbacks go in LIFO order (insert at front)
vm.state
.atexit_funcs
.lock()
.insert(0, PyRc::new((func.clone(), args)));
func
}
#[pyfunction]
fn _clear(vm: &VirtualMachine) {
vm.state.atexit_funcs.lock().clear();
}
#[pyfunction]
fn unregister(func: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
// Iterate backward (oldest to newest in LIFO list).
// Release the lock during comparison so __eq__ can call atexit functions.
let mut i = {
let funcs = vm.state.atexit_funcs.lock();
funcs.len() as isize - 1
};
while i >= 0 {
let entry = {
let funcs = vm.state.atexit_funcs.lock();
if i as usize >= funcs.len() {
i = funcs.len() as isize;
i -= 1;
continue;
}
// Keep the entry alive for as long as it is being compared, so
// it cannot be dropped and have its address handed to a
// callback registered from within __eq__.
funcs[i as usize].clone()
};
// Lock released: __eq__ can safely call atexit functions
let eq = vm.bool_eq(&func, &entry.0)?;
if eq {
// The entry may have moved during __eq__. Search backward by identity.
let mut funcs = vm.state.atexit_funcs.lock();
let mut j = (funcs.len() as isize - 1).min(i);
while j >= 0 {
if PyRc::ptr_eq(funcs.get(j as usize).unwrap(), &entry) {
funcs.remove(j as usize);
i = j;
break;
}
j -= 1;
}
}
{
let funcs = vm.state.atexit_funcs.lock();
if i as usize >= funcs.len() {
i = funcs.len() as isize;
}
}
i -= 1;
}
Ok(())
}
#[pyfunction]
pub fn _run_exitfuncs(vm: &VirtualMachine) {
let funcs: Vec<_> = core::mem::take(&mut *vm.state.atexit_funcs.lock());
// Callbacks stored in LIFO order, iterate forward
for entry in funcs {
let (func, args) = PyRc::try_unwrap(entry).unwrap_or_else(|e| (*e).clone());
if let Err(e) = func.call(args, vm) {
let exit = e.fast_isinstance(vm.ctx.exceptions.system_exit);
let msg = func
.repr(vm)
.ok()
.map(|r| format!("Exception ignored in atexit callback {}", r.as_wtf8()));
vm.run_unraisable(e, msg, vm.ctx.none());
if exit {
break;
}
}
}
}
#[pyfunction]
fn _ncallbacks(vm: &VirtualMachine) -> usize {
vm.state.atexit_funcs.lock().len()
}
}