Catch incorrect uses of Profile::add_allocation_sample. (#629)

Jeff was running into wrong stacks because he was calling
Profile::add_allocation_sample for a thread which was not the main
thread, and then we were using stack handles which aren't valid on the
main thread, and the existing assertions were insufficient to catch
this.

Change the API to accept a process instead of a thread, and widen the
assertions.

Author: mstange

fxprof-processed-profile/src/process.rs

@@ -15,6 +15,7 @@ pub struct Process {
     pid: String,
     name: String,
     threads: Vec<ThreadHandle>,
+    main_thread: Option<ThreadHandle>,
     start_time: Timestamp,
     end_time: Option<Timestamp>,
     libs: LibMappings<LibraryHandle>,
@@ -25,6 +26,7 @@ impl Process {
         Self {
             pid,
             threads: Vec::new(),
+            main_thread: None,
             libs: LibMappings::new(),
             start_time,
             end_time: None,
@@ -33,7 +35,7 @@ impl Process {
     }
 
     pub fn thread_handle_for_allocations(&self) -> Option<ThreadHandle> {
-        self.threads.first().cloned()
+        self.main_thread
     }
 
     pub fn set_start_time(&mut self, start_time: Timestamp) {
@@ -60,8 +62,11 @@ impl Process {
         &self.name
     }
 
-    pub fn add_thread(&mut self, thread: ThreadHandle) {
+    pub fn add_thread(&mut self, thread: ThreadHandle, is_main: bool) {
         self.threads.push(thread);
+        if is_main && self.main_thread.is_none() {
+            self.main_thread = Some(thread)
+        }
     }
 
     pub fn pid(&self) -> &str {

fxprof-processed-profile/src/profile.rs

@@ -481,7 +481,7 @@ impl Profile {
         let handle = ThreadHandle(self.threads.len());
         self.threads
             .push(Thread::new(process, tid, start_time, is_main));
-        self.processes[process.0].add_thread(handle);
+        self.processes[process.0].add_thread(handle, is_main);
         handle
     }
 
@@ -926,8 +926,12 @@ impl Profile {
         self.threads[thread.0].add_sample_same_stack_zero_cpu(timestamp, weight);
     }
 
-    /// Add an allocation or deallocation sample to the given thread. This is used
-    /// to collect stacks showing where allocations and deallocations happened.
+    /// Add an allocation or deallocation sample to the *main* thread of the given
+    /// process. This is used to collect stacks showing where allocations and
+    /// deallocations happened.
+    ///
+    /// CANNOT BE CALLED BEFORE THE MAIN THREAD FOR `process` HAS BEEN CREATED.
+    /// `stack` MUST BE A STACK HANDLE WHICH IS VALID FOR THAT MAIN THREAD.
     ///
     /// When loading profiles with allocation samples in the Firefox Profiler, the
     /// UI will display a dropdown above the call tree to switch between regular
@@ -951,40 +955,40 @@ impl Profile {
     ///
     /// To get the stack handle, you can use [`Profile::handle_for_stack`] or
     /// [`Profile::handle_for_stack_frames`].
+    ///
+    /// ## Main thread requirement
+    ///
+    /// Allocations are per-process, because you can allocate something one one thread
+    /// and then free it on a different thread, and you'll still want those two operations
+    /// to be matched up in a view that shows the retained memory.
+    ///
+    /// Unfortunately, `StackHandle` is currently per-thread. This method will become more
+    /// ergonomic once the profile format has changed so that `StackHandle`s can be used
+    /// across all threads of a profile. In the meantime, unfortunately you must manually
+    /// ensure that you create the stack handle for the main thread of the given process.
     pub fn add_allocation_sample(
         &mut self,
-        thread: ThreadHandle,
+        process: ProcessHandle,
         timestamp: Timestamp,
         stack: Option<StackHandle>,
         allocation_address: u64,
         allocation_size: i64,
     ) {
-        // The profile format strictly separates sample data from different threads.
-        // For allocation samples, this separation is a bit unfortunate, especially
-        // when it comes to the "Retained Memory" panel which shows allocation stacks
-        // for just objects that haven't been deallocated yet. This panel is per-thread,
-        // and it needs to know about deallocations even if they happened on a different
-        // thread from the allocation.
-        // To resolve this conundrum, for now, we will put all allocation and deallocation
-        // samples on a single thread per process, regardless of what thread they actually
-        // happened on.
-        // The Gecko profiler puts all allocation samples on the main thread, for example.
-        // Here in fxprof-processed-profile, we just deem the first thread of each process
-        // as the processes "allocation thread".
-        let process_handle = self.threads[thread.0].process();
-        let process = &self.processes[process_handle.0];
-        let allocation_thread_handle = process.thread_handle_for_allocations().unwrap();
+        let process = &self.processes[process.0];
+        let Some(allocation_thread) = process.thread_handle_for_allocations() else {
+            panic!("Profile::add_allocation_sample called for a thread whose process does not have a main thread");
+        };
         let stack_index = match stack {
-            Some(StackHandle(stack_thread_handle, stack_index)) => {
+            Some(StackHandle(stack_thread, stack_index)) => {
                 assert_eq!(
-                    stack_thread_handle, thread,
-                    "StackHandle from different thread passed to Profile::add_sample"
+                    stack_thread, allocation_thread,
+                    "StackHandle from different thread passed to Profile::add_allocation_sample"
                 );
                 Some(stack_index)
             }
             None => None,
         };
-        self.threads[allocation_thread_handle.0].add_allocation_sample(
+        self.threads[allocation_thread.0].add_allocation_sample(
             timestamp,
             stack_index,
             allocation_address,