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>

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)
+        LayoutNode::layout_data(&self).map(|data| {
-            .map(|data| data.downcast_ref::<DOMLayoutData>().unwrap().0.borrow())
+            data.as_any()
+                .downcast_ref::<DOMLayoutData>()
+                .unwrap()
+                .0
+                .borrow()
+        })
     }
 
     fn element_box_slot(&self) -> BoxSlot<'dom> {

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>),

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

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>>>,
 }

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

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.

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,
+    FragmentType, GenericLayoutData, GenericLayoutDataTrait, HTMLCanvasData, HTMLMediaData,
-    StyleData,
+    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`.