script: Measure stored layout data memory usage. (#36664)

We previously ignored the opaque layout data field inside each node when
measuring a DOM node's memory usage. While some of the reachable memory
was accounted for by measuring the layout's box tree, measuring it via
the node ensures that we don't miss anything. Since there are often Arc
values involved, this means that the layout-thread box tree measurements
now look quite small, while reported JS heap usage has increased.

Testing: Manually compared about:memory for servo.org.

---------

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

7 files changed, 34 insertions, 14 deletions

diff --git a/components/layout/dom.rs b/components/layout/dom.rs
index 6db4dbccd41..a81a6eec545 100644
--- a/components/layout/dom.rs
+++ b/components/layout/dom.rs
@@ -2,18 +2,21 @@
  * 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::any::Any;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
 use atomic_refcell::{AtomicRef, AtomicRefCell, AtomicRefMut};
 use base::id::{BrowsingContextId, PipelineId};
 use html5ever::{local_name, ns};
+use malloc_size_of_derive::MallocSizeOf;
 use pixels::Image;
 use script_layout_interface::wrapper_traits::{
     LayoutDataTrait, LayoutNode, ThreadSafeLayoutElement, ThreadSafeLayoutNode,
 };
 use script_layout_interface::{
-    HTMLCanvasDataSource, LayoutElementType, LayoutNodeType as ScriptLayoutNodeType,
+    GenericLayoutDataTrait, HTMLCanvasDataSource, LayoutElementType,
+    LayoutNodeType as ScriptLayoutNodeType,
 };
 use servo_arc::Arc as ServoArc;
 use style::properties::ComputedValues;
@@ -31,7 +34,7 @@ use crate::table::TableLevelBox;
 use crate::taffy::TaffyItemBox;
 
 /// The data that is stored in each DOM node that is used by layout.
-#[derive(Default)]
+#[derive(Default, MallocSizeOf)]
 pub struct InnerDOMLayoutData {
     pub(super) self_box: ArcRefCell<Option<LayoutBox>>,
     pub(super) pseudo_before_box: ArcRefCell<Option<LayoutBox>>,
@@ -54,6 +57,7 @@ impl InnerDOMLayoutData {
 }
 
 /// A box that is stored in one of the `DOMLayoutData` slots.
+#[derive(MallocSizeOf)]
 pub(super) enum LayoutBox {
     DisplayContents,
     BlockLevel(ArcRefCell<BlockLevelBox>),
@@ -98,11 +102,16 @@ impl LayoutBox {
 /// A wrapper for [`InnerDOMLayoutData`]. This is necessary to give the entire data
 /// structure interior mutability, as we will need to mutate the layout data of
 /// non-mutable DOM nodes.
-#[derive(Default)]
+#[derive(Default, MallocSizeOf)]
 pub struct DOMLayoutData(AtomicRefCell<InnerDOMLayoutData>);
 
 // The implementation of this trait allows the data to be stored in the DOM.
 impl LayoutDataTrait for DOMLayoutData {}
+impl GenericLayoutDataTrait for DOMLayoutData {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}
 
 pub struct BoxSlot<'dom> {
     pub(crate) slot: Option<ArcRefCell<Option<LayoutBox>>>,
@@ -255,6 +264,7 @@ where
         }
         LayoutNode::layout_data(&self)
             .unwrap()
+            .as_any()
             .downcast_ref::<DOMLayoutData>()
             .unwrap()
             .0
@@ -262,8 +272,13 @@ where
     }
 
     fn layout_data(self) -> Option<AtomicRef<'dom, InnerDOMLayoutData>> {
-        LayoutNode::layout_data(&self)
-            .map(|data| data.downcast_ref::<DOMLayoutData>().unwrap().0.borrow())
+        LayoutNode::layout_data(&self).map(|data| {
+            data.as_any()
+                .downcast_ref::<DOMLayoutData>()
+                .unwrap()
+                .0
+                .borrow()
+        })
     }
 
     fn element_box_slot(&self) -> BoxSlot<'dom> {
diff --git a/components/layout/table/mod.rs b/components/layout/table/mod.rs
index 120270fc7cf..fe7f90437b8 100644
--- a/components/layout/table/mod.rs
+++ b/components/layout/table/mod.rs
@@ -346,6 +346,7 @@ pub(crate) struct TableLayoutStyle<'a> {
 /// Table parts that are stored in the DOM. This is used in order to map from
 /// the DOM to the box tree and will eventually be important for incremental
 /// layout.
+#[derive(MallocSizeOf)]
 pub(crate) enum TableLevelBox {
     Caption(ArcRefCell<TableCaption>),
     Cell(ArcRefCell<TableSlotCell>),
diff --git a/components/malloc_size_of/lib.rs b/components/malloc_size_of/lib.rs
index b7f677f8044..52523af7cb1 100644
--- a/components/malloc_size_of/lib.rs
+++ b/components/malloc_size_of/lib.rs
@@ -746,6 +746,7 @@ malloc_size_of_is_0!(std::sync::atomic::AtomicUsize);
 malloc_size_of_is_0!(std::time::Duration);
 malloc_size_of_is_0!(std::time::Instant);
 malloc_size_of_is_0!(std::time::SystemTime);
+malloc_size_of_is_0!(style::data::ElementData);
 malloc_size_of_is_0!(style::font_face::SourceList);
 malloc_size_of_is_0!(style::properties::ComputedValues);
 malloc_size_of_is_0!(style::queries::values::PrefersColorScheme);
diff --git a/components/script/dom/node.rs b/components/script/dom/node.rs
index 2a01370085a..45a107ae673 100644
--- a/components/script/dom/node.rs
+++ b/components/script/dom/node.rs
@@ -167,7 +167,6 @@ pub struct Node {
 
     /// Layout data for this node. This is populated during layout and can
     /// be used for incremental relayout and script queries.
-    #[ignore_malloc_size_of = "trait object"]
     #[no_trace]
     layout_data: DomRefCell<Option<Box<GenericLayoutData>>>,
 }
diff --git a/components/script/script_thread.rs b/components/script/script_thread.rs
index c9b27bb6c56..231de76df1d 100644
--- a/components/script/script_thread.rs
+++ b/components/script/script_thread.rs
@@ -2439,8 +2439,6 @@ impl ScriptThread {
 
         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()
@@ -2448,6 +2446,9 @@ impl ScriptThread {
                     .collect_reports(&mut reports, ops);
             }
 
+            let prefix = format!("url({urls})");
+            reports.extend(self.get_cx().get_reports(prefix.clone(), ops));
+
             reports.push(self.image_cache.memory_report(&prefix, ops));
         });
 
diff --git a/components/shared/script_layout/lib.rs b/components/shared/script_layout/lib.rs
index 69e577e139d..233d02dfd47 100644
--- a/components/shared/script_layout/lib.rs
+++ b/components/shared/script_layout/lib.rs
@@ -27,7 +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::{MallocSizeOf as MallocSizeOfTrait, MallocSizeOfOps};
 use malloc_size_of_derive::MallocSizeOf;
 use net_traits::image_cache::{ImageCache, PendingImageId};
 use pixels::Image;
@@ -51,7 +51,11 @@ use style::selector_parser::{PseudoElement, RestyleDamage, Snapshot};
 use style::stylesheets::Stylesheet;
 use webrender_api::ImageKey;
 
-pub type GenericLayoutData = dyn Any + Send + Sync;
+pub trait GenericLayoutDataTrait: Any + MallocSizeOfTrait {
+    fn as_any(&self) -> &dyn Any;
+}
+
+pub type GenericLayoutData = dyn GenericLayoutDataTrait + Send + Sync;
 
 #[derive(MallocSizeOf)]
 pub struct StyleData {
@@ -59,7 +63,6 @@ pub struct StyleData {
     /// style system is being used standalone, this is all that hangs
     /// off the node. This must be first to permit the various
     /// transmutations between ElementData and PersistentLayoutData.
-    #[ignore_malloc_size_of = "This probably should not be ignored"]
     pub element_data: AtomicRefCell<ElementData>,
 
     /// Information needed during parallel traversals.
diff --git a/components/shared/script_layout/wrapper_traits.rs b/components/shared/script_layout/wrapper_traits.rs
index be27050a42f..6c4de339c1b 100644
--- a/components/shared/script_layout/wrapper_traits.rs
+++ b/components/shared/script_layout/wrapper_traits.rs
@@ -25,11 +25,11 @@ use style::selector_parser::{PseudoElement, PseudoElementCascadeType, SelectorIm
 use style::stylist::RuleInclusion;
 
 use crate::{
-    FragmentType, GenericLayoutData, HTMLCanvasData, HTMLMediaData, LayoutNodeType, SVGSVGData,
-    StyleData,
+    FragmentType, GenericLayoutData, GenericLayoutDataTrait, HTMLCanvasData, HTMLMediaData,
+    LayoutNodeType, SVGSVGData, StyleData,
 };
 
-pub trait LayoutDataTrait: Default + Send + Sync + 'static {}
+pub trait LayoutDataTrait: GenericLayoutDataTrait + Default + Send + Sync + 'static {}
 
 /// A wrapper so that layout can access only the methods that it should have access to. Layout must
 /// only ever see these and must never see instances of `LayoutDom`.