allocator: Add optional heap allocation measurement tracking. (#38727)

Add an off-by-default allocator mode that tracks all live allocations
with sizes and associated stack traces. We also track if each allocation
is visited as part of a measuring heap usage in `about:memory`, allowing
us to report on allocations that are not tracked yet. Right now the list
of untracked allocations is dumped to stdout; I have a python script
coming in a separate PR which makes it easier to perform analysis on the
massive output.

Testing: Manually tested with `./mach build -d --features
servo_allocator/allocation-tracking` and visiting about:memory.
Part of: #11559

---------

Signed-off-by: Josh Matthews <josh@joshmatthews.net>

Cargo.lock

@@ -6551,6 +6551,7 @@ dependencies = [
  "regex",
  "serde",
  "serde_json",
+ "servo_allocator",
  "servo_config",
  "task_info",
  "tikv-jemalloc-sys",
@@ -7655,6 +7656,8 @@ dependencies = [
 name = "servo_allocator"
 version = "0.0.1"
 dependencies = [
+ "backtrace",
+ "fnv",
  "libc",
  "tikv-jemalloc-sys",
  "tikv-jemallocator",

components/allocator/Cargo.toml

@@ -11,8 +11,13 @@ rust-version.workspace = true
 path = "lib.rs"
 
 [features]
+allocation-tracking = ["backtrace", "fnv"]
 use-system-allocator = ["libc"]
 
+[dependencies]
+backtrace = { workspace = true, optional = true }
+fnv = { workspace = true, optional = true }
+
 [target.'cfg(not(any(windows, target_env = "ohos")))'.dependencies]
 libc = { workspace = true, optional = true }
 tikv-jemalloc-sys = { workspace = true }

components/allocator/lib.rs

@@ -2,13 +2,53 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
-//! Selecting the default global allocator for Servo
+//! Selecting the default global allocator for Servo, and exposing common
+//! allocator introspection APIs for memory profiling.
 
+use std::os::raw::c_void;
+
+#[cfg(not(feature = "allocation-tracking"))]
 #[global_allocator]
 static ALLOC: Allocator = Allocator;
 
+#[cfg(feature = "allocation-tracking")]
+#[global_allocator]
+static ALLOC: crate::tracking::AccountingAlloc<Allocator> =
+    crate::tracking::AccountingAlloc::with_allocator(Allocator);
+
+#[cfg(feature = "allocation-tracking")]
+mod tracking;
+
+pub fn dump_unmeasured() {
+    #[cfg(feature = "allocation-tracking")]
+    ALLOC.dump_unmeasured_allocations();
+}
+
 pub use crate::platform::*;
 
+type EnclosingSizeFn = unsafe extern "C" fn(*const c_void) -> usize;
+
+/// # Safety
+/// No restrictions. The passed pointer is never dereferenced.
+/// This function is only marked unsafe because the MallocSizeOfOps APIs
+/// requires an unsafe function pointer.
+#[cfg(feature = "allocation-tracking")]
+unsafe extern "C" fn enclosing_size_impl(ptr: *const c_void) -> usize {
+    let (adjusted, size) = crate::ALLOC.enclosing_size(ptr);
+    if size != 0 {
+        crate::ALLOC.note_allocation(adjusted, size);
+    }
+    size
+}
+
+#[allow(non_upper_case_globals)]
+#[cfg(feature = "allocation-tracking")]
+pub static enclosing_size: Option<EnclosingSizeFn> = Some(crate::enclosing_size_impl);
+
+#[allow(non_upper_case_globals)]
+#[cfg(not(feature = "allocation-tracking"))]
+pub static enclosing_size: Option<EnclosingSizeFn> = None;
+
 #[cfg(not(any(windows, feature = "use-system-allocator", target_env = "ohos")))]
 mod platform {
     use std::os::raw::c_void;
@@ -21,7 +61,10 @@ mod platform {
     ///
     /// Passing a non-heap allocated pointer to this function results in undefined behavior.
     pub unsafe extern "C" fn usable_size(ptr: *const c_void) -> usize {
-        unsafe { tikv_jemallocator::usable_size(ptr) }
+        let size = unsafe { tikv_jemallocator::usable_size(ptr) };
+        #[cfg(feature = "allocation-tracking")]
+        crate::ALLOC.note_allocation(ptr, size);
+        size
     }
 
     /// Memory allocation APIs compatible with libc
@@ -46,12 +89,18 @@ mod platform {
     pub unsafe extern "C" fn usable_size(ptr: *const c_void) -> usize {
         #[cfg(target_vendor = "apple")]
         unsafe {
-            return libc::malloc_size(ptr);
+            let size = libc::malloc_size(ptr);
+            #[cfg(feature = "allocation-tracking")]
+            crate::ALLOC.note_allocation(ptr, size);
+            size
         }
 
         #[cfg(not(target_vendor = "apple"))]
         unsafe {
-            return libc::malloc_usable_size(ptr as *mut _);
+            let size = libc::malloc_usable_size(ptr as *mut _);
+            #[cfg(feature = "allocation-tracking")]
+            crate::ALLOC.note_allocation(ptr, size);
+            size
         }
     }
 
@@ -81,7 +130,10 @@ mod platform {
                 ptr = *(ptr as *const *const c_void).offset(-1)
             }
 
-            HeapSize(heap, 0, ptr) as usize
+            let size = HeapSize(heap, 0, ptr) as usize;
+            #[cfg(feature = "allocation-tracking")]
+            crate::ALLOC.note_allocation(ptr, size);
+            size
         }
     }
 }

components/allocator/tracking.rs

@@ -0,0 +1,211 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
+
+//! An allocator wrapper that records metadata about each live allocation.
+//! This metadata can then be used to identify allocations that are not visible
+//! through any existing MallocSizeOf path.
+
+use std::alloc::{GlobalAlloc, Layout, System};
+use std::cell::Cell;
+use std::collections::hash_map::Entry;
+use std::os::raw::c_void;
+use std::sync::Mutex;
+
+use fnv::{FnvBuildHasher, FnvHashMap};
+
+/// The maximum number of allocations that we'll keep track of. Once the limit
+/// is reached, we'll evict the first allocation that is smaller than the new addition.
+const MAX_TRACKED_ALLOCATIONS: usize = usize::MAX;
+/// Cap the number of stack frames that we'll store per allocation.
+const MAX_FRAMES: usize = 50;
+/// A certain number of frames at the top of the allocation stack are just
+/// just internal liballoc implementation details or AccountingAlloc functions.
+/// We can skip them without losing any meaningful information.
+const SKIPPED_FRAMES: usize = 5;
+/// The minimum size of allocation that we'll track. Can be used to reduce overhead
+/// by skipping bookkeeping for small allocations.
+const MIN_SIZE: usize = 0;
+
+thread_local! {
+    static IN_ALLOCATION: Cell<bool> = Cell::new(false);
+}
+
+#[derive(PartialEq, Eq)]
+struct AllocSite {
+    /// The stack at the time this allocation was recorded.
+    frames: [*mut std::ffi::c_void; MAX_FRAMES],
+    /// The start of the allocated memory.
+    ptr: *mut u8,
+    /// The size of the allocated memory.
+    size: usize,
+    /// True if this allocation site's size is ever queried after the initial
+    /// allocation. If false, it means that the allocation is not visible from
+    /// any of the MallocSizeOf roots.
+    noted: bool,
+}
+
+impl AllocSite {
+    fn contains(&self, ptr: *mut u8) -> bool {
+        ptr >= self.ptr && ptr < unsafe { self.ptr.offset(self.size as isize) }
+    }
+}
+
+unsafe impl Send for AllocSite {}
+
+/// A map of pointers to allocation callsite metadata.
+static ALLOCATION_SITES: Mutex<FnvHashMap<usize, AllocSite>> =
+    const { Mutex::new(FnvHashMap::with_hasher(FnvBuildHasher::new())) };
+
+#[derive(Default)]
+pub struct AccountingAlloc<A = System> {
+    allocator: A,
+}
+
+impl<A> AccountingAlloc<A> {
+    pub const fn with_allocator(allocator: A) -> Self {
+        Self { allocator }
+    }
+
+    fn remove_allocation(&self, ptr: *const c_void, size: usize) {
+        if size < MIN_SIZE {
+            return;
+        }
+        let old = IN_ALLOCATION.with(|status| status.replace(true));
+        if old {
+            return;
+        }
+        let mut sites = ALLOCATION_SITES.lock().unwrap();
+        if let Entry::Occupied(e) = sites.entry(ptr as usize) {
+            e.remove();
+        }
+        IN_ALLOCATION.with(|status| status.set(old));
+    }
+
+    fn record_allocation(&self, ptr: *mut u8, size: usize) {
+        if size < MIN_SIZE {
+            return;
+        }
+        let old = IN_ALLOCATION.with(|status| status.replace(true));
+        if old {
+            return;
+        }
+        let mut num_skipped = 0;
+        let mut num_frames = 0;
+        let mut frames = [std::ptr::null_mut(); MAX_FRAMES];
+        backtrace::trace(|frame| {
+            if num_skipped < SKIPPED_FRAMES {
+                num_skipped += 1;
+                return true;
+            }
+            frames[num_frames] = frame.ip();
+            num_frames += 1;
+            num_frames < MAX_FRAMES
+        });
+        let site = AllocSite {
+            frames,
+            size,
+            ptr,
+            noted: false,
+        };
+        let mut sites = ALLOCATION_SITES.lock().unwrap();
+
+        if sites.len() < MAX_TRACKED_ALLOCATIONS {
+            sites.insert(ptr as usize, site);
+        } else if let Some(key) = sites
+            .iter()
+            .find(|(_, s)| s.size < site.size)
+            .map(|(k, _)| k.clone())
+        {
+            sites.remove(&key);
+            sites.insert(ptr as usize, site);
+        }
+
+        IN_ALLOCATION.with(|status| status.set(old));
+    }
+
+    pub(crate) fn enclosing_size(&self, ptr: *const c_void) -> (*const c_void, usize) {
+        ALLOCATION_SITES
+            .lock()
+            .unwrap()
+            .iter()
+            .find(|(_, site)| site.contains(ptr.cast_mut().cast()))
+            .map_or((std::ptr::null_mut(), 0), |(_, site)| {
+                (site.ptr.cast(), site.size)
+            })
+    }
+
+    pub(crate) fn note_allocation(&self, ptr: *const c_void, size: usize) {
+        if size < MIN_SIZE {
+            return;
+        }
+        IN_ALLOCATION.with(|status| status.set(true));
+        if let Some(site) = ALLOCATION_SITES.lock().unwrap().get_mut(&(ptr as usize)) {
+            site.noted = true;
+        }
+        IN_ALLOCATION.with(|status| status.set(false));
+    }
+
+    pub(crate) fn dump_unmeasured_allocations(&self) {
+        // Ensure that we ignore all allocations triggered while processing
+        // the existing allocation data.
+        IN_ALLOCATION.with(|status| status.set(true));
+        {
+            let sites = ALLOCATION_SITES.lock().unwrap();
+            let default = "???";
+            for site in sites
+                .values()
+                .filter(|site| site.size > MIN_SIZE && !site.noted)
+            {
+                let mut resolved = vec![];
+                for ip in site.frames.iter().filter(|ip| !ip.is_null()) {
+                    backtrace::resolve(*ip, |symbol| {
+                        resolved.push((
+                            symbol.filename().map(|f| f.to_owned()),
+                            symbol.lineno(),
+                            symbol.name().map(|n| format!("{}", n)),
+                        ));
+                    });
+                }
+
+                println!("---\n{}\n", site.size);
+                for (filename, line, symbol) in &resolved {
+                    let fname = filename.as_ref().map(|f| f.display().to_string());
+                    println!(
+                        "{}:{} ({})",
+                        fname.as_deref().unwrap_or(default),
+                        line.unwrap_or_default(),
+                        symbol.as_deref().unwrap_or(default),
+                    );
+                }
+            }
+        }
+        IN_ALLOCATION.with(|status| status.set(false));
+    }
+}
+
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAlloc<A> {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let ptr = unsafe { self.allocator.alloc(layout) };
+        self.record_allocation(ptr, layout.size());
+        ptr
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        unsafe { self.allocator.dealloc(ptr, layout) };
+        self.remove_allocation(ptr.cast(), layout.size());
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        let ptr = unsafe { self.allocator.alloc_zeroed(layout) };
+        self.record_allocation(ptr, layout.size());
+        ptr
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        self.remove_allocation(ptr.cast(), layout.size());
+        let ptr = unsafe { self.allocator.realloc(ptr, layout, new_size) };
+        self.record_allocation(ptr, new_size);
+        ptr
+    }
+}

components/profile/Cargo.toml

@@ -18,6 +18,7 @@ log = { workspace = true }
 profile_traits = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
+servo_allocator = { path = "../allocator" }
 servo_config = { path = "../config" }
 time = { workspace = true }
 

components/profile/mem.rs

@@ -112,6 +112,7 @@ impl Profiler {
                     })
                     .collect();
                 let _ = sender.send(MemoryReportResult { results });
+                servo_allocator::dump_unmeasured();
                 true
             },
             ProfilerMsg::Exit => false,

components/shared/profile/mem.rs

@@ -293,7 +293,7 @@ pub fn perform_memory_report<F: FnOnce(&mut MallocSizeOfOps)>(f: F) {
     let seen_pointer = move |ptr| SEEN_POINTERS.with(|pointers| !pointers.borrow_mut().insert(ptr));
     let mut ops = MallocSizeOfOps::new(
         servo_allocator::usable_size,
-        None,
+        servo_allocator::enclosing_size,
         Some(Box::new(seen_pointer)),
     );
     f(&mut ops);

ports/servoshell/Cargo.toml

@@ -121,7 +121,7 @@ headers = { workspace = true }
 net = { path = "../../components/net" }
 net_traits = { workspace = true }
 serde_json = { workspace = true }
-# For optional feature servo_allocator/use-system-allocator
+# For optional feature servo_allocator/use-system-allocator and servo_allocator/allocation-tracking
 servo_allocator = { path = "../../components/allocator" }
 surfman = { workspace = true, features = ["sm-raw-window-handle-06", "sm-x11"] }
 winit = { workspace = true }