//! zstd allocates through c malloc, which this doesn't see. the zstd rows count only the rust side.
//! run with `cargo bench -p dasl-star --bench alloc [filter]`
#[allow(dead_code)]
mod common;
use std::alloc::{GlobalAlloc, Layout, System};
use std::hint::black_box;
use std::io::{self, BufReader};
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering};
struct CountingAlloc;
#[global_allocator]
static GLOBAL: CountingAlloc = CountingAlloc;
static MEASURING: AtomicBool = AtomicBool::new(false);
static ALLOCS: AtomicU64 = AtomicU64::new(0);
static REALLOCS: AtomicU64 = AtomicU64::new(0);
static TOTAL_BYTES: AtomicU64 = AtomicU64::new(0);
static LIVE_BYTES: AtomicI64 = AtomicI64::new(0);
static PEAK_BYTES: AtomicI64 = AtomicI64::new(0);
fn grow(size: usize) {
TOTAL_BYTES.fetch_add(size as u64, Ordering::Relaxed);
let live = LIVE_BYTES.fetch_add(size as i64, Ordering::Relaxed) + size as i64;
PEAK_BYTES.fetch_max(live, Ordering::Relaxed);
}
unsafe impl GlobalAlloc for CountingAlloc {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
let ptr = unsafe { System.alloc(layout) };
if !ptr.is_null() && MEASURING.load(Ordering::Relaxed) {
ALLOCS.fetch_add(1, Ordering::Relaxed);
grow(layout.size());
}
ptr
}
unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
let ptr = unsafe { System.alloc_zeroed(layout) };
if !ptr.is_null() && MEASURING.load(Ordering::Relaxed) {
ALLOCS.fetch_add(1, Ordering::Relaxed);
grow(layout.size());
}
ptr
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
if MEASURING.load(Ordering::Relaxed) {
LIVE_BYTES.fetch_sub(layout.size() as i64, Ordering::Relaxed);
}
unsafe { System.dealloc(ptr, layout) };
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
let new_ptr = unsafe { System.realloc(ptr, layout, new_size) };
if !new_ptr.is_null() && MEASURING.load(Ordering::Relaxed) {
REALLOCS.fetch_add(1, Ordering::Relaxed);
match new_size.checked_sub(layout.size()) {
Some(more) => grow(more),
None => {
LIVE_BYTES.fetch_sub((layout.size() - new_size) as i64, Ordering::Relaxed);
}
}
}
new_ptr
}
}
type Run<'a> = Box<dyn Fn() -> Stats + 'a>;
struct Stats {
allocs: u64,
reallocs: u64,
total_bytes: u64,
peak_bytes: u64,
}
fn measure<T>(f: impl FnOnce() -> T) -> Stats {
for counter in [&ALLOCS, &REALLOCS, &TOTAL_BYTES] {
counter.store(0, Ordering::SeqCst);
}
LIVE_BYTES.store(0, Ordering::SeqCst);
PEAK_BYTES.store(0, Ordering::SeqCst);
MEASURING.store(true, Ordering::SeqCst);
black_box(f());
MEASURING.store(false, Ordering::SeqCst);
Stats {
allocs: ALLOCS.load(Ordering::SeqCst),
reallocs: REALLOCS.load(Ordering::SeqCst),
total_bytes: TOTAL_BYTES.load(Ordering::SeqCst),
peak_bytes: PEAK_BYTES.load(Ordering::SeqCst).max(0) as u64,
}
}
fn main() {
let filter = std::env::args().skip(1).find(|arg| !arg.starts_with('-'));
println!(
"{:<28} {:>9} {:>9} {:>14} {:>12}",
"bench", "allocs", "reallocs", "total_bytes", "peak_bytes"
);
for w in common::workloads() {
let plain = common::encode(&w);
let compressed = common::encode_zstd(&w);
let rows: [(&str, Run); 5] = [
(
"write/vec",
Box::new(|| {
let mut out = Vec::with_capacity(plain.len());
measure(|| common::write(&mut out, &w))
}),
),
(
"write/zstd",
Box::new(|| {
measure(|| {
let mut out = zstd::Encoder::new(io::sink(), common::ZSTD_LEVEL).unwrap();
common::write(&mut out, &w);
out.finish().unwrap()
})
}),
),
(
"read/slice",
Box::new(|| measure(|| common::read(&plain[..]))),
),
(
"read/owned",
Box::new(|| measure(|| common::read_owned(&plain[..]))),
),
(
"read/zstd",
Box::new(|| {
measure(|| {
let decoder = zstd::Decoder::new(&compressed[..]).unwrap();
common::read(BufReader::new(decoder))
})
}),
),
];
for (name, run) in rows {
let name = format!("{}/{name}", w.name);
if filter
.as_ref()
.is_some_and(|pat| !name.contains(pat.as_str()))
{
continue;
}
let s = run();
println!(
"{name:<28} {:>9} {:>9} {:>14} {:>12}",
s.allocs, s.reallocs, s.total_bytes, s.peak_bytes
);
}
}
}