Refactor common infrastructure for creating memory reports. (#36579)

This removes a bunch of duplicated code needed to support
ConditionalMallocSizeOf correctly, and fixes multiple places where that
code was subtly wrong (the seen pointers hashset was never cleared).

Testing: Measuring https://www.nist.gov/image-gallery lots of times.

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

11 files changed, 68 insertions, 57 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 429abb44e50..2d3a667e976 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5835,6 +5835,7 @@ dependencies = [
  "log",
  "malloc_size_of_derive",
  "serde",
+ "servo_allocator",
  "servo_config",
  "servo_malloc_size_of",
  "signpost",
diff --git a/components/layout_thread_2020/lib.rs b/components/layout_thread_2020/lib.rs
index f564f93790b..6a9922701e4 100644
--- a/components/layout_thread_2020/lib.rs
+++ b/components/layout_thread_2020/lib.rs
@@ -399,11 +399,7 @@ impl Layout for LayoutThread {
 
     fn exit_now(&mut self) {}
 
-    fn collect_reports(&self, reports: &mut Vec<Report>) {
-        // Servo uses vanilla jemalloc, which doesn't have a
-        // malloc_enclosing_size_of function.
-        let mut ops = MallocSizeOfOps::new(servo_allocator::usable_size, None, None);
-
+    fn collect_reports(&self, reports: &mut Vec<Report>, ops: &mut MallocSizeOfOps) {
         // TODO: Measure more than just display list, stylist, and font context.
         let formatted_url = &format!("url({})", self.url);
         reports.push(Report {
@@ -415,13 +411,13 @@ impl Layout for LayoutThread {
         reports.push(Report {
             path: path![formatted_url, "layout-thread", "stylist"],
             kind: ReportKind::ExplicitJemallocHeapSize,
-            size: self.stylist.size_of(&mut ops),
+            size: self.stylist.size_of(ops),
         });
 
         reports.push(Report {
             path: path![formatted_url, "layout-thread", "font-context"],
             kind: ReportKind::ExplicitJemallocHeapSize,
-            size: self.font_context.size_of(&mut ops),
+            size: self.font_context.size_of(ops),
         });
     }
 
diff --git a/components/net/image_cache.rs b/components/net/image_cache.rs
index 078380bb466..55a2c86cae8 100644
--- a/components/net/image_cache.rs
+++ b/components/net/image_cache.rs
@@ -2,10 +2,8 @@
  * 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/. */
 
-use std::cell::{LazyCell, RefCell};
+use std::collections::HashMap;
 use std::collections::hash_map::Entry::{Occupied, Vacant};
-use std::collections::{HashMap, HashSet};
-use std::ffi::c_void;
 use std::sync::{Arc, Mutex};
 use std::{mem, thread};
 
@@ -460,15 +458,8 @@ impl ImageCache for ImageCacheImpl {
         }
     }
 
-    fn memory_report(&self, prefix: &str) -> Report {
-        let seen_pointer =
-            move |ptr| SEEN_POINTERS.with(|pointers| !pointers.borrow_mut().insert(ptr));
-        let mut ops = MallocSizeOfOps::new(
-            servo_allocator::usable_size,
-            None,
-            Some(Box::new(seen_pointer)),
-        );
-        let size = self.store.lock().unwrap().size_of(&mut ops);
+    fn memory_report(&self, prefix: &str, ops: &mut MallocSizeOfOps) -> Report {
+        let size = self.store.lock().unwrap().size_of(ops);
         Report {
             path: path![prefix, "image-cache"],
             kind: ReportKind::ExplicitSystemHeapSize,
@@ -678,7 +669,3 @@ impl ImageCacheImpl {
         warn!("Couldn't find cached entry for listener {:?}", id);
     }
 }
-
-thread_local!(static SEEN_POINTERS: LazyCell<RefCell<HashSet<*const c_void>>> = const {
-    LazyCell::new(|| RefCell::new(HashSet::new()))
-});
diff --git a/components/net/resource_thread.rs b/components/net/resource_thread.rs
index 7235c362d66..b6f885f29b7 100644
--- a/components/net/resource_thread.rs
+++ b/components/net/resource_thread.rs
@@ -21,7 +21,6 @@ use embedder_traits::EmbedderProxy;
 use hyper_serde::Serde;
 use ipc_channel::ipc::{self, IpcReceiver, IpcReceiverSet, IpcSender};
 use log::{debug, trace, warn};
-use malloc_size_of::MallocSizeOfOps;
 use net_traits::blob_url_store::parse_blob_url;
 use net_traits::filemanager_thread::FileTokenCheck;
 use net_traits::request::{Destination, RequestBuilder, RequestId};
@@ -32,7 +31,9 @@ use net_traits::{
     FetchChannels, FetchTaskTarget, ResourceFetchTiming, ResourceThreads, ResourceTimingType,
     WebSocketDomAction, WebSocketNetworkEvent,
 };
-use profile_traits::mem::{ProcessReports, ProfilerChan as MemProfilerChan, ReportsChan};
+use profile_traits::mem::{
+    ProcessReports, ProfilerChan as MemProfilerChan, ReportsChan, perform_memory_report,
+};
 use profile_traits::time::ProfilerChan;
 use rustls::RootCertStore;
 use serde::{Deserialize, Serialize};
@@ -280,10 +281,11 @@ impl ResourceChannelManager {
         public_http_state: &Arc<HttpState>,
         private_http_state: &Arc<HttpState>,
     ) {
-        let mut ops = MallocSizeOfOps::new(servo_allocator::usable_size, None, None);
-        let mut reports = public_http_state.memory_reports("public", &mut ops);
-        reports.extend(private_http_state.memory_reports("private", &mut ops));
-        msg.send(ProcessReports::new(reports));
+        perform_memory_report(|ops| {
+            let mut reports = public_http_state.memory_reports("public", ops);
+            reports.extend(private_http_state.memory_reports("private", ops));
+            msg.send(ProcessReports::new(reports));
+        })
     }
 
     fn cancellation_listener(&self, request_id: RequestId) -> Option<Arc<CancellationListener>> {
diff --git a/components/script/dom/workerglobalscope.rs b/components/script/dom/workerglobalscope.rs
index 275a9cb0741..5775375c385 100644
--- a/components/script/dom/workerglobalscope.rs
+++ b/components/script/dom/workerglobalscope.rs
@@ -25,7 +25,7 @@ use net_traits::request::{
     CredentialsMode, Destination, InsecureRequestsPolicy, ParserMetadata,
     RequestBuilder as NetRequestInit,
 };
-use profile_traits::mem::ProcessReports;
+use profile_traits::mem::{ProcessReports, perform_memory_report};
 use servo_url::{MutableOrigin, ServoUrl};
 use timers::TimerScheduler;
 use uuid::Uuid;
@@ -546,8 +546,10 @@ impl WorkerGlobalScope {
             CommonScriptMsg::Task(_, task, _, _) => task.run_box(),
             CommonScriptMsg::CollectReports(reports_chan) => {
                 let cx = self.get_cx();
-                let reports = cx.get_reports(format!("url({})", self.get_url()));
-                reports_chan.send(ProcessReports::new(reports));
+                perform_memory_report(|ops| {
+                    let reports = cx.get_reports(format!("url({})", self.get_url()), ops);
+                    reports_chan.send(ProcessReports::new(reports));
+                });
             },
         }
         true
diff --git a/components/script/script_runtime.rs b/components/script/script_runtime.rs
index b23be19720c..d6832a644ec 100644
--- a/components/script/script_runtime.rs
+++ b/components/script/script_runtime.rs
@@ -8,8 +8,7 @@
 #![allow(dead_code)]
 
 use core::ffi::c_char;
-use std::cell::{Cell, LazyCell, RefCell};
-use std::collections::HashSet;
+use std::cell::Cell;
 use std::ffi::CString;
 use std::io::{Write, stdout};
 use std::ops::Deref;
@@ -818,9 +817,7 @@ fn in_range<T: PartialOrd + Copy>(val: T, min: T, max: T) -> Option<T> {
     }
 }
 
-thread_local!(static SEEN_POINTERS: LazyCell<RefCell<HashSet<*const c_void>>> = const {
-    LazyCell::new(|| RefCell::new(HashSet::new()))
-});
+thread_local!(static MALLOC_SIZE_OF_OPS: Cell<*mut MallocSizeOfOps> = const { Cell::new(ptr::null_mut()) });
 
 #[allow(unsafe_code)]
 unsafe extern "C" fn get_size(obj: *mut JSObject) -> usize {
@@ -831,14 +828,8 @@ unsafe extern "C" fn get_size(obj: *mut JSObject) -> usize {
             if dom_object.is_null() {
                 return 0;
             }
-            let seen_pointer =
-                move |ptr| SEEN_POINTERS.with(|pointers| !pointers.borrow_mut().insert(ptr));
-            let mut ops = MallocSizeOfOps::new(
-                servo_allocator::usable_size,
-                None,
-                Some(Box::new(seen_pointer)),
-            );
-            (v.malloc_size_of)(&mut ops, dom_object)
+            let ops = MALLOC_SIZE_OF_OPS.get();
+            (v.malloc_size_of)(&mut *ops, dom_object)
         },
         Err(_e) => 0,
     }
@@ -943,13 +934,13 @@ pub(crate) use script_bindings::script_runtime::JSContext;
 /// Extra methods for the JSContext type defined in script_bindings, when
 /// the methods are only called by code in the script crate.
 pub(crate) trait JSContextHelper {
-    fn get_reports(&self, path_seg: String) -> Vec<Report>;
+    fn get_reports(&self, path_seg: String, ops: &mut MallocSizeOfOps) -> Vec<Report>;
 }
 
 impl JSContextHelper for JSContext {
     #[allow(unsafe_code)]
-    fn get_reports(&self, path_seg: String) -> Vec<Report> {
-        SEEN_POINTERS.with(|pointers| pointers.borrow_mut().clear());
+    fn get_reports(&self, path_seg: String, ops: &mut MallocSizeOfOps) -> Vec<Report> {
+        MALLOC_SIZE_OF_OPS.with(|ops_tls| ops_tls.set(ops));
         let stats = unsafe {
             let mut stats = ::std::mem::zeroed();
             if !CollectServoSizes(**self, &mut stats, Some(get_size)) {
@@ -957,6 +948,7 @@ impl JSContextHelper for JSContext {
             }
             stats
         };
+        MALLOC_SIZE_OF_OPS.with(|ops| ops.set(ptr::null_mut()));
 
         let mut reports = vec![];
         let mut report = |mut path_suffix, kind, size| {
diff --git a/components/script/script_thread.rs b/components/script/script_thread.rs
index a31b43ca609..fad228efed0 100644
--- a/components/script/script_thread.rs
+++ b/components/script/script_thread.rs
@@ -78,7 +78,7 @@ use net_traits::{
     ResourceFetchTiming, ResourceThreads, ResourceTimingType,
 };
 use percent_encoding::percent_decode;
-use profile_traits::mem::{ProcessReports, ReportsChan};
+use profile_traits::mem::{ProcessReports, ReportsChan, perform_memory_report};
 use profile_traits::time::ProfilerCategory;
 use profile_traits::time_profile;
 use script_layout_interface::{
@@ -2435,13 +2435,19 @@ impl ScriptThread {
         let documents = self.documents.borrow();
         let urls = itertools::join(documents.iter().map(|(_, d)| d.url().to_string()), ", ");
 
-        let prefix = format!("url({urls})");
-        let mut reports = self.get_cx().get_reports(prefix.clone());
-        for (_, document) in documents.iter() {
-            document.window().layout().collect_reports(&mut reports);
-        }
+        let mut reports = vec![];
+        perform_memory_report(|ops| {
+            let prefix = format!("url({urls})");
+            reports.extend(self.get_cx().get_reports(prefix.clone(), ops));
+            for (_, document) in documents.iter() {
+                document
+                    .window()
+                    .layout()
+                    .collect_reports(&mut reports, ops);
+            }
 
-        reports.push(self.image_cache.memory_report(&prefix));
+            reports.push(self.image_cache.memory_report(&prefix, ops));
+        });
 
         reports_chan.send(ProcessReports::new(reports));
     }
diff --git a/components/shared/net/image_cache.rs b/components/shared/net/image_cache.rs
index 46ed85d280f..8fb329d304f 100644
--- a/components/shared/net/image_cache.rs
+++ b/components/shared/net/image_cache.rs
@@ -8,6 +8,7 @@ use base::id::PipelineId;
 use compositing_traits::CrossProcessCompositorApi;
 use ipc_channel::ipc::IpcSender;
 use log::debug;
+use malloc_size_of::MallocSizeOfOps;
 use malloc_size_of_derive::MallocSizeOf;
 use pixels::{Image, ImageMetadata};
 use profile_traits::mem::Report;
@@ -116,7 +117,7 @@ pub trait ImageCache: Sync + Send {
     where
         Self: Sized;
 
-    fn memory_report(&self, prefix: &str) -> Report;
+    fn memory_report(&self, prefix: &str, ops: &mut MallocSizeOfOps) -> Report;
 
     /// Definitively check whether there is a cached, fully loaded image available.
     fn get_image(
diff --git a/components/shared/profile/Cargo.toml b/components/shared/profile/Cargo.toml
index 645ab24b4c0..8086bf17ba2 100644
--- a/components/shared/profile/Cargo.toml
+++ b/components/shared/profile/Cargo.toml
@@ -22,6 +22,7 @@ log = { workspace = true }
 malloc_size_of = { workspace = true }
 malloc_size_of_derive = { workspace = true }
 serde = { workspace = true }
+servo_allocator = { path = "../../allocator" }
 servo_config = { path = "../../config" }
 signpost = { git = "https://github.com/pcwalton/signpost.git" }
 strum_macros = { workspace = true }
diff --git a/components/shared/profile/mem.rs b/components/shared/profile/mem.rs
index 62fc30d9438..1be4eb5abc4 100644
--- a/components/shared/profile/mem.rs
+++ b/components/shared/profile/mem.rs
@@ -6,12 +6,16 @@
 
 #![deny(missing_docs)]
 
+use std::cell::{LazyCell, RefCell};
+use std::collections::HashSet;
+use std::ffi::c_void;
 use std::marker::Send;
 
 use crossbeam_channel::Sender;
 use ipc_channel::ipc::{self, IpcSender};
 use ipc_channel::router::ROUTER;
 use log::warn;
+use malloc_size_of::MallocSizeOfOps;
 use serde::{Deserialize, Serialize};
 
 /// A trait to abstract away the various kinds of message senders we use.
@@ -266,3 +270,21 @@ pub enum ProfilerMsg {
     /// Triggers sending back the memory profiling metrics,
     Report(IpcSender<MemoryReportResult>),
 }
+
+thread_local!(static SEEN_POINTERS: LazyCell<RefCell<HashSet<*const c_void>>> = const {
+    LazyCell::new(Default::default)
+});
+
+/// Invoke the provided function after initializing the memory profile tools.
+/// The function is expected to call all the desired [MallocSizeOf::size_of]
+/// for allocations reachable from the current thread.
+pub fn perform_memory_report<F: FnOnce(&mut MallocSizeOfOps)>(f: F) {
+    SEEN_POINTERS.with(|pointers| pointers.borrow_mut().clear());
+    let seen_pointer = move |ptr| SEEN_POINTERS.with(|pointers| !pointers.borrow_mut().insert(ptr));
+    let mut ops = MallocSizeOfOps::new(
+        servo_allocator::usable_size,
+        None,
+        Some(Box::new(seen_pointer)),
+    );
+    f(&mut ops);
+}
diff --git a/components/shared/script_layout/lib.rs b/components/shared/script_layout/lib.rs
index ac1e332e3b6..9b052642c32 100644
--- a/components/shared/script_layout/lib.rs
+++ b/components/shared/script_layout/lib.rs
@@ -27,6 +27,7 @@ use fonts::{FontContext, SystemFontServiceProxy};
 use fxhash::FxHashMap;
 use ipc_channel::ipc::IpcSender;
 use libc::c_void;
+use malloc_size_of::MallocSizeOfOps;
 use malloc_size_of_derive::MallocSizeOf;
 use net_traits::image_cache::{ImageCache, PendingImageId};
 use pixels::Image;
@@ -217,7 +218,7 @@ pub trait Layout {
 
     /// Requests that layout measure its memory usage. The resulting reports are sent back
     /// via the supplied channel.
-    fn collect_reports(&self, reports: &mut Vec<Report>);
+    fn collect_reports(&self, reports: &mut Vec<Report>, ops: &mut MallocSizeOfOps);
 
     /// Sets quirks mode for the document, causing the quirks mode stylesheet to be used.
     fn set_quirks_mode(&mut self, quirks_mode: QuirksMode);