222 lines
6.7 KiB
Rust
222 lines
6.7 KiB
Rust
// This Source Code Form is subject to the terms of the Mozilla Public
|
|
// License, v. 2.0. If a copy of the MPL was not distributed with this
|
|
// file, You can obtain one at https://mozilla.org/MPL/2.0/.
|
|
|
|
//! Custom Python memory allocators.
|
|
|
|
#[cfg(feature = "jemalloc-sys")]
|
|
use jemalloc_sys as jemallocffi;
|
|
use libc::{c_void, size_t};
|
|
use python3_sys as pyffi;
|
|
use std::alloc;
|
|
use std::collections::HashMap;
|
|
#[cfg(feature = "jemalloc-sys")]
|
|
use std::ptr::null_mut;
|
|
|
|
const MIN_ALIGN: usize = 16;
|
|
|
|
type RawAllocatorState = HashMap<*mut u8, alloc::Layout>;
|
|
|
|
/// Holds state for the raw memory allocator.
|
|
///
|
|
/// Ideally we wouldn't need to track state. But Rust's dealloc() API
|
|
/// requires passing in a Layout that matches the allocation. This means
|
|
/// we need to track the Layout for each allocation. This data structure
|
|
/// facilitates that.
|
|
///
|
|
/// TODO HashMap isn't thread safe and the Python raw allocator doesn't
|
|
/// hold the GIL. So we need a thread safe map or a mutex guarding access.
|
|
pub struct RawAllocator {
|
|
pub allocator: pyffi::PyMemAllocatorEx,
|
|
_state: Box<RawAllocatorState>,
|
|
}
|
|
|
|
extern "C" fn raw_rust_malloc(ctx: *mut c_void, size: size_t) -> *mut c_void {
|
|
// PyMem_RawMalloc()'s docs say: Requesting zero bytes returns a distinct
|
|
// non-NULL pointer if possible, as if PyMem_RawMalloc(1) had been called
|
|
// instead.
|
|
let size = match size {
|
|
0 => 1,
|
|
val => val,
|
|
};
|
|
|
|
unsafe {
|
|
let state = ctx as *mut RawAllocatorState;
|
|
let layout = alloc::Layout::from_size_align_unchecked(size, MIN_ALIGN);
|
|
let res = alloc::alloc(layout);
|
|
|
|
(*state).insert(res, layout);
|
|
|
|
//println!("allocated {} bytes to {:?}", size, res);
|
|
res as *mut c_void
|
|
}
|
|
}
|
|
|
|
extern "C" fn raw_rust_calloc(ctx: *mut c_void, nelem: size_t, elsize: size_t) -> *mut c_void {
|
|
// PyMem_RawCalloc()'s docs say: Requesting zero elements or elements of
|
|
// size zero bytes returns a distinct non-NULL pointer if possible, as if
|
|
// PyMem_RawCalloc(1, 1) had been called instead.
|
|
let size = match nelem * elsize {
|
|
0 => 1,
|
|
val => val,
|
|
};
|
|
|
|
unsafe {
|
|
let state = ctx as *mut RawAllocatorState;
|
|
let layout = alloc::Layout::from_size_align_unchecked(size, MIN_ALIGN);
|
|
let res = alloc::alloc_zeroed(layout);
|
|
|
|
(*state).insert(res, layout);
|
|
|
|
//println!("zero allocated {} bytes to {:?}", size, res);
|
|
|
|
res as *mut c_void
|
|
}
|
|
}
|
|
|
|
extern "C" fn raw_rust_realloc(
|
|
ctx: *mut c_void,
|
|
ptr: *mut c_void,
|
|
new_size: size_t,
|
|
) -> *mut c_void {
|
|
//println!("reallocating {:?} to {} bytes", ptr as *mut u8, new_size);
|
|
|
|
// PyMem_RawRealloc()'s docs say: If p is NULL, the call is equivalent to
|
|
// PyMem_RawMalloc(n); else if n is equal to zero, the memory block is
|
|
// resized but is not freed, and the returned pointer is non-NULL.
|
|
if ptr.is_null() {
|
|
return raw_rust_malloc(ctx, new_size);
|
|
}
|
|
|
|
let new_size = match new_size {
|
|
0 => 1,
|
|
val => val,
|
|
};
|
|
|
|
unsafe {
|
|
let state = ctx as *mut RawAllocatorState;
|
|
let layout = alloc::Layout::from_size_align_unchecked(new_size, MIN_ALIGN);
|
|
|
|
let key = ptr as *mut u8;
|
|
let old_layout = (*state)
|
|
.remove(&key)
|
|
.expect("original memory address not tracked");
|
|
|
|
let res = alloc::realloc(ptr as *mut u8, old_layout, new_size);
|
|
|
|
(*state).insert(res, layout);
|
|
|
|
res as *mut c_void
|
|
}
|
|
}
|
|
|
|
extern "C" fn raw_rust_free(ctx: *mut c_void, ptr: *mut c_void) {
|
|
if ptr.is_null() {
|
|
return;
|
|
}
|
|
|
|
//println!("freeing {:?}", ptr as *mut u8);
|
|
unsafe {
|
|
let state = ctx as *mut RawAllocatorState;
|
|
|
|
let key = ptr as *mut u8;
|
|
let layout = (*state)
|
|
.get(&key)
|
|
.expect(format!("could not find allocated memory record: {:?}", key).as_str());
|
|
|
|
alloc::dealloc(key, *layout);
|
|
(*state).remove(&key);
|
|
}
|
|
}
|
|
|
|
pub fn make_raw_rust_memory_allocator() -> RawAllocator {
|
|
// We need to allocate the HashMap on the heap so the pointer doesn't refer
|
|
// to the stack. We rebox and add the Box to our struct so lifetimes are
|
|
// managed.
|
|
let alloc = Box::new(HashMap::<*mut u8, alloc::Layout>::new());
|
|
let state = Box::into_raw(alloc);
|
|
|
|
let allocator = pyffi::PyMemAllocatorEx {
|
|
ctx: state as *mut c_void,
|
|
malloc: Some(raw_rust_malloc),
|
|
calloc: Some(raw_rust_calloc),
|
|
realloc: Some(raw_rust_realloc),
|
|
free: Some(raw_rust_free),
|
|
};
|
|
|
|
RawAllocator {
|
|
allocator,
|
|
_state: unsafe { Box::from_raw(state) },
|
|
}
|
|
}
|
|
|
|
// Now let's define a raw memory allocator that interfaces directly with jemalloc.
|
|
// This avoids the overhead of going through Rust's allocation layer.
|
|
|
|
#[cfg(feature = "jemalloc-sys")]
|
|
extern "C" fn raw_jemalloc_malloc(_ctx: *mut c_void, size: size_t) -> *mut c_void {
|
|
// PyMem_RawMalloc()'s docs say: Requesting zero bytes returns a distinct
|
|
// non-NULL pointer if possible, as if PyMem_RawMalloc(1) had been called
|
|
// instead.
|
|
let size = match size {
|
|
0 => 1,
|
|
val => val,
|
|
};
|
|
|
|
unsafe { jemallocffi::mallocx(size, 0) }
|
|
}
|
|
|
|
#[cfg(feature = "jemalloc-sys")]
|
|
extern "C" fn raw_jemalloc_calloc(_ctx: *mut c_void, nelem: size_t, elsize: size_t) -> *mut c_void {
|
|
// PyMem_RawCalloc()'s docs say: Requesting zero elements or elements of
|
|
// size zero bytes returns a distinct non-NULL pointer if possible, as if
|
|
// PyMem_RawCalloc(1, 1) had been called instead.
|
|
let size = match nelem * elsize {
|
|
0 => 1,
|
|
val => val,
|
|
};
|
|
|
|
unsafe { jemallocffi::mallocx(size, jemallocffi::MALLOCX_ZERO) }
|
|
}
|
|
|
|
#[cfg(feature = "jemalloc-sys")]
|
|
extern "C" fn raw_jemalloc_realloc(
|
|
ctx: *mut c_void,
|
|
ptr: *mut c_void,
|
|
new_size: size_t,
|
|
) -> *mut c_void {
|
|
// PyMem_RawRealloc()'s docs say: If p is NULL, the call is equivalent to
|
|
// PyMem_RawMalloc(n); else if n is equal to zero, the memory block is
|
|
// resized but is not freed, and the returned pointer is non-NULL.
|
|
if ptr.is_null() {
|
|
return raw_jemalloc_malloc(ctx, new_size);
|
|
}
|
|
|
|
let new_size = match new_size {
|
|
0 => 1,
|
|
val => val,
|
|
};
|
|
|
|
unsafe { jemallocffi::rallocx(ptr, new_size, 0) }
|
|
}
|
|
|
|
#[cfg(feature = "jemalloc-sys")]
|
|
extern "C" fn raw_jemalloc_free(_ctx: *mut c_void, ptr: *mut c_void) {
|
|
if ptr.is_null() {
|
|
return;
|
|
}
|
|
|
|
unsafe { jemallocffi::dallocx(ptr, 0) }
|
|
}
|
|
|
|
#[cfg(feature = "jemalloc-sys")]
|
|
pub fn make_raw_jemalloc_allocator() -> pyffi::PyMemAllocatorEx {
|
|
pyffi::PyMemAllocatorEx {
|
|
ctx: null_mut(),
|
|
malloc: Some(raw_jemalloc_malloc),
|
|
calloc: Some(raw_jemalloc_calloc),
|
|
realloc: Some(raw_jemalloc_realloc),
|
|
free: Some(raw_jemalloc_free),
|
|
}
|
|
}
|