Skip to content

Instantly share code, notes, and snippets.

@danielhenrymantilla
Created December 15, 2021 00:50
Show Gist options
  • Save danielhenrymantilla/a75235cacd716dd5340e6d24f709e165 to your computer and use it in GitHub Desktop.
Save danielhenrymantilla/a75235cacd716dd5340e6d24f709e165 to your computer and use it in GitHub Desktop.
Poorman's alloc-instrumentation
struct MyAllocator;
#[global_allocator]
static A: MyAllocator = MyAllocator;
const _: () = {
use ::std::{
alloc::{
GlobalAlloc,
Layout,
System,
},
cell::Cell,
ops::Not as _,
sync::{
Arc,
atomic::{
self,
AtomicUsize,
},
},
};
use ::dashmap::DashMap;
type Trace = Arc<[(String, Option<String>, Option<u32>)]>;
::lazy_static::lazy_static! {
static ref TRACKED_ALLOC_SIZES: DashMap<Trace, usize> = DashMap::new();
static ref TRACKED_ALLOC_PTRS: DashMap<usize, Trace> = DashMap::new();
}
static TOTAL_SIZE: AtomicUsize = AtomicUsize::new(0);
#[track_caller]
fn get_trace() -> Trace {
let mut ret = Vec::new();
::backtrace::trace(|frame| {
// let ip = frame.ip();
// let symbol_address = frame.symbol_address();
// Resolve this instruction pointer to a symbol name
::backtrace::resolve_frame(frame, |symbol| {
ret.push((
symbol
.name()
.as_ref()
.map_or_else(
|| "<unresolved>".into(),
|it| it.to_string(),
)
,
symbol.filename().map(|it| it.display().to_string()),
symbol.lineno(),
));
});
// keep going to the next frame while:
ret.len() < 15
});
ret.into_boxed_slice().into()
}
fn not_reentrantly (f: impl FnOnce())
{
thread_local! {
static REENTRANT: Cell<bool> = false.into();
}
if REENTRANT.with(Cell::get).not() {
REENTRANT.with(|it| it.set(true));
let () = f();
REENTRANT.with(|it| it.set(false));
}
}
static MAX_MEMORY: AtomicUsize = AtomicUsize::new(usize::MAX);
#[ffi_export]
fn my_alloc_set_limit (bytes: usize)
{
MAX_MEMORY.store(dbg!(bytes), atomic::Ordering::Release);
}
#[ffi_export]
fn my_alloc_get_allocated ()
-> usize
{
TOTAL_SIZE.load(atomic::Ordering::Acquire)
}
// #[ffi_export]
fn my_alloc_dump_traced_alloc ()
{
(|| -> ::std::io::Result<()> {
use ::std::io::Write;
let mut entries = TRACKED_ALLOC_SIZES.iter().map(|guard| {
let (trace, &total_alloc_size) = guard.pair();
(trace.clone(), total_alloc_size)
}).collect::<Vec<_>>();
entries.sort_unstable_by_key(|&(_, sz)| sz);
let mut file = ::std::fs::File::create("allocs.trace")?;
for (trace, mut total_alloc_size) in entries.into_iter().rev() {
let size = {
let total_alloc_mb = total_alloc_size / (1024 * 1024);
if total_alloc_mb > 0 {
format!("{}.… MiB", total_alloc_mb)
} else {
let total_alloc_kb = total_alloc_size / 1024;
if total_alloc_kb > 0 {
format!("{}.… KiB", total_alloc_kb)
} else {
format!("{} bytes", total_alloc_size)
}
}
};
let () = write!(file, "\n{} from:\n", size)?;
for caller in &trace[..] {
let (fname, mb_file, mb_line) = &*caller;
let () = write!(
file,
concat!(
" {fname}\n",
" at {file}:{line}\n",
),
fname = fname,
file = mb_file.as_deref().unwrap_or("<unknown>"),
line = mb_line.unwrap_or(0),
)?;
}
}
let () = file.flush()?;
let () = file.sync_data()?;
Ok(())
})().unwrap();
}
unsafe
impl GlobalAlloc for MyAllocator {
unsafe
fn alloc (self: &'_ MyAllocator, layout: Layout)
-> *mut u8
{
let size = layout.size();
not_reentrantly(|| while {
let prev = TOTAL_SIZE.load(atomic::Ordering::Relaxed);
let new = prev.saturating_add(size);
if new > MAX_MEMORY.load(atomic::Ordering::Acquire) {
MAX_MEMORY.store(usize::MAX, atomic::Ordering::Relaxed);
my_alloc_dump_traced_alloc();
panic!("exhausted allowed memory!");
}
TOTAL_SIZE
.compare_exchange_weak(
prev,
new,
atomic::Ordering::AcqRel,
atomic::Ordering::Relaxed,
)
.is_err()
} {});
let ptr = System.alloc(layout);
not_reentrantly(|| if ptr.is_null() {
TOTAL_SIZE.fetch_sub(size, atomic::Ordering::Release);
} else {
let trace = get_trace();
*TRACKED_ALLOC_SIZES.entry(trace.clone()).or_insert(0)
+= size
;
TRACKED_ALLOC_PTRS.insert(ptr as usize, trace);
});
ptr
}
unsafe
fn dealloc (self: &'_ MyAllocator, ptr: *mut u8, layout: Layout)
{
not_reentrantly(|| if ptr.is_null().not() {
let size = layout.size();
TOTAL_SIZE.fetch_sub(size, atomic::Ordering::Release);
let (_, trace) = TRACKED_ALLOC_PTRS.remove(&(ptr as usize)).unwrap();
*TRACKED_ALLOC_SIZES.get_mut(&trace).unwrap()
-= size
;
});
System.dealloc(ptr, layout)
}
}
};
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment