-
-
Notifications
You must be signed in to change notification settings - Fork 1.6k
/
memory.rs
75 lines (63 loc) · 2.79 KB
/
memory.rs
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
//! Utilities for dealing with memory allocations.
use std::alloc::GlobalAlloc;
use libc::{c_int, c_uint, size_t, uintptr_t};
// When debug_assertions is enabled, use Python's tracemalloc to track memory
// allocations. This is a useful feature for production use too, but has a
// potential performance impact and so would need additional benchmarking. In
// addition, these APIs are not part of the limited Python ABI Polars uses,
// though they are unchanged between 3.7 and 3.12.
#[cfg(not(target_os = "windows"))]
extern "C" {
fn PyTraceMalloc_Track(domain: c_uint, ptr: uintptr_t, size: size_t) -> c_int;
fn PyTraceMalloc_Untrack(domain: c_uint, ptr: uintptr_t) -> c_int;
}
// Windows has issues linking to the tracemalloc APIs, so the functionality is
// disabled. We have fake implementations just to make sure we don't have
// issues building.
#[cfg(target_os = "windows")]
#[allow(non_snake_case)]
fn PyTraceMalloc_Track(_domain: c_uint, _ptr: uintptr_t, _size: size_t) -> c_int {
-2
}
#[cfg(target_os = "windows")]
#[allow(non_snake_case)]
fn PyTraceMalloc_Untrack(_domain: c_uint, _ptr: uintptr_t) -> c_int {
-2
}
/// Allocations require a domain to identify them when registering with
/// tracemalloc. Following NumPy's lead, we just pick a random constant that is
/// unlikely to clash with anyone else.
const TRACEMALLOC_DOMAIN: c_uint = 36740582;
/// Wrap an existing allocator, and register allocations and frees with Python's
/// `tracemalloc`. Registration functionality is disabled on Windows.
pub struct TracemallocAllocator<A: GlobalAlloc> {
wrapped_alloc: A,
}
impl<A: GlobalAlloc> TracemallocAllocator<A> {
/// Wrap the allocator such that allocations are registered with
/// tracemalloc.
pub const fn new(wrapped_alloc: A) -> Self {
Self { wrapped_alloc }
}
}
unsafe impl<A: GlobalAlloc> GlobalAlloc for TracemallocAllocator<A> {
unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
let result = self.wrapped_alloc.alloc(layout);
PyTraceMalloc_Track(TRACEMALLOC_DOMAIN, result as uintptr_t, layout.size());
result
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
PyTraceMalloc_Untrack(TRACEMALLOC_DOMAIN, ptr as uintptr_t);
self.wrapped_alloc.dealloc(ptr, layout)
}
unsafe fn alloc_zeroed(&self, layout: std::alloc::Layout) -> *mut u8 {
let result = self.wrapped_alloc.alloc_zeroed(layout);
PyTraceMalloc_Track(TRACEMALLOC_DOMAIN, result as uintptr_t, layout.size());
result
}
unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
let result = self.wrapped_alloc.realloc(ptr, layout, new_size);
PyTraceMalloc_Track(TRACEMALLOC_DOMAIN, result as uintptr_t, new_size);
result
}
}