Allocators & Measuring Memory

use std::alloc::Layout;

#[allow(dead_code)]
struct Order {
    id: u64,
    qty: u32,
    express: bool,
}

fn main() {
    let order = Layout::new::<Order>();
    println!("Box<Order>:          size {:>3}, align {}", order.size(), order.align());

    let ids = Layout::array::<u32>(10).unwrap();
    println!("Vec<u32> of cap 10:  size {:>3}, align {}", ids.size(), ids.align());

    let bytes = Layout::from_size_align(100, 1).unwrap();
    println!("String of cap 100:   size {:>3}, align {}", bytes.size(), bytes.align());
}

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};

struct Counting;
static ALLOCS: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        ALLOCS.fetch_add(1, Relaxed);
        unsafe { System.alloc(layout) }
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) }
    }
}

#[global_allocator]
static GLOBAL: Counting = Counting;

fn allocs<R>(f: impl FnOnce() -> R) -> usize {
    let before = ALLOCS.load(Relaxed);
    let r = f();
    let n = ALLOCS.load(Relaxed) - before;
    drop(r);
    n
}

fn main() {
    let rows = [
        ("Vec::new()", allocs(|| Vec::<u64>::new())),
        ("String::new()", allocs(String::new)),
        ("Box::new(())", allocs(|| Box::new(()))),
        ("[0u8; 4096] on the stack", allocs(|| [0u8; 4096])),
        ("\"literal\"", allocs(|| "literal")),
        ("Box::new(42)", allocs(|| Box::new(42))),
        ("String::from(\"x\")", allocs(|| String::from("x"))),
        ("format!(\"{}-{}\", 1, 2)", allocs(|| format!("{}-{}", 1, 2))),
        ("vec![0; 1000]", allocs(|| vec![0u8; 1000])),
        ("(0..1000).collect::<Vec<_>>()", allocs(|| (0..1000).collect::<Vec<u32>>())),
        ("push 1000 times, no capacity", allocs(|| { let mut v = Vec::new(); for i in 0..1000 { v.push(i) } v })),
    ];
    for (what, n) in rows {
        println!("{what:<32} {n}");
    }
}

use mimalloc::MiMalloc;

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

fn main() {
    // Every Box, Vec, String, and HashMap in the program now uses mimalloc.
    let v: Vec<u64> = (0..1_000).collect();
    println!("{}", v.len());
}

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};

struct Tracking;
static CURRENT: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Tracking {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let now = CURRENT.fetch_add(layout.size(), Relaxed) + layout.size();
        PEAK.fetch_max(now, Relaxed);
        unsafe { System.alloc(layout) }
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        CURRENT.fetch_sub(layout.size(), Relaxed);
        unsafe { System.dealloc(ptr, layout) }
    }
}

#[global_allocator]
static GLOBAL: Tracking = Tracking;

/// Runs `f` and returns its result plus the peak bytes allocated while it ran.
fn peak_of<R>(f: impl FnOnce() -> R) -> (R, usize) {
    let base = CURRENT.load(Relaxed);
    PEAK.store(base, Relaxed);
    let r = f();
    (r, PEAK.load(Relaxed) - base)
}

fn main() {
    let input: String = (0..100_000).map(|i| format!("{i}\n")).collect();

    // Version 1: collect every line into a Vec<String>, then sum.
    let (sum1, peak1) = peak_of(|| {
        let lines: Vec<String> = input.lines().map(String::from).collect();
        lines.iter().map(|l| l.parse::<u64>().unwrap()).sum::<u64>()
    });

    // Version 2: stream over borrowed lines.
    let (sum2, peak2) = peak_of(|| input.lines().map(|l| l.parse::<u64>().unwrap()).sum::<u64>());

    println!("collect then sum: sum {sum1}, peak {peak1} bytes");
    println!("stream:           sum {sum2}, peak {peak2} bytes");
}

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering::Relaxed};

struct Histogram;
static ENABLED: AtomicBool = AtomicBool::new(false);
// Buckets: <=16, <=64, <=256, <=4096, larger.
static BUCKETS: [AtomicUsize; 5] = [const { AtomicUsize::new(0) }; 5];

unsafe impl GlobalAlloc for Histogram {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        if ENABLED.load(Relaxed) {
            let i = match layout.size() {
                0..=16 => 0,
                17..=64 => 1,
                65..=256 => 2,
                257..=4096 => 3,
                _ => 4,
            };
            BUCKETS[i].fetch_add(1, Relaxed);
        }
        unsafe { System.alloc(layout) }
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) }
    }
}

#[global_allocator]
static GLOBAL: Histogram = Histogram;

struct Event {
    user: String,
    tags: Vec<String>,
}

fn main() {
    ENABLED.store(true, Relaxed);
    let events: Vec<Event> = (0..10_000)
        .map(|i| Event { user: format!("u{}", i % 500), tags: vec!["web".to_string(), "eu".to_string()] })
        .collect();
    ENABLED.store(false, Relaxed);

    for (label, bucket) in ["<=16", "<=64", "<=256", "<=4096", "larger"].iter().zip(&BUCKETS) {
        println!("{label:>7}: {}", bucket.load(Relaxed));
    }
    println!("events: {}", events.len());
}

use std::fs;

fn rss_mib() -> u64 {
    let statm = fs::read_to_string("/proc/self/statm").unwrap();
    let pages: u64 = statm.split_whitespace().nth(1).unwrap().parse().unwrap();
    pages * 4096 / (1024 * 1024)
}

fn main() {
    println!("start:            {} MiB", rss_mib());
    let mut keep = Vec::new();
    let mut spike: Vec<Vec<u8>> = Vec::new();
    for i in 0..200_000 {
        spike.push(vec![1u8; 512]);
        if i % 100 == 0 {
            keep.push(vec![2u8; 512]); // a few long-lived blocks among the spike
        }
    }
    println!("during spike:     {} MiB", rss_mib());
    drop(spike);
    println!("after the spike:  {} MiB (live data is only ~1 MiB)", rss_mib());
    println!("kept {} blocks", keep.len());
}

fn read_payload(declared_len: usize) -> Result<Vec<u8>, String> {
    let mut buf = Vec::new();
    buf.try_reserve_exact(declared_len).map_err(|e| format!("cannot allocate {declared_len} bytes: {e}"))?;
    buf.resize(declared_len, 0); // stands in for reading from the socket
    Ok(buf)
}

fn main() {
    match read_payload(1024) {
        Ok(buf) => println!("ok: {} bytes", buf.len()),
        Err(e) => println!("error: {e}"),
    }
    match read_payload(usize::MAX / 2) {
        Ok(buf) => println!("ok: {} bytes", buf.len()),
        Err(e) => println!("error: {e}"),
    }
}

use std::sync::atomic::{AtomicUsize, Ordering};

struct MemoryBudget {
    limit: usize,
    used: AtomicUsize,
}

struct Reservation<'a> {
    budget: &'a MemoryBudget,
    bytes: usize,
}

impl MemoryBudget {
    const fn new(limit: usize) -> Self {
        Self { limit, used: AtomicUsize::new(0) }
    }

    fn try_reserve(&self, bytes: usize) -> Option<Reservation<'_>> {
        self.used
            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |used| used.checked_add(bytes).filter(|&total| total <= self.limit))
            .ok()
            .map(|_| Reservation { budget: self, bytes })
    }
}

impl Drop for Reservation<'_> {
    fn drop(&mut self) {
        self.budget.used.fetch_sub(self.bytes, Ordering::AcqRel); // released when the request finishes
    }
}

static BUDGET: MemoryBudget = MemoryBudget::new(100 * 1024 * 1024);

fn start_upload(name: &str, size: usize) -> Option<Reservation<'static>> {
    let r = BUDGET.try_reserve(size);
    let used = BUDGET.used.load(Ordering::Acquire) / (1024 * 1024);
    println!("{name:<8} {:>3} MiB: {} (in use: {used} MiB)", size / (1024 * 1024), if r.is_some() { "accepted" } else { "rejected, try later" });
    r
}

fn main() {
    let a = start_upload("video-a", 60 * 1024 * 1024);
    let _b = start_upload("video-b", 30 * 1024 * 1024);
    let _c = start_upload("video-c", 20 * 1024 * 1024);
    drop(a); // video-a finished
    let _d = start_upload("video-d", 20 * 1024 * 1024);
}

use std::collections::HashMap;
use std::hash::{BuildHasher, RandomState};
use std::ops::Range;

/// Stores many strings in one buffer. Handles are indices, not pointers.
#[derive(Default)]
struct StrArena {
    bytes: String,                  // every string's bytes, back to back
    spans: Vec<Range<usize>>,       // handle -> position in `bytes`
    lookup: HashMap<u64, Vec<u32>>, // content hash -> handles, for interning
    hasher: RandomState,
}

impl StrArena {
    fn intern(&mut self, s: &str) -> u32 {
        let h = self.hasher.hash_one(s);
        if let Some(ids) = self.lookup.get(&h) {
            if let Some(&id) = ids.iter().find(|&&id| self.get(id) == s) {
                return id; // already stored: no allocation at all
            }
        }
        let start = self.bytes.len();
        self.bytes.push_str(s);
        self.spans.push(start..self.bytes.len());
        let id = (self.spans.len() - 1) as u32;
        self.lookup.entry(h).or_default().push(id);
        id
    }

    fn get(&self, id: u32) -> &str {
        &self.bytes[self.spans[id as usize].clone()]
    }
}

fn main() {
    let mut arena = StrArena::default();
    let mut name = String::new(); // one reused scratch buffer instead of format! per event
    let events: Vec<u32> = (0..10_000)
        .map(|i| {
            use std::fmt::Write;
            name.clear();
            write!(name, "u{}", i % 500).unwrap();
            arena.intern(&name)
        })
        .collect();

    println!("events: {}, distinct users stored: {}", events.len(), arena.spans.len());
    println!("all user bytes live in one buffer of {} bytes", arena.bytes.len());
    println!("event 1234 is {}", arena.get(events[1234]));
}

use std::ops::{Deref, DerefMut};
use std::sync::Mutex;

struct BufferPool {
    free: Mutex<Vec<Vec<u8>>>,
}

struct PooledBuf<'a> {
    pool: &'a BufferPool,
    buf: Vec<u8>,
}

impl BufferPool {
    fn get(&self) -> PooledBuf<'_> {
        let buf = self.free.lock().unwrap().pop().unwrap_or_else(|| Vec::with_capacity(64 * 1024));
        PooledBuf { pool: self, buf }
    }
    fn idle(&self) -> usize {
        self.free.lock().unwrap().len()
    }
}

impl Drop for PooledBuf<'_> {
    fn drop(&mut self) {
        let mut buf = std::mem::take(&mut self.buf);
        buf.clear(); // keep the capacity, forget the contents
        self.pool.free.lock().unwrap().push(buf);
    }
}

impl Deref for PooledBuf<'_> {
    type Target = Vec<u8>;
    fn deref(&self) -> &Vec<u8> {
        &self.buf
    }
}

impl DerefMut for PooledBuf<'_> {
    fn deref_mut(&mut self) -> &mut Vec<u8> {
        &mut self.buf
    }
}

fn handle(pool: &BufferPool, request: &str) -> usize {
    let mut buf = pool.get();
    buf.extend_from_slice(b"HTTP/1.1 200 OK\r\n\r\n");
    buf.extend_from_slice(request.to_uppercase().as_bytes());
    buf.len()
}

fn main() {
    let pool = BufferPool { free: Mutex::new(Vec::new()) };
    for req in ["hello", "pooled", "buffers"] {
        let len = handle(&pool, req);
        let first_ptr = pool.get().as_ptr() as usize; // borrow one to see which buffer comes back
        println!("{req}: {len} bytes, idle buffers {}, reused same buffer: {}", pool.idle(), first_ptr == pool.get().as_ptr() as usize);
    }
}

use std::alloc::{GlobalAlloc, Layout, System};
use std::fmt::Write as _;
use std::io::Write as _;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};

struct Counting;
static ALLOCS: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        ALLOCS.fetch_add(1, Relaxed);
        unsafe { System.alloc(layout) }
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) }
    }
}

#[global_allocator]
static GLOBAL: Counting = Counting;

fn main() {
    let (user, latency_ms) = ("ana", 42);

    // 1. format!: a new String per call.
    let before = ALLOCS.load(Relaxed);
    let keys: Vec<usize> = (0..1000).map(|i| format!("{user}:{latency_ms}:{i}").len()).collect();
    println!("format!:          {} allocations", ALLOCS.load(Relaxed) - before - 1);

    // 2. One String reused with clear() + write!.
    let before = ALLOCS.load(Relaxed);
    let mut line = String::with_capacity(64);
    for i in 0..1000 {
        line.clear();
        write!(line, "{user}:{latency_ms}:{i}").unwrap();
    }
    println!("reused String:    {} allocation", ALLOCS.load(Relaxed) - before);

    // 3. A fixed stack buffer: no heap at all.
    let before = ALLOCS.load(Relaxed);
    let mut total = 0;
    for i in 0..1000 {
        let mut stack = [0u8; 64];
        let mut cursor = &mut stack[..];
        write!(cursor, "{user}:{latency_ms}:{i}").unwrap();
        total += 64 - cursor.len();
    }
    println!("stack buffer:     {} allocations", ALLOCS.load(Relaxed) - before);
    println!("same lengths: {}", total == keys.iter().sum::<usize>());
}

use smallvec::{SmallVec, smallvec};

struct Event {
    tags: SmallVec<[&'static str; 4]>, // no allocation for up to 4 tags
}

fn main() {
    let e = Event { tags: smallvec!["web", "eu"] };
    println!("{} tags, spilled to heap: {}", e.tags.len(), e.tags.spilled());
}

#[cfg(feature = "dhat-heap")]
#[global_allocator]
static ALLOC: dhat::Alloc = dhat::Alloc;

fn main() {
    #[cfg(feature = "dhat-heap")]
    let _profiler = dhat::Profiler::new_heap(); // writes dhat-heap.json when dropped

    run_the_service();
}

use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::Cell;

// Count per thread, so tests running in parallel don't disturb each other.
thread_local!(static ALLOCS: Cell<usize> = const { Cell::new(0) });

struct Counting;

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        ALLOCS.with(|c| c.set(c.get() + 1));
        unsafe { System.alloc(layout) }
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) }
    }
}

#[global_allocator]
static GLOBAL: Counting = Counting;

fn allocations_in(f: impl FnOnce()) -> usize {
    let before = ALLOCS.with(Cell::get);
    f();
    ALLOCS.with(Cell::get) - before
}

/// Parses "key=value;key=value" and sums the values of keys starting with `prefix`.
pub fn sum_matching(input: &str, prefix: &str) -> u64 {
    input
        .split(';')
        .filter_map(|pair| pair.split_once('='))
        .filter(|(k, _)| k.starts_with(prefix))
        .filter_map(|(_, v)| v.parse::<u64>().ok())
        .sum()
}

#[test]
fn sum_matching_is_correct() {
    assert_eq!(sum_matching("cpu.user=3;cpu.sys=4;mem=100", "cpu."), 7);
}

#[test]
fn sum_matching_does_not_allocate() {
    let input = "cpu.user=3;cpu.sys=4;mem=100;".repeat(100);
    let n = allocations_in(|| {
        std::hint::black_box(sum_matching(&input, "cpu."));
    });
    assert_eq!(n, 0, "sum_matching allocated {n} times");
}