horizon-ctl memory: the live heap by startup step, span and job

A build with the alloc-census feature runs mimalloc with a tag in front of
each block: the startup step, timed span or pool job that allocated it, else
its thread's name. horizon-ctl memory then lists the live heap by tag
(heap_by_tag_mb). heaptrack's interpreter ran out of 7 GB on the game; this
needs nothing beside it. Without the feature nothing changes: the tags are a
constant None.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

crates/horizon-client-render/src/spans.rs

@@ -61,6 +61,7 @@ fn cpu_if_enabled() -> Option<Duration> {
 /// Run `f`, recording its duration under `name` when spans are enabled.
 #[inline]
 pub fn time<T>(name: &'static str, f: impl FnOnce() -> T) -> T {
+    let _heap = horizon_threads::heap_tag(name);
     if !recording_here() {
         return f();
     }
@@ -291,9 +292,15 @@ pub mod startup {
     pub struct Step {
         name: &'static str,
         started: Instant,
+        traced: bool,
+        /// A heap-census build's tag for the step's allocations.
+        _heap: Option<horizon_threads::heap_census::Tag>,
     }
     impl Drop for Step {
         fn drop(&mut self) {
+            if !self.traced {
+                return;
+            }
             let end = Instant::now();
             write(format_args!(
                 r#""ph":"X","name":"{}","ts":{},"dur":{}"#,
@@ -307,9 +314,12 @@ pub mod startup {
     /// Starts a step named `name` (a literal: no quotes or backslashes).
     #[inline]
     pub fn step(name: &'static str) -> Option<Step> {
-        on().then(|| Step {
+        let traced = on();
+        (traced || horizon_threads::HEAP_CENSUS).then(|| Step {
             name,
             started: Instant::now(),
+            traced,
+            _heap: horizon_threads::heap_tag(name),
         })
     }
 

crates/horizon-client-render/src/streaming.rs

@@ -774,12 +774,23 @@ impl<G: ChunkGenerator> Source for Jobs<G> {
             };
             (kind, crate::spans::thread_cpu().unwrap_or_default())
         });
+        let _tag = horizon_threads::heap_tag(match &job.job {
+            Job::Generate(..) => "stream.generate",
+            Job::Decorate(..) => "stream.decorate",
+            Job::Finish(..) => "stream.finish",
+            Job::Light(..) => "stream.light",
+            Job::Mesh(..) => "stream.mesh",
+            _ => "stream.other",
+        });
         let done = self.locals.with(
             worker,
-            || WorkerLocal {
-                worker: self.generator.worker(),
-                faces: None,
-                mesh_scratch: MeshScratch::default(),
+            || {
+                let _tag = horizon_threads::heap_tag("stream.worker_state");
+                WorkerLocal {
+                    worker: self.generator.worker(),
+                    faces: None,
+                    mesh_scratch: MeshScratch::default(),
+                }
             },
             |local| self.run_job(job.job, local),
         );

crates/horizon-client/Cargo.toml

@@ -117,6 +117,9 @@ rng-trace = ["horizon-core/rng-trace", "horizon-client-lib/rng-trace"]
 # Rust's allocations go to the system heap instead of mimalloc, where heap profilers that replace
 # malloc (heaptrack, valgrind's massif) see them. For profiling only: mimalloc is faster.
 system-allocator = []
+# mimalloc with a tag on each block: `horizon-ctl memory` lists the live heap by startup step and
+# thread (`heap_by_tag_mb`). For measuring only: each block costs 16 bytes more.
+alloc-census = ["horizon-threads/alloc-census"]
 # Compatibility with earlier experiment commands: main already uses mimalloc.
 windows-mimalloc = []
 # Opt-in Windows heap experiment; see docs/windows-allocators.md.

crates/horizon-client/src/bin/horizon-ctl/commands.rs

@@ -221,7 +221,7 @@ pub const COMMANDS: &[Cmd] = &[
         "Replies protocol (MAJOR.MINOR.PATCH of the control protocol, docs/api/stability.md), release (the game's version, the string MCP serverInfo carries too) and features (the command groups the game answers, so a script can test for one before calling it).",
         ["horizon-ctl hello"], &[], mcp: true, local: false),
     cmd!("memory", Session, "memory", "what holds the game's memory: resident set, allocator, client and server columns by part, renderer, GPU memory by resource",
-        "Columns count blocks, biomes, light, heightmaps and records, each shared section once; server_columns are the single-player server thread's own levels. GPU memory by label needs Vulkan or DirectX 12. What the census doesn't reach is the resident set less its parts.", ["horizon-ctl memory"], &[], mcp: true, local: false),
+        "Columns count blocks, biomes, light, heightmaps and records, each shared section once; server_columns are the single-player server thread's own levels. GPU memory by label needs Vulkan or DirectX 12. What the census doesn't reach is the resident set less its parts. A build with the alloc-census feature adds heap_by_tag_mb: the live heap by startup step, timed span, chunk job and thread.", ["horizon-ctl memory"], &[], mcp: true, local: false),
     cmd!("worlds", Session, "worlds", "the saved worlds, newest first: folder, name, mode, last played, version",
         "folder is what open, world-rename, world-backup and world-delete take.", ["horizon-ctl worlds"], &[], mcp: true, local: false),
     cmd!("world-rename", Session, "world-rename FOLDER NAME", "rename a world, as Edit World does (its folder stays)",

crates/horizon-client/src/game/memory.rs

@@ -141,6 +141,15 @@ impl Game<'_> {
                 json!({ "active": active, "cached": cached })
             }),
             "renderer": renderer,
+            // A build with the `alloc-census` feature: the live heap by the
+            // startup step, span or pool job that allocated it, else its thread.
+            "heap_by_tag_mb": horizon_threads::HEAP_CENSUS.then(|| {
+                horizon_threads::heap_census::census()
+                    .into_iter()
+                    .take(60)
+                    .map(|(tag, bytes)| json!([tag, mb(bytes.max(0) as u64)]))
+                    .collect::<Vec<_>>()
+            }),
         })
     }
 }

crates/horizon-client/src/main.rs

@@ -1,9 +1,12 @@
 //! Horizon: native Rust client for Minecraft Java Edition 26.3 gameplay.
 #![forbid(unsafe_code)]
 
-#[cfg(not(feature = "system-allocator"))]
+#[cfg(not(any(feature = "system-allocator", feature = "alloc-census")))]
 #[global_allocator]
 static ALLOCATOR: horizon_threads::Allocator = horizon_threads::Allocator;
+#[cfg(all(feature = "alloc-census", not(feature = "system-allocator")))]
+#[global_allocator]
+static ALLOCATOR: horizon_threads::CensusAllocator = horizon_threads::CensusAllocator;
 
 /// The mod API level this game serves, kept in the program for the launcher
 /// to read without running it (`horizon_mods::api_level::MARK`). `main`

crates/horizon-threads/Cargo.toml

@@ -16,5 +16,9 @@ libc = "0.2"
 [target.'cfg(windows)'.dependencies]
 windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_Graphics_Gdi", "Win32_System_Console", "Win32_System_IO", "Win32_System_Memory", "Win32_System_ProcessStatus", "Win32_System_Threading", "Win32_System_WindowsProgramming", "Win32_UI_WindowsAndMessaging"] }
 
+[features]
+# Counts live heap bytes by startup step and thread (`CensusAllocator`); for measuring only.
+alloc-census = []
+
 [lints]
 workspace = true

crates/horizon-threads/src/heap_census.rs

@@ -0,0 +1,269 @@
+//! Who holds the heap, for `horizon-ctl memory` (feature `alloc-census`, for
+//! measuring only): mimalloc with a tag in front of each block, and the live
+//! bytes of each tag. A block's tag is what ran on its thread when it was
+//! allocated: the innermost [`tag`] (a startup step, a timed span or a pool
+//! job: `startup::time`, `spans::time` and `jobs` set them), else the
+//! thread's name. A block keeps its tag when another thread frees or grows
+//! it, so a column generated by a pool job and kept by the server thread
+//! counts for the job.
+//!
+//! Each block costs a header of 16 bytes (its alignment when larger) and an
+//! atomic add: heaptrack's interpreter needed more than 7 GB for the game
+//! (18 GB with line tables), this needs nothing beside it.
+use std::{
+    alloc::{GlobalAlloc, Layout},
+    cell::Cell,
+    sync::atomic::{AtomicIsize, AtomicPtr, AtomicUsize, Ordering::Relaxed},
+};
+
+const TAGS: usize = 512;
+/// Blocks allocated while this thread names itself, and past the table's end.
+const UNTAGGED: u32 = 0;
+const NONE: u32 = u32::MAX;
+
+static LIVE: [AtomicIsize; TAGS] = [const { AtomicIsize::new(0) }; TAGS];
+static NAME_PTR: [AtomicPtr<u8>; TAGS] = [const { AtomicPtr::new(std::ptr::null_mut()) }; TAGS];
+static NAME_LEN: [AtomicUsize; TAGS] = [const { AtomicUsize::new(0) }; TAGS];
+/// The next free slot; slot 0 is [`UNTAGGED`].
+static NEXT: AtomicUsize = AtomicUsize::new(1);
+
+thread_local! {
+    static STEP: Cell<u32> = const { Cell::new(NONE) };
+    static THREAD: Cell<u32> = const { Cell::new(NONE) };
+    /// Allocations this thread made before its tag settled: std names a
+    /// thread after its first few, so the name is read again until then.
+    static NAMING: Cell<u32> = const { Cell::new(0) };
+    static BUSY: Cell<bool> = const { Cell::new(false) };
+}
+
+/// The slot of `name`, registering it. A race may give one name two slots;
+/// [`census`] adds them up.
+fn slot(name: &'static str) -> u32 {
+    let n = NEXT.load(Relaxed).min(TAGS);
+    for i in 1..n {
+        let ptr = NAME_PTR[i].load(Relaxed);
+        if !ptr.is_null() && NAME_LEN[i].load(Relaxed) == name.len() {
+            // SAFETY: the slot holds a leaked or static string of that length.
+            let s = unsafe { std::slice::from_raw_parts(ptr, name.len()) };
+            if s == name.as_bytes() {
+                return i as u32;
+            }
+        }
+    }
+    let i = NEXT.fetch_add(1, Relaxed);
+    if i >= TAGS {
+        return UNTAGGED;
+    }
+    NAME_LEN[i].store(name.len(), Relaxed);
+    NAME_PTR[i].store(name.as_ptr().cast_mut(), Relaxed);
+    i as u32
+}
+
+/// The OS name of the calling thread, read without allocating.
+#[cfg(any(target_vendor = "apple", target_os = "linux"))]
+fn thread_name(buf: &mut [u8; 64]) -> &[u8] {
+    // SAFETY: writes at most `buf.len()` bytes, NUL included, into `buf`.
+    let ok = unsafe {
+        libc::pthread_getname_np(libc::pthread_self(), buf.as_mut_ptr().cast(), buf.len())
+    } == 0;
+    let len = buf.iter().position(|&b| b == 0).unwrap_or(0);
+    if ok { &buf[..len] } else { &[] }
+}
+#[cfg(not(any(target_vendor = "apple", target_os = "linux")))]
+fn thread_name(_: &mut [u8; 64]) -> &[u8] {
+    &[]
+}
+
+/// The calling thread's own tag: its name, looked up again for its first
+/// allocations. Allocations made while it registers a name are untagged.
+fn thread_tag() -> u32 {
+    let known = THREAD.get();
+    if known != NONE && NAMING.get() >= 256 {
+        return known;
+    }
+    if BUSY.replace(true) {
+        return UNTAGGED;
+    }
+    NAMING.set(NAMING.get().saturating_add(1));
+    let mut buf = [0u8; 64];
+    let name = match std::str::from_utf8(thread_name(&mut buf)) {
+        Ok("") | Err(_) => "unnamed thread",
+        Ok(name) => name,
+    };
+    let tag = if known != NONE && slot_name(known) == Some(name) {
+        known
+    } else {
+        // Thread names are few: each is leaked once.
+        let name: &'static str = match name {
+            "unnamed thread" => "unnamed thread",
+            name => Box::leak(name.to_owned().into_boxed_str()),
+        };
+        slot(name)
+    };
+    THREAD.set(tag);
+    BUSY.set(false);
+    tag
+}
+
+fn slot_name(i: u32) -> Option<&'static str> {
+    let i = i as usize;
+    let ptr = NAME_PTR.get(i)?.load(Relaxed);
+    if ptr.is_null() {
+        return None;
+    }
+    // SAFETY: registered slots hold leaked or static UTF-8 of that length.
+    Some(unsafe {
+        std::str::from_utf8_unchecked(std::slice::from_raw_parts(ptr, NAME_LEN[i].load(Relaxed)))
+    })
+}
+
+fn current() -> u32 {
+    match STEP.get() {
+        NONE => thread_tag(),
+        step => step,
+    }
+}
+
+/// While it lives, the calling thread's blocks are tagged `name`.
+pub struct Tag(u32);
+impl Drop for Tag {
+    fn drop(&mut self) {
+        STEP.set(self.0);
+    }
+}
+/// Tag the calling thread's blocks `name` until the guard drops.
+pub fn tag(name: &'static str) -> Tag {
+    let previous = STEP.get();
+    let was = BUSY.replace(true);
+    let tag = slot(name);
+    BUSY.set(was);
+    STEP.set(tag);
+    Tag(previous)
+}
+
+/// Live bytes by tag, largest first.
+pub fn census() -> Vec<(&'static str, isize)> {
+    let mut by_name: Vec<(&'static str, isize)> = Vec::new();
+    for i in 0..NEXT.load(Relaxed).min(TAGS) {
+        let bytes = LIVE[i].load(Relaxed);
+        let name = slot_name(i as u32).unwrap_or("untagged");
+        match by_name.iter_mut().find(|(n, _)| *n == name) {
+            Some(entry) => entry.1 += bytes,
+            None => by_name.push((name, bytes)),
+        }
+    }
+    by_name.sort_by_key(|(_, bytes)| std::cmp::Reverse(*bytes));
+    by_name
+}
+
+/// mimalloc with a tag in front of each block (see the module).
+pub struct CensusAllocator;
+
+/// The header in front of a block of `layout`: 16 bytes, or its alignment.
+fn header(layout: Layout) -> usize {
+    layout.align().max(16)
+}
+
+// SAFETY: forwards every call to mimalloc with the header added; the tag sits
+// in the header's last four bytes, inside the block mimalloc returned.
+unsafe impl GlobalAlloc for CensusAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let h = header(layout);
+        let Ok(outer) = Layout::from_size_align(layout.size() + h, layout.align().max(16)) else {
+            return std::ptr::null_mut();
+        };
+        // SAFETY: `outer` has a nonzero size.
+        let base = unsafe { mimalloc::MiMalloc.alloc(outer) };
+        if base.is_null() {
+            return base;
+        }
+        let tag = current();
+        LIVE[tag as usize].fetch_add(layout.size() as isize, Relaxed);
+        // SAFETY: `h` bytes precede the caller's block inside the allocation.
+        unsafe {
+            let block = base.add(h);
+            block.sub(4).cast::<u32>().write_unaligned(tag);
+            block
+        }
+    }
+    unsafe fn dealloc(&self, block: *mut u8, layout: Layout) {
+        let h = header(layout);
+        // SAFETY: `block` came from `alloc` with this layout.
+        unsafe {
+            let tag = block.sub(4).cast::<u32>().read_unaligned();
+            LIVE[tag as usize].fetch_sub(layout.size() as isize, Relaxed);
+            let outer =
+                Layout::from_size_align_unchecked(layout.size() + h, layout.align().max(16));
+            mimalloc::MiMalloc.dealloc(block.sub(h), outer);
+        }
+    }
+    unsafe fn realloc(&self, block: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        let h = header(layout);
+        // SAFETY: `block` came from `alloc` with this layout; the header moves with it.
+        unsafe {
+            let outer =
+                Layout::from_size_align_unchecked(layout.size() + h, layout.align().max(16));
+            let tag = block.sub(4).cast::<u32>().read_unaligned();
+            let base = mimalloc::MiMalloc.realloc(block.sub(h), outer, new_size + h);
+            if base.is_null() {
+                return base;
+            }
+            LIVE[tag as usize].fetch_add(new_size as isize - layout.size() as isize, Relaxed);
+            base.add(h)
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{CensusAllocator, census, tag};
+    use std::alloc::{GlobalAlloc, Layout};
+
+    fn live(name: &str) -> isize {
+        census().iter().find(|(n, _)| *n == name).map_or(0, |e| e.1)
+    }
+
+    #[test]
+    fn a_block_counts_for_its_step_until_another_thread_frees_it() {
+        let layout = Layout::from_size_align(1 << 20, 64).unwrap();
+        let block = {
+            let _tag = tag("test.step");
+            // SAFETY: a nonzero layout.
+            unsafe { CensusAllocator.alloc(layout) }
+        };
+        assert_eq!(live("test.step"), 1 << 20);
+        // SAFETY: `block` came from this allocator with `layout`.
+        let block = unsafe { CensusAllocator.realloc(block, layout, 3 << 20) } as usize;
+        assert_eq!(live("test.step"), 3 << 20, "growing keeps the step's tag");
+        let grown = Layout::from_size_align(3 << 20, 64).unwrap();
+        std::thread::spawn(move || {
+            let _other = tag("test.other");
+            // SAFETY: `block` came from this allocator with `grown`.
+            unsafe { CensusAllocator.dealloc(block as *mut u8, grown) };
+        })
+        .join()
+        .unwrap();
+        assert_eq!(live("test.step"), 0);
+        assert_eq!(live("test.other"), 0);
+    }
+
+    #[test]
+    fn without_a_step_a_block_counts_for_its_thread() {
+        let layout = Layout::from_size_align(4096, 8).unwrap();
+        let (name, block) = std::thread::Builder::new()
+            .name("census-test".into())
+            .spawn(move || {
+                // SAFETY: a nonzero layout.
+                let block = unsafe { CensusAllocator.alloc(layout) } as usize;
+                (std::thread::current().name().unwrap().to_owned(), block)
+            })
+            .unwrap()
+            .join()
+            .unwrap();
+        // Linux keeps 15 bytes of a thread's name; the test's is shorter.
+        assert_eq!(live(&name), 4096);
+        // SAFETY: `block` came from this allocator with `layout`.
+        unsafe { CensusAllocator.dealloc(block as *mut u8, layout) };
+        assert_eq!(live(&name), 0);
+    }
+}

crates/horizon-threads/src/jobs.rs

@@ -65,6 +65,10 @@ pub trait Source: Send + Sync {
     fn idle(&self, _worker: usize) {}
     /// Every worker is idle: free what the workers share.
     fn all_idle(&self) {}
+    /// What a heap-census build tags its jobs' allocations (`crate::heap_tag`).
+    fn label(&self) -> &'static str {
+        std::any::type_name::<Self>()
+    }
 }
 
 /// How long a worker waits with no job before it frees its buffers.
@@ -384,7 +388,11 @@ fn work(inner: &Inner, worker: usize) {
             }
             // A panicking job loses its own work, not the worker: the panic
             // hook has already reported it.
-            let _ = catch_unwind(AssertUnwindSafe(|| entry.source.run(worker)));
+            let _ = catch_unwind(AssertUnwindSafe(|| {
+                let _tag =
+                    crate::HEAP_CENSUS.then(|| crate::heap_census::tag(entry.source.label()));
+                entry.source.run(worker)
+            }));
             if far {
                 inner.far.fetch_sub(1, Ordering::SeqCst);
             }

crates/horizon-threads/src/lib.rs

@@ -4,6 +4,7 @@
 /// The game window's boot screen, before wgpu's first frame paints over it.
 #[cfg(target_os = "windows")]
 pub mod boot_paint;
+pub mod heap_census;
 pub mod jobs;
 pub mod memory_limit;
 
@@ -207,6 +208,18 @@ pub fn catch_stop_signals() {}
 /// 16% of a chunk's generation time), mimalloc keeps a heap per thread.
 pub use mimalloc::MiMalloc as Allocator;
 
+/// The global allocator of a build with the `alloc-census` feature: live
+/// heap bytes by startup step and thread, for `horizon-ctl memory`.
+pub use heap_census::CensusAllocator;
+/// Whether this build counts the heap by tag (feature `alloc-census`).
+pub const HEAP_CENSUS: bool = cfg!(feature = "alloc-census");
+/// Tag the calling thread's allocations `name` while the guard lives, in a
+/// build that counts the heap ([`HEAP_CENSUS`]); nothing otherwise.
+#[inline]
+pub fn heap_tag(name: &'static str) -> Option<heap_census::Tag> {
+    HEAP_CENSUS.then(|| heap_census::tag(name))
+}
+
 /// `mi_option_purge_delay` in mimalloc.h (index 15 in v2 and v3); the bindings
 /// don't name it.
 const PURGE_DELAY: libmimalloc_sys::mi_option_t = 15;