5 files changed, 70 insertions, 3 deletions

diff --git a/components/layout/context.rs b/components/layout/context.rs
index 6e0aaa83dd6..b5ed8b729bc 100644
--- a/components/layout/context.rs
+++ b/components/layout/context.rs
@@ -29,6 +29,7 @@ use std::sync::mpsc::{channel, Sender};
 use style::selector_matching::Stylist;
 use url::Url;
 use util::geometry::Au;
+use util::mem::HeapSizeOf;
 use util::opts;
 
 struct LocalLayoutContext {
@@ -37,8 +38,21 @@ struct LocalLayoutContext {
     style_sharing_candidate_cache: RefCell<StyleSharingCandidateCache>,
 }
 
+impl HeapSizeOf for LocalLayoutContext {
+    // FIXME(njn): measure other fields eventually.
+    fn heap_size_of_children(&self) -> usize {
+        self.font_context.heap_size_of_children()
+    }
+}
+
 thread_local!(static LOCAL_CONTEXT_KEY: RefCell<Option<Rc<LocalLayoutContext>>> = RefCell::new(None));
 
+pub fn heap_size_of_local_context() -> usize {
+    LOCAL_CONTEXT_KEY.with(|r| {
+        r.borrow().clone().map_or(0, |context| context.heap_size_of_children())
+    })
+}
+
 fn create_or_get_local_context(shared_layout_context: &SharedLayoutContext)
                                -> Rc<LocalLayoutContext> {
     LOCAL_CONTEXT_KEY.with(|r| {
diff --git a/components/layout/layout_task.rs b/components/layout/layout_task.rs
index 2756780dfb2..f769f7ee3d7 100644
--- a/components/layout/layout_task.rs
+++ b/components/layout/layout_task.rs
@@ -9,7 +9,7 @@
 
 use animation;
 use construct::ConstructionResult;
-use context::{SharedLayoutContext, SharedLayoutContextWrapper};
+use context::{SharedLayoutContext, SharedLayoutContextWrapper, heap_size_of_local_context};
 use css::node_style::StyledNode;
 use data::{LayoutDataAccess, LayoutDataWrapper};
 use display_list_builder::ToGfxColor;
@@ -575,10 +575,28 @@ impl LayoutTask {
         let rw_data = self.lock_rw_data(possibly_locked_rw_data);
         let stacking_context = rw_data.stacking_context.as_ref();
         reports.push(Report {
-            path: path!["pages", format!("url({})", self.url), "display-list"],
+            path: path!["pages", format!("url({})", self.url), "layout-task", "display-list"],
             size: stacking_context.map_or(0, |sc| sc.heap_size_of_children()),
         });
 
+        // The LayoutTask has a context in TLS...
+        reports.push(Report {
+            path: path!["pages", format!("url({})", self.url), "layout-task", "local-context"],
+            size: heap_size_of_local_context(),
+        });
+
+        // ... as do each of the LayoutWorkers, if present.
+        if let Some(ref traversal) = rw_data.parallel_traversal {
+            let sizes = traversal.heap_size_of_tls(heap_size_of_local_context);
+            for (i, size) in sizes.iter().enumerate() {
+                reports.push(Report {
+                    path: path!["pages", format!("url({})", self.url),
+                                format!("layout-worker-{}-local-context", i)],
+                    size: *size
+                });
+            }
+        }
+
         reports_chan.send(reports);
     }
 
diff --git a/components/profile/mem.rs b/components/profile/mem.rs
index 79437c2ba17..e3330531db0 100644
--- a/components/profile/mem.rs
+++ b/components/profile/mem.rs
@@ -224,7 +224,7 @@ impl ReportsTree {
         }
 
         let mebi = 1024f64 * 1024f64;
-        let count_str = if self.count > 1 { format!(" {}", self.count) } else { "".to_owned() };
+        let count_str = if self.count > 1 { format!(" [{}]", self.count) } else { "".to_owned() };
         println!("|{}{:8.2} MiB -- {}{}",
                  indent_str, (self.size as f64) / mebi, self.path_seg, count_str);
 
diff --git a/components/util/mem.rs b/components/util/mem.rs
index 5ecd511bdcf..0170a1b359f 100644
--- a/components/util/mem.rs
+++ b/components/util/mem.rs
@@ -5,6 +5,7 @@
 //! Data structure measurement.
 
 use libc::{c_void, size_t};
+use std::cell::RefCell;
 use std::collections::LinkedList;
 use std::mem::transmute;
 use std::sync::Arc;
@@ -92,6 +93,12 @@ impl<T: HeapSizeOf> HeapSizeOf for Arc<T> {
     }
 }
 
+impl<T: HeapSizeOf> HeapSizeOf for RefCell<T> {
+    fn heap_size_of_children(&self) -> usize {
+        self.borrow().heap_size_of_children()
+    }
+}
+
 impl<T: HeapSizeOf> HeapSizeOf for Vec<T> {
     fn heap_size_of_children(&self) -> usize {
         heap_size_of(self.as_ptr() as *const c_void) +
diff --git a/components/util/workqueue.rs b/components/util/workqueue.rs
index 2a76adb51eb..2f0523a09ef 100644
--- a/components/util/workqueue.rs
+++ b/components/util/workqueue.rs
@@ -36,6 +36,8 @@ enum WorkerMsg<QueueData: 'static, WorkData: 'static> {
     Start(Worker<WorkUnit<QueueData, WorkData>>, *mut AtomicUsize, *const QueueData),
     /// Tells the worker to stop. It can be restarted again with a `WorkerMsg::Start`.
     Stop,
+    /// Tells the worker to measure the heap size of its TLS using the supplied function.
+    HeapSizeOfTLS(fn() -> usize),
     /// Tells the worker thread to terminate.
     Exit,
 }
@@ -45,6 +47,7 @@ unsafe impl<QueueData: 'static, WorkData: 'static> Send for WorkerMsg<QueueData,
 /// Messages to the supervisor.
 enum SupervisorMsg<QueueData: 'static, WorkData: 'static> {
     Finished,
+    HeapSizeOfTLS(usize),
     ReturnDeque(usize, Worker<WorkUnit<QueueData, WorkData>>),
 }
 
@@ -102,6 +105,10 @@ impl<QueueData: Send, WorkData: Send> WorkerThread<QueueData, WorkData> {
                 WorkerMsg::Start(deque, ref_count, queue_data) => (deque, ref_count, queue_data),
                 WorkerMsg::Stop => panic!("unexpected stop message"),
                 WorkerMsg::Exit => return,
+                WorkerMsg::HeapSizeOfTLS(f) => {
+                    self.chan.send(SupervisorMsg::HeapSizeOfTLS(f())).unwrap();
+                    continue;
+                }
             };
 
             let mut back_off_sleep = 0 as u32;
@@ -335,11 +342,32 @@ impl<QueueData: Send, WorkData: Send> WorkQueue<QueueData, WorkData> {
         for _ in 0..self.workers.len() {
             match self.port.recv().unwrap() {
                 SupervisorMsg::ReturnDeque(index, deque) => self.workers[index].deque = Some(deque),
+                SupervisorMsg::HeapSizeOfTLS(_) => panic!("unexpected HeapSizeOfTLS message"),
                 SupervisorMsg::Finished => panic!("unexpected finished message!"),
             }
         }
     }
 
+    /// Synchronously measure memory usage of any thread-local storage.
+    pub fn heap_size_of_tls(&self, f: fn() -> usize) -> Vec<usize> {
+        // Tell the workers to measure themselves.
+        for worker in self.workers.iter() {
+            worker.chan.send(WorkerMsg::HeapSizeOfTLS(f)).unwrap()
+        }
+
+        // Wait for the workers to finish measuring themselves.
+        let mut sizes = vec![];
+        for _ in 0..self.workers.len() {
+            match self.port.recv().unwrap() {
+                SupervisorMsg::HeapSizeOfTLS(size) => {
+                    sizes.push(size);
+                }
+                _ => panic!("unexpected message!"),
+            }
+        }
+        sizes
+    }
+
     pub fn shutdown(&mut self) {
         for worker in self.workers.iter() {
             worker.chan.send(WorkerMsg::Exit).unwrap()