84 lines
3.2 KiB
Rust
84 lines
3.2 KiB
Rust
//! Observe actual CPython allocation failures before AnySystem replaces a
|
|
//! Python exception with a generic panic. Student stdout is not a result channel.
|
|
use std::ffi::c_void;
|
|
use std::mem::MaybeUninit;
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
use std::sync::Once;
|
|
|
|
use pyo3::ffi::{self, PyMemAllocatorDomain, PyMemAllocatorEx};
|
|
use pyo3::Python;
|
|
|
|
static INSTALLED: Once = Once::new();
|
|
static FAILED: AtomicBool = AtomicBool::new(false);
|
|
|
|
fn observe(pointer: *mut c_void, nonzero: bool) -> *mut c_void {
|
|
if nonzero && pointer.is_null() {
|
|
FAILED.store(true, Ordering::Relaxed);
|
|
}
|
|
pointer
|
|
}
|
|
|
|
extern "C" fn malloc(ctx: *mut c_void, size: usize) -> *mut c_void {
|
|
// SAFETY: ctx points to an immutable original allocator, leaked for the
|
|
// interpreter's entire lifetime. Calls retain CPython's allocator contract.
|
|
let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
|
|
observe((allocator.malloc.unwrap())(allocator.ctx, size), size != 0)
|
|
}
|
|
extern "C" fn calloc(ctx: *mut c_void, count: usize, size: usize) -> *mut c_void {
|
|
// SAFETY: same lifetime and callback contract as malloc above.
|
|
let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
|
|
observe(
|
|
(allocator.calloc.unwrap())(allocator.ctx, count, size),
|
|
count != 0 && size != 0,
|
|
)
|
|
}
|
|
extern "C" fn realloc(ctx: *mut c_void, pointer: *mut c_void, size: usize) -> *mut c_void {
|
|
// SAFETY: same lifetime and callback contract as malloc above.
|
|
let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
|
|
observe(
|
|
(allocator.realloc.unwrap())(allocator.ctx, pointer, size),
|
|
size != 0,
|
|
)
|
|
}
|
|
extern "C" fn free(ctx: *mut c_void, pointer: *mut c_void) {
|
|
// SAFETY: same lifetime and callback contract as malloc above.
|
|
let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
|
|
(allocator.free.unwrap())(allocator.ctx, pointer);
|
|
}
|
|
|
|
pub fn install() {
|
|
Python::attach(|_| {
|
|
INSTALLED.call_once(|| {
|
|
for domain in [
|
|
PyMemAllocatorDomain::PYMEM_DOMAIN_RAW,
|
|
PyMemAllocatorDomain::PYMEM_DOMAIN_MEM,
|
|
PyMemAllocatorDomain::PYMEM_DOMAIN_OBJ,
|
|
] {
|
|
// SAFETY: installation holds the GIL before any student code runs.
|
|
// CPython requires post-initialization hooks to wrap the previous
|
|
// allocator. Its four callbacks are supplied by the interpreter.
|
|
unsafe {
|
|
let mut original = MaybeUninit::uninit();
|
|
ffi::PyMem_GetAllocator(domain, original.as_mut_ptr());
|
|
let original = Box::into_raw(Box::new(original.assume_init()));
|
|
let mut wrapper = PyMemAllocatorEx {
|
|
ctx: original.cast(),
|
|
malloc: Some(malloc),
|
|
calloc: Some(calloc),
|
|
realloc: Some(realloc),
|
|
free: Some(free),
|
|
};
|
|
ffi::PyMem_SetAllocator(domain, &mut wrapper);
|
|
}
|
|
}
|
|
})
|
|
});
|
|
}
|
|
|
|
pub fn reset() {
|
|
FAILED.store(false, Ordering::Relaxed);
|
|
}
|
|
pub fn failed() -> bool {
|
|
FAILED.load(Ordering::Relaxed)
|
|
}
|