trace memory allocations and show in tracing (#6866)

### Description * collect memory allocation data and report it when tracing is enable * Add value mode to trace server to show memory info Closes PACK-2184
vercel · Jan 10, 2024 · 5d2b0fe · 5d2b0fe · vercel · Jan 10, 2024
1 parent b611f41
commit 5d2b0fe
Show file tree

Hide file tree

Showing 11 changed files with 458 additions and 41 deletions.
diff --git a/crates/turbo-tasks-malloc/src/counter.rs b/crates/turbo-tasks-malloc/src/counter.rs
@@ -13,16 +13,37 @@ const TARGET_BUFFER: usize = 100 * KB;
 /// global counter.
 const MAX_BUFFER: usize = 200 * KB;
 
+#[derive(Default)]
+pub struct AllocationInfo {
+    pub allocations: usize,
+    pub deallocations: usize,
+    pub allocation_count: usize,
+    pub deallocation_count: usize,
+}
+
+impl AllocationInfo {
+    pub fn is_empty(&self) -> bool {
+        self.allocations == 0
+            && self.deallocations == 0
+            && self.allocation_count == 0
+            && self.deallocation_count == 0
+    }
+}
+
+#[derive(Default)]
 struct ThreadLocalCounter {
     /// Thread-local buffer of allocated bytes that have been added to the
     /// global counter desprite not being allocated yet. It is unsigned so that
     /// means the global counter is always equal or greater than the real
     /// value.
     buffer: usize,
+    allocation_info: AllocationInfo,
 }
 
 impl ThreadLocalCounter {
     fn add(&mut self, size: usize) {
+        self.allocation_info.allocations += size;
+        self.allocation_info.allocation_count += 1;
         if self.buffer >= size {
             self.buffer -= size;
         } else {
@@ -33,6 +54,8 @@ impl ThreadLocalCounter {
     }
 
     fn remove(&mut self, size: usize) {
+        self.allocation_info.deallocations += size;
+        self.allocation_info.deallocation_count += 1;
         self.buffer += size;
         if self.buffer > MAX_BUFFER {
             let offset = self.buffer - TARGET_BUFFER;
@@ -50,19 +73,23 @@ impl ThreadLocalCounter {
 }
 
 thread_local! {
-  static LOCAL_COUNTER: UnsafeCell<ThreadLocalCounter> = UnsafeCell::new(ThreadLocalCounter { buffer: 0 });
+  static LOCAL_COUNTER: UnsafeCell<ThreadLocalCounter> = UnsafeCell::new(ThreadLocalCounter::default());
 }
 
 pub fn get() -> usize {
     ALLOCATED.load(Ordering::Relaxed)
 }
 
-fn with_local_counter(f: impl FnOnce(&mut ThreadLocalCounter)) {
+pub fn pop_allocations() -> AllocationInfo {
+    with_local_counter(|local| std::mem::take(&mut local.allocation_info))
+}
+
+fn with_local_counter<T>(f: impl FnOnce(&mut ThreadLocalCounter) -> T) -> T {
     LOCAL_COUNTER.with(|local| {
         let ptr = local.get();
         // SAFETY: This is a thread local.
         let mut local = unsafe { NonNull::new_unchecked(ptr) };
-        f(unsafe { local.as_mut() });
+        f(unsafe { local.as_mut() })
     })
 }
 

diff --git a/crates/turbo-tasks-malloc/src/lib.rs b/crates/turbo-tasks-malloc/src/lib.rs
@@ -16,6 +16,10 @@ impl TurboMalloc {
     pub fn thread_stop() {
         flush();
     }
+
+    pub fn pop_allocations() -> self::counter::AllocationInfo {
+        self::counter::pop_allocations()
+    }
 }
 
 #[cfg(all(

diff --git a/crates/turbopack-convert-trace/src/main.rs b/crates/turbopack-convert-trace/src/main.rs
@@ -344,6 +344,9 @@ fn main() {
                 let parent = &mut spans[internal_parent];
                 parent.items.push(SpanItem::Child(internal_id));
             }
+            TraceRow::Allocation { .. } => {
+                // ignore
+            }
         }
     }
 

diff --git a/crates/turbopack-trace-server/src/reader.rs b/crates/turbopack-trace-server/src/reader.rs
@@ -253,6 +253,7 @@ fn process(store: &mut StoreWriteGuard, state: &mut ReaderState, row: TraceRow<'
                 .remove("name")
                 .and_then(|v| v.as_str().map(|s| s.to_string()))
                 .unwrap_or("event".into());
+
             let id = store.add_span(
                 parent,
                 ts,
@@ -271,6 +272,34 @@ fn process(store: &mut StoreWriteGuard, state: &mut ReaderState, row: TraceRow<'
                 &mut state.outdated_spans,
             );
         }
+        TraceRow::Allocation {
+            ts: _,
+            thread_id,
+            allocations,
+            allocation_count,
+            deallocations,
+            deallocation_count,
+        } => {
+            let stack = state.thread_stacks.entry(thread_id).or_default();
+            if let Some(&id) = stack.last() {
+                if allocations > 0 {
+                    store.add_allocation(
+                        id,
+                        allocations,
+                        allocation_count,
+                        &mut state.outdated_spans,
+                    );
+                }
+                if deallocations > 0 {
+                    store.add_deallocation(
+                        id,
+                        deallocations,
+                        deallocation_count,
+                        &mut state.outdated_spans,
+                    );
+                }
+            }
+        }
     }
 }
 

diff --git a/crates/turbopack-trace-server/src/server.rs b/crates/turbopack-trace-server/src/server.rs
@@ -41,6 +41,9 @@ pub enum ServerToClientMessage {
         is_graph: bool,
         start: u64,
         end: u64,
+        duration: u64,
+        allocations: u64,
+        deallocations: u64,
         args: Vec<(String, String)>,
         path: Vec<String>,
     },
@@ -90,6 +93,7 @@ pub struct ViewRect {
     pub horizontal_pixels: u64,
     pub query: String,
     pub view_mode: String,
+    pub value_mode: String,
 }
 
 struct ConnectionState {
@@ -133,6 +137,7 @@ pub fn serve(store: Arc<StoreContainer>) -> Result<()> {
                         horizontal_pixels: 1,
                         query: String::new(),
                         view_mode: "aggregated".to_string(),
+                        value_mode: "duration".to_string(),
                     },
                     last_update_generation: 0,
                 }));
@@ -270,6 +275,9 @@ pub fn serve(store: Arc<StoreContainer>) -> Result<()> {
                                     {
                                         let span_start = span.start();
                                         let span_end = span.end();
+                                        let duration = span.corrected_total_time();
+                                        let allocations = span.total_allocations();
+                                        let deallocations = span.total_deallocations();
                                         let args = span
                                             .args()
                                             .map(|(k, v)| (k.to_string(), v.to_string()))
@@ -286,6 +294,9 @@ pub fn serve(store: Arc<StoreContainer>) -> Result<()> {
                                             is_graph,
                                             start: span_start,
                                             end: span_end,
+                                            duration,
+                                            allocations,
+                                            deallocations,
                                             args,
                                             path,
                                         }
@@ -295,6 +306,9 @@ pub fn serve(store: Arc<StoreContainer>) -> Result<()> {
                                             is_graph: false,
                                             start: 0,
                                             end: 0,
+                                            duration: 0,
+                                            allocations: 0,
+                                            deallocations: 0,
                                             args: Vec::new(),
                                             path: Vec::new(),
                                         }

diff --git a/crates/turbopack-trace-server/src/span.rs b/crates/turbopack-trace-server/src/span.rs
@@ -23,13 +23,21 @@ pub struct Span {
     // These values are computed automatically:
     pub self_end: u64,
     pub self_time: u64,
+    pub self_allocations: u64,
+    pub self_allocation_count: u64,
+    pub self_deallocations: u64,
+    pub self_deallocation_count: u64,
 
     // These values are computed when accessed (and maybe deleted during writing):
     pub end: OnceLock<u64>,
     pub nice_name: OnceLock<(String, String)>,
     pub group_name: OnceLock<String>,
     pub max_depth: OnceLock<u32>,
     pub total_time: OnceLock<u64>,
+    pub total_allocations: OnceLock<u64>,
+    pub total_deallocations: OnceLock<u64>,
+    pub total_persistent_allocations: OnceLock<u64>,
+    pub total_allocation_count: OnceLock<u64>,
     pub corrected_self_time: OnceLock<u64>,
     pub corrected_total_time: OnceLock<u64>,
     pub graph: OnceLock<Vec<SpanGraphEvent>>,
@@ -63,7 +71,15 @@ pub struct SpanGraph {
     pub max_depth: OnceLock<u32>,
     pub events: OnceLock<Vec<SpanGraphEvent>>,
     pub self_time: OnceLock<u64>,
+    pub self_allocations: OnceLock<u64>,
+    pub self_deallocations: OnceLock<u64>,
+    pub self_persistent_allocations: OnceLock<u64>,
+    pub self_allocation_count: OnceLock<u64>,
     pub total_time: OnceLock<u64>,
+    pub total_allocations: OnceLock<u64>,
+    pub total_deallocations: OnceLock<u64>,
+    pub total_persistent_allocations: OnceLock<u64>,
+    pub total_allocation_count: OnceLock<u64>,
     pub corrected_self_time: OnceLock<u64>,
     pub corrected_total_time: OnceLock<u64>,
 }