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>

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

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| {

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));
     }