turbo-tasks-malloc: report memory from mimalloc

@@ -129,8 +129,11 @@ impl EvictionControl {         evict     } -    /// Call after completing an eviction cycle. Seeds the memory floor with the-    /// post-eviction usage; later cycles lower it further as memory settles.+    /// Call after completing an eviction cycle and after mimalloc is cleaned up+    /// with [`TurboMalloc::collect`], since the freed memory is only reflected+    /// in [`TurboMalloc::memory_usage`] once it has been. Seeds the memory floor+    /// with the post-eviction usage; later cycles lower it further as memory+    /// settles.     pub(crate) fn record_eviction(&mut self) {         self.memory_floor = Some(TurboMalloc::memory_usage());     }

@@ -3028,8 +3028,6 @@ impl TurboTasksBackend {                                     // memory so racing with execution is as likely to save time as                                     // cost it.                                     self.storage.evict_after_snapshot(background_span.id());-                                    // Sample the post-eviction floor as the new baseline.-                                    eviction_control.record_eviction();                                     true                                 } else {                                     false@@ -3082,6 +3080,12 @@ impl TurboTasksBackend {                                 {                                     TurboMalloc::collect(true);                                 }++                                // Sample the new baseline after the collect above, which is what+                                // makes the evicted memory show up in `memory_usage`.+                                if ran_eviction {+                                    eviction_control.record_eviction();+                                }                             }                         }                     }

@@ -1,60 +1,140 @@-use std::{-    cell::UnsafeCell,-    ptr::NonNull,-    sync::atomic::{AtomicUsize, Ordering},-};+//! Allocation accounting.+//!+//! Every build tracks per-thread allocation totals, which the tracing layer reads through+//! [`allocation_counters`] to attribute allocations to spans.+//!+//! Builds without the `custom_allocator` feature additionally maintain a process-wide counter of+//! live bytes, which backs [`crate::TurboMalloc::memory_usage`]. With mimalloc that figure comes+//! from the allocator instead, so [`global`] is not compiled in: the atomic would otherwise be+//! contended by every thread on every allocation, and the thread-local buffering that makes it+//! affordable is inlined into every allocation site in the binary. +use std::{cell::UnsafeCell, ptr::NonNull};++#[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]+pub use self::global::get; use crate::AllocationCounters; -/// Tracks the current total amount of memory allocated through all the [ThreadLocalCounter]-/// instances.  This is an overestimate as individual threads 'preallocate' a [TARGET_BUFFER] bytes-/// to reduce the number of global synchronizations.  This means at any given time this might-/// overcount by up to [MAX_BUFFER] bytes for each thread.-static ALLOCATED: AtomicUsize = AtomicUsize::new(0);-const KB: usize = 1024;-/// When global counter is updates we will keep a thread-local buffer of this-/// size.-const TARGET_BUFFER: usize = 100 * KB;-/// When the thread-local buffer would exceed this size, we will update the-/// global counter.-const MAX_BUFFER: usize = 200 * KB;+/// The process-wide live-bytes counter, and the buffering that keeps updating it affordable.+///+/// Only compiled without the `custom_allocator` feature; see the module docs. Each thread holds+/// its buffer in its own [`ThreadLocalCounter`] and passes it in, so the counter's state lives in+/// exactly one place.+#[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]+mod global {+    use std::sync::atomic::{AtomicUsize, Ordering};++    /// Tracks the current total amount of memory allocated through all the+    /// [`super::ThreadLocalCounter`] instances.  This is an overestimate as individual threads+    /// 'preallocate' a [TARGET_BUFFER] bytes to reduce the number of global synchronizations.+    /// This means at any given time this might overcount by up to [MAX_BUFFER] bytes for each+    /// thread.+    static ALLOCATED: AtomicUsize = AtomicUsize::new(0);+    const KB: usize = 1024;+    /// When global counter is updates we will keep a thread-local buffer of this+    /// size.+    pub const TARGET_BUFFER: usize = 100 * KB;+    /// When the thread-local buffer would exceed this size, we will update the+    /// global counter.+    pub const MAX_BUFFER: usize = 200 * KB;++    /// Live bytes (allocations minus deallocations) across all threads.+    pub fn get() -> usize {+        ALLOCATED.load(Ordering::Relaxed)+    }++    /// Takes `size` from the global counter, refilling `buffer` while it is there.+    ///+    /// Kept out of the allocator's inlined hot path: the buffer means this runs about once per+    /// [`TARGET_BUFFER`] bytes rather than once per allocation.+    #[inline(never)]+    pub fn refill(buffer: &mut usize, size: usize) {+        debug_assert!(*buffer < size);+        let offset = size - *buffer + TARGET_BUFFER;+        *buffer = TARGET_BUFFER;+        ALLOCATED.fetch_add(offset, Ordering::Relaxed);+    }++    /// Returns everything buffered above [`TARGET_BUFFER`] to the global counter.+    #[inline(never)]+    pub fn flush_excess(buffer: &mut usize) {+        debug_assert!(*buffer > MAX_BUFFER);+        let offset = *buffer - TARGET_BUFFER;+        *buffer = TARGET_BUFFER;+        ALLOCATED.fetch_sub(offset, Ordering::Relaxed);+    }++    /// Returns everything buffered, for a thread that is going away.+    pub fn flush_all(buffer: &mut usize) {+        if *buffer > 0 {+            ALLOCATED.fetch_sub(*buffer, Ordering::Relaxed);+            *buffer = 0;+        }+    }++    impl super::ThreadLocalCounter {+        /// Charges `size` against this thread's buffer, refilling it from the global counter when+        /// it runs dry. Does nothing with `custom_allocator`, where there is no global+        /// counter.+        #[inline(always)]+        pub(super) fn buffered_add(&mut self, size: usize) {+            if self.buffer >= size {+                self.buffer -= size;+            } else {+                refill(&mut self.buffer, size);+            }+        }++        /// Returns `size` to this thread's buffer, flushing the excess to the global counter once+        /// the buffer grows past [`global::MAX_BUFFER`]. Does nothing with+        /// `custom_allocator`.+        #[inline(always)]+        pub(super) fn buffered_remove(&mut self, size: usize) {+            self.buffer += size;+            if self.buffer > MAX_BUFFER {+                flush_excess(&mut self.buffer);+            }+        }+    }+} +/// Per-thread allocation and deallocation totals. #[derive(Default)] struct ThreadLocalCounter {     /// Thread-local buffer of allocated bytes that have been added to the     /// global counter desprite not being allocated yet. It is unsigned so that     /// means the global counter is always equal or greater than the real     /// value.+    #[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]     buffer: usize,     allocation_counters: AllocationCounters, }  impl ThreadLocalCounter {     const fn new() -> Self {         Self {+            #[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]             buffer: 0,             allocation_counters: AllocationCounters::new(),         }     }+     #[inline(always)]     fn add(&mut self, size: usize) {         self.allocation_counters.allocations += size;         self.allocation_counters.allocation_count += 1;-        if self.buffer >= size {-            self.buffer -= size;-        } else {-            add_slow(self, size);-        }++        #[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]+        self.buffered_add(size);     }      #[inline(always)]     fn remove(&mut self, size: usize) {         self.allocation_counters.deallocations += size;         self.allocation_counters.deallocation_count += 1;-        self.buffer += size;-        if self.buffer > MAX_BUFFER {-            remove_slow(self);-        }++        #[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]+        self.buffered_remove(size);     }      #[inline(always)]@@ -63,62 +143,28 @@ impl ThreadLocalCounter {         self.allocation_counters.deallocation_count += 1;         self.allocation_counters.allocations += new_size;         self.allocation_counters.allocation_count += 1;-        match old_size.cmp(&new_size) {-            std::cmp::Ordering::Equal => {}-            std::cmp::Ordering::Less => {-                let size = new_size - old_size;-                if self.buffer >= size {-                    self.buffer -= size;-                } else {-                    add_slow(self, size);-                }-            }-            std::cmp::Ordering::Greater => {-                let size = old_size - new_size;-                self.buffer += size;-                if self.buffer > MAX_BUFFER {-                    remove_slow(self);-                }++        #[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]+        {+            match old_size.cmp(&new_size) {+                std::cmp::Ordering::Equal => {}+                std::cmp::Ordering::Less => self.buffered_add(new_size - old_size),+                std::cmp::Ordering::Greater => self.buffered_remove(old_size - new_size),             }         }     }      fn unload(&mut self) {-        if self.buffer > 0 {-            ALLOCATED.fetch_sub(self.buffer, Ordering::Relaxed);-            self.buffer = 0;-        }+        #[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]+        global::flush_all(&mut self.buffer);         self.allocation_counters = AllocationCounters::default();     } } -// Keep the uncommon atomic updates out of the allocator's inlined hot path.-#[cold]-#[inline(never)]-fn add_slow(local: &mut ThreadLocalCounter, size: usize) {-    debug_assert!(local.buffer < size);-    let offset = size - local.buffer + TARGET_BUFFER;-    local.buffer = TARGET_BUFFER;-    ALLOCATED.fetch_add(offset, Ordering::Relaxed);-}--#[cold]-#[inline(never)]-fn remove_slow(local: &mut ThreadLocalCounter) {-    debug_assert!(local.buffer > MAX_BUFFER);-    let offset = local.buffer - TARGET_BUFFER;-    local.buffer = TARGET_BUFFER;-    ALLOCATED.fetch_sub(offset, Ordering::Relaxed);-}- thread_local! {   static LOCAL_COUNTER: UnsafeCell<ThreadLocalCounter> = const {UnsafeCell::new(ThreadLocalCounter::new())}; } -pub fn get() -> usize {-    ALLOCATED.load(Ordering::Relaxed)-}- pub fn allocation_counters() -> AllocationCounters {     with_local_counter(|local| local.allocation_counters.clone()) }@@ -155,8 +201,8 @@ pub fn update(old_size: usize, new_size: usize) {     with_local_counter(|local| local.update(old_size, new_size)); } -/// Flushes the thread-local buffer to the global counter. This should be called-/// e. g. when a thread is stopped or goes to sleep for a long time.+/// Clears this thread's counters. Called when a thread stops, so a recycled thread does not+/// inherit the previous occupant's totals. pub fn flush() {     with_local_counter(|local| local.unload()); }@@ -166,42 +212,115 @@ mod tests {     use super::*;      #[test]-    fn counting() {-        let mut expected = get();-        add(100);-        // Initial change should fill up the buffer-        expected += TARGET_BUFFER + 100;-        assert_eq!(get(), expected);+    fn counts_allocations_and_deallocations() {+        let start = allocation_counters();+         add(100);-        // Further changes should use the buffer-        assert_eq!(get(), expected);-        add(MAX_BUFFER);-        // Large changes should require more buffer space-        expected += 100 + MAX_BUFFER;-        assert_eq!(get(), expected);+        add(250);         remove(100);-        // Small changes should use the buffer-        // buffer size is now TARGET_BUFFER + 100-        assert_eq!(get(), expected);-        remove(MAX_BUFFER);-        // The buffer should not grow over MAX_BUFFER-        // buffer size would be TARGET_BUFFER + 100 + MAX_BUFFER-        // but it will be reduce to TARGET_BUFFER-        // this means the global counter should reduce by 100 + MAX_BUFFER-        expected -= MAX_BUFFER + 100;-        assert_eq!(get(), expected);--        update(100, 200);-        // Small reallocations should use the buffer.-        assert_eq!(get(), expected);-        update(0, MAX_BUFFER);-        // Growing beyond the buffer should require more buffer space. The prior small growth-        // consumed another 100 bytes from the buffer.-        expected += MAX_BUFFER + 100;-        assert_eq!(get(), expected);-        update(MAX_BUFFER + 1, 0);-        // Shrinking beyond MAX_BUFFER should flush the excess.-        expected -= MAX_BUFFER + 1;-        assert_eq!(get(), expected);++        let after = allocation_counters();+        assert_eq!(after.allocations - start.allocations, 350);+        assert_eq!(after.allocation_count - start.allocation_count, 2);+        assert_eq!(after.deallocations - start.deallocations, 100);+        assert_eq!(after.deallocation_count - start.deallocation_count, 1);+    }++    #[test]+    fn update_counts_both_sides() {+        let start = allocation_counters();++        update(40, 100);++        let after = allocation_counters();+        assert_eq!(after.allocations - start.allocations, 100);+        assert_eq!(after.allocation_count - start.allocation_count, 1);+        assert_eq!(after.deallocations - start.deallocations, 40);+        assert_eq!(after.deallocation_count - start.deallocation_count, 1);+    }++    /// `reset_allocation_counters` restores a previously captured value, which is how the tracing+    /// layer excludes its own writes from a span's totals.+    #[test]+    fn reset_restores_a_captured_value() {+        let start = allocation_counters();+        add(4096);+        assert!(allocation_counters().allocations > start.allocations);++        reset_allocation_counters(start.clone());+        assert_eq!(allocation_counters().allocations, start.allocations);+        assert_eq!(+            allocation_counters().allocation_count,+            start.allocation_count+        );+    }++    /// `flush` is called when a thread stops so a thread reusing the slot starts clean.+    #[test]+    fn flush_clears_this_threads_counters() {+        std::thread::spawn(|| {+            add(1234);+            assert!(allocation_counters().allocations >= 1234);+            flush();+            let cleared = allocation_counters();+            assert_eq!(cleared.allocations, 0);+            assert_eq!(cleared.allocation_count, 0);+            assert_eq!(cleared.deallocations, 0);+            assert_eq!(cleared.deallocation_count, 0);+        })+        .join()+        .unwrap();+    }++    /// The buffered global counter only exists without the `custom_allocator` feature.+    ///+    /// Asserts the buffering arithmetic on a [`ThreadLocalCounter`] directly: how much a thread+    /// keeps buffered, and therefore when it has to touch the global. The global itself is not+    /// read — it is process-wide, and this binary installs [`crate::TurboMalloc`] as its global+    /// allocator, so every other thread moves it concurrently. `buffer` is thread-local and+    /// exact.+    #[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]+    #[test]+    fn counting() {+        use super::global::{MAX_BUFFER, TARGET_BUFFER};++        let mut local = ThreadLocalCounter::new();++        // A fresh counter has nothing buffered, so the first allocation has to reach the global,+        // taking a full TARGET_BUFFER while it is there.+        local.add(100);+        assert_eq!(local.buffer, TARGET_BUFFER);++        // Further small allocations come straight out of the buffer.+        local.add(100);+        assert_eq!(local.buffer, TARGET_BUFFER - 100);++        // An allocation larger than the buffer refills it from the global.+        local.add(MAX_BUFFER);+        assert_eq!(local.buffer, TARGET_BUFFER);++        // Frees go back into the buffer while it stays under MAX_BUFFER.+        local.remove(100);+        assert_eq!(local.buffer, TARGET_BUFFER + 100);++        // Past MAX_BUFFER the excess is flushed back to the global, down to TARGET_BUFFER.+        local.remove(MAX_BUFFER);+        assert_eq!(local.buffer, TARGET_BUFFER);++        // A reallocation that grows by less than the buffer is served locally.+        local.update(100, 200);+        assert_eq!(local.buffer, TARGET_BUFFER - 100);++        // One that grows beyond it refills.+        local.update(0, MAX_BUFFER);+        assert_eq!(local.buffer, TARGET_BUFFER);++        // One that shrinks beyond MAX_BUFFER flushes the excess.+        local.update(MAX_BUFFER + 1, 0);+        assert_eq!(local.buffer, TARGET_BUFFER);++        // Unloading returns whatever is still buffered.+        local.unload();+        assert_eq!(local.buffer, 0);     } }

@@ -7,7 +7,7 @@ use std::{     ops::{Add, AddAssign}, }; -use self::counter::{add, flush, get, remove, update};+use self::counter::{add, flush, remove, update};  #[derive(Default, Clone, Debug)] pub struct AllocationInfo {@@ -85,16 +85,52 @@ impl AllocationCounters { pub struct TurboMalloc;  impl TurboMalloc {-    /// Returns the current amount of live memory (bytes allocated minus freed)-    /// tracked across all threads.+    /// Returns the bytes the allocator currently has committed from the OS.     ///-    /// For efficiency reasons every thread only synchronizes with this counter after ~100K bytes of-    /// allocations or deallocations.  So this could be off by as much as 100K*number of thread in-    /// either direction.+    /// This is the allocator's own accounting, not a per-OS query, so it means the same thing on+    /// every platform. It counts what mimalloc has taken from the OS, which includes allocator+    /// overhead and fragmentation, and excludes anything mimalloc did not hand out — the binary,+    /// mmap'd files, and any memory allocated by the embedding process. It is a measure of what+    /// this allocator holds, not of the process's total footprint.+    ///+    /// It does not track frees in lock step. mimalloc reuses and purges pages on its own+    /// schedule, so the figure lags a burst of frees, and memory abandoned by threads that have+    /// since exited is only reclaimed by a forcing [`Self::collect`].+    ///+    /// Without the `custom_allocator` feature this is a process-wide counter of live bytes+    /// (allocations minus deallocations), maintained by [`self::counter`]. That figure is+    /// approximate: threads buffer their updates, so it can be off by up to a fixed amount per+    /// thread in either direction.     pub fn memory_usage() -> usize {-        get()+        #[cfg(all(feature = "custom_allocator", not(target_family = "wasm")))]+        {+            // `current_commit` is a relaxed atomic load, but `mi_process_info` also calls+            // `_mi_prim_process_info`, which is a `getrusage` (plus a `task_info` on macOS). All+            // eight out-params are optional, so ask only for the one we use.+            let mut current_commit = 0usize;+            // Safety: every out-param is either null or a valid `usize` we own.+            unsafe {+                libmimalloc_sys::mi_process_info(+                    std::ptr::null_mut(),+                    std::ptr::null_mut(),+                    std::ptr::null_mut(),+                    std::ptr::null_mut(),+                    std::ptr::null_mut(),+                    &mut current_commit,+                    std::ptr::null_mut(),+                    std::ptr::null_mut(),+                );+            }+            current_commit+        }+        #[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]+        {+            self::counter::get()+        }     } +    /// Clears the calling thread's allocation counters. Call this when a thread is about to stop,+    /// so a thread that reuses its slot does not inherit the previous totals.     pub fn thread_stop() {         flush();     }@@ -205,6 +241,40 @@ unsafe impl GlobalAlloc for TurboMalloc { mod tests {     use super::TurboMalloc; +    // `memory_usage` reports what *this* allocator has committed, so the test binary has to+    // actually route its allocations through it. Without this the `vec!` below goes to the+    // system allocator and mimalloc's counter never moves.+    #[global_allocator]+    static ALLOC: TurboMalloc = TurboMalloc;++    /// Also guards against the counter silently becoming unavailable. mimalloc's `committed`+    /// stat is maintained even at `MI_STAT 0` (which is what a release build compiles, since+    /// `build.rs` sets `MI_DEBUG=0`) because the `mi_os_stat_*` macros are not gated on+    /// `MI_STAT` — an internal detail rather than a documented guarantee, so a+    /// `libmimalloc-sys` bump could zero it out. If that happens, this fails.+    #[test]+    fn memory_usage_is_reported_and_tracks_a_large_allocation() {+        let before = TurboMalloc::memory_usage();+        assert!(before > 0, "a running process has live memory");++        // Large enough to dwarf whatever else the test process does concurrently, and written to+        // so the pages are actually committed.+        const SIZE: usize = 256 * 1024 * 1024;+        let mut buffer = vec![0u8; SIZE];+        for chunk in buffer.chunks_mut(4096) {+            chunk[0] = 1;+        }+        std::hint::black_box(&buffer);++        let after = TurboMalloc::memory_usage();+        assert!(+            after >= before + SIZE / 2,+            "expected a rise of at least {} bytes, got {before} -> {after}",+            SIZE / 2+        );+        drop(buffer);+    }+     #[test]     fn memory_pressure_is_in_range() {         let value = TurboMalloc::memory_pressure();