Refactor memory handling and add markers

tracing-trace/src/entry.rs

@@ -1,5 +1,4 @@
 use std::borrow::Cow;
-use std::ops::Sub;
 
 use serde::{Deserialize, Serialize};
 use tracing::span::Id as TracingId;
@@ -123,18 +122,16 @@ impl From<stats_alloc::Stats> for MemoryStats {
     }
 }
 
-impl Sub for MemoryStats {
+impl MemoryStats {
-    type Output = Self;
+    pub fn checked_sub(self, other: Self) -> Option<Self> {
-
+        Some(Self {
-    fn sub(self, other: Self) -> Self::Output {
+            allocations: self.allocations.checked_sub(other.allocations)?,
-        Self {
+            deallocations: self.deallocations.checked_sub(other.deallocations)?,
-            allocations: self.allocations - other.allocations,
+            reallocations: self.reallocations.checked_sub(other.reallocations)?,
-            deallocations: self.deallocations - other.deallocations,
+            bytes_allocated: self.bytes_allocated.checked_sub(other.bytes_allocated)?,
-            reallocations: self.reallocations - other.reallocations,
+            bytes_deallocated: self.bytes_deallocated.checked_sub(other.bytes_deallocated)?,
-            bytes_allocated: self.bytes_allocated - other.bytes_allocated,
+            bytes_reallocated: self.bytes_reallocated.checked_sub(other.bytes_reallocated)?,
-            bytes_deallocated: self.bytes_deallocated - other.bytes_deallocated,
+        })
-            bytes_reallocated: self.bytes_reallocated - other.bytes_reallocated,
-        }
     }
 }
 

tracing-trace/src/processor/firefox_profiler.rs

@@ -1,9 +1,9 @@
 use std::collections::HashMap;
 
 use fxprof_processed_profile::{
-    CategoryPairHandle, CpuDelta, Frame, FrameFlags, FrameInfo, MarkerDynamicField,
+    CategoryPairHandle, CounterHandle, CpuDelta, Frame, FrameFlags, FrameInfo, MarkerDynamicField,
-    MarkerFieldFormat, MarkerLocation, MarkerSchema, MarkerSchemaField, Profile, ProfilerMarker,
+    MarkerFieldFormat, MarkerLocation, MarkerSchema, MarkerSchemaField, ProcessHandle, Profile,
-    ReferenceTimestamp, SamplingInterval, StringHandle, Timestamp,
+    ProfilerMarker, ReferenceTimestamp, SamplingInterval, StringHandle, Timestamp,
 };
 use serde_json::json;
 
@@ -33,18 +33,8 @@ pub fn to_firefox_profile<R: std::io::Read>(
     let subcategory = profile.add_subcategory(category, "subcategory");
 
     let mut current_memory = MemoryStats::default();
-    let init_allocations = |profile: &mut Profile| {
+
-        profile.add_counter(main, "mimmalloc", "Memory", "Amount of allocation calls")
+    let mut memory_counters = None;
-    };
-    let init_deallocations = |profile: &mut Profile| {
-        profile.add_counter(main, "mimmalloc", "Memory", "Amount of deallocation calls")
-    };
-    let init_reallocations = |profile: &mut Profile| {
-        profile.add_counter(main, "mimmalloc", "Memory", "Amount of reallocation calls")
-    };
-    let mut allocations_counter = None;
-    let mut deallocations_counter = None;
-    let mut reallocations_counter = None;
 
     for entry in trace {
         let entry = entry?;
@@ -75,59 +65,29 @@ pub fn to_firefox_profile<R: std::io::Read>(
                     continue;
                 };
 
-                *status = SpanStatus::Inside(time);
+                *status = SpanStatus::Inside { time, memory };
 
                 last_timestamp = Timestamp::from_nanos_since_reference(time.as_nanos() as u64);
 
-                if let Some(stats) = memory {
+                add_memory_samples(
-                    let MemoryStats {
+                    &mut profile,
-                        allocations,
+                    main,
-                        deallocations,
+                    memory,
-                        reallocations,
+                    last_timestamp,
-                        bytes_allocated,
+                    &mut memory_counters,
-                        bytes_deallocated,
+                    &mut current_memory,
-                        bytes_reallocated,
+                );
-                    } = current_memory - stats;
-
-                    let counter =
-                        *allocations_counter.get_or_insert_with(|| init_allocations(&mut profile));
-                    profile.add_counter_sample(
-                        counter,
-                        last_timestamp,
-                        bytes_allocated as f64,
-                        allocations.try_into().unwrap(),
-                    );
-
-                    let counter = *deallocations_counter
-                        .get_or_insert_with(|| init_deallocations(&mut profile));
-                    profile.add_counter_sample(
-                        counter,
-                        last_timestamp,
-                        bytes_deallocated as f64,
-                        deallocations.try_into().unwrap(),
-                    );
-
-                    let counter = *reallocations_counter
-                        .get_or_insert_with(|| init_reallocations(&mut profile));
-                    profile.add_counter_sample(
-                        counter,
-                        last_timestamp,
-                        bytes_reallocated as f64,
-                        reallocations.try_into().unwrap(),
-                    );
-
-                    current_memory = stats;
-                }
             }
             Entry::SpanExit(SpanExit { id, time, memory }) => {
                 let (span, status) = spans.get_mut(&id).unwrap();
 
-                let SpanStatus::Inside(begin) = status else {
+                let SpanStatus::Inside { time: begin, memory: begin_memory } = status else {
                     continue;
                 };
                 last_timestamp = Timestamp::from_nanos_since_reference(time.as_nanos() as u64);
 
                 let begin = *begin;
+                let begin_memory = *begin_memory;
 
                 *status = SpanStatus::Outside;
 
@@ -151,49 +111,20 @@ pub fn to_firefox_profile<R: std::io::Read>(
                     1,
                 );
 
-                if let Some(stats) = memory {
+                add_memory_samples(
-                    let MemoryStats {
+                    &mut profile,
-                        allocations,
+                    main,
-                        deallocations,
+                    memory,
-                        reallocations,
+                    last_timestamp,
-                        bytes_allocated,
+                    &mut memory_counters,
-                        bytes_deallocated,
+                    &mut current_memory,
-                        bytes_reallocated,
+                );
-                    } = current_memory - stats;
-
-                    let counter =
-                        *allocations_counter.get_or_insert_with(|| init_allocations(&mut profile));
-                    profile.add_counter_sample(
-                        counter,
-                        last_timestamp,
-                        bytes_allocated as f64,
-                        allocations.try_into().unwrap(),
-                    );
-
-                    let counter = *deallocations_counter
-                        .get_or_insert_with(|| init_deallocations(&mut profile));
-                    profile.add_counter_sample(
-                        counter,
-                        last_timestamp,
-                        bytes_deallocated as f64,
-                        deallocations.try_into().unwrap(),
-                    );
-
-                    let counter = *reallocations_counter
-                        .get_or_insert_with(|| init_reallocations(&mut profile));
-                    profile.add_counter_sample(
-                        counter,
-                        last_timestamp,
-                        bytes_reallocated as f64,
-                        reallocations.try_into().unwrap(),
-                    );
-
-                    current_memory = stats;
-                }
 
                 let (callsite, _) = calls.get(&span.call_id).unwrap();
 
-                let marker = SpanMarker { callsite, span: &span };
+                let memory_delta =
+                    begin_memory.zip(memory).and_then(|(begin, end)| end.checked_sub(begin));
+                let marker = SpanMarker { callsite, span: &span, memory_delta };
 
                 profile.add_marker_with_stack(
                     *thread_handle,
@@ -216,6 +147,77 @@ pub fn to_firefox_profile<R: std::io::Read>(
     Ok(profile)
 }
 
+struct MemoryCounterHandles {
+    allocations: CounterHandle,
+    deallocations: CounterHandle,
+    reallocations: CounterHandle,
+}
+
+impl MemoryCounterHandles {
+    fn new(profile: &mut Profile, main: ProcessHandle) -> Self {
+        let allocations =
+            profile.add_counter(main, "mimmalloc", "Memory", "Amount of allocated memory");
+        let deallocations =
+            profile.add_counter(main, "mimmalloc", "Memory", "Amount of deallocated memory");
+        let reallocations =
+            profile.add_counter(main, "mimmalloc", "Memory", "Amount of reallocated memory");
+        Self { allocations, deallocations, reallocations }
+    }
+}
+
+fn add_memory_samples(
+    profile: &mut Profile,
+    main: ProcessHandle,
+    memory: Option<MemoryStats>,
+    last_timestamp: Timestamp,
+    memory_counters: &mut Option<MemoryCounterHandles>,
+    current_memory: &mut MemoryStats,
+) {
+    let Some(stats) = memory else {
+        return;
+    };
+
+    let Some(MemoryStats {
+        allocations,
+        deallocations,
+        reallocations,
+        bytes_allocated,
+        bytes_deallocated,
+        bytes_reallocated,
+    }) = stats.checked_sub(*current_memory)
+    else {
+        // since spans are recorded out-of-order it is possible they are not always monotonic.
+        // We ignore spans that made no difference.
+        return;
+    };
+
+    let memory_counters =
+        memory_counters.get_or_insert_with(|| MemoryCounterHandles::new(profile, main));
+
+    profile.add_counter_sample(
+        memory_counters.allocations,
+        last_timestamp,
+        bytes_allocated as f64,
+        allocations.try_into().unwrap(),
+    );
+
+    profile.add_counter_sample(
+        memory_counters.deallocations,
+        last_timestamp,
+        bytes_deallocated as f64,
+        deallocations.try_into().unwrap(),
+    );
+
+    profile.add_counter_sample(
+        memory_counters.reallocations,
+        last_timestamp,
+        bytes_reallocated as f64,
+        reallocations.try_into().unwrap(),
+    );
+
+    *current_memory = stats;
+}
+
 fn to_timestamp(time: std::time::Duration) -> Timestamp {
     Timestamp::from_nanos_since_reference(time.as_nanos() as u64)
 }
@@ -252,12 +254,13 @@ fn make_frame(
 #[derive(Debug, Clone, Copy)]
 enum SpanStatus {
     Outside,
-    Inside(std::time::Duration),
+    Inside { time: std::time::Duration, memory: Option<MemoryStats> },
 }
 
 struct SpanMarker<'a> {
     span: &'a NewSpan,
     callsite: &'a NewCallsite,
+    memory_delta: Option<MemoryStats>,
 }
 
 impl<'a> ProfilerMarker for SpanMarker<'a> {
@@ -295,6 +298,42 @@ impl<'a> ProfilerMarker for SpanMarker<'a> {
                 format: MarkerFieldFormat::Integer,
                 searchable: true,
             }),
+            MarkerSchemaField::Dynamic(MarkerDynamicField {
+                key: "allocations",
+                label: "Number of allocation operations while this function was executing",
+                format: MarkerFieldFormat::Integer,
+                searchable: false,
+            }),
+            MarkerSchemaField::Dynamic(MarkerDynamicField {
+                key: "deallocations",
+                label: "Number of deallocation operations while this function was executing",
+                format: MarkerFieldFormat::Integer,
+                searchable: false,
+            }),
+            MarkerSchemaField::Dynamic(MarkerDynamicField {
+                key: "reallocations",
+                label: "Number of reallocation operations while this function was executing",
+                format: MarkerFieldFormat::Integer,
+                searchable: false,
+            }),
+            MarkerSchemaField::Dynamic(MarkerDynamicField {
+                key: "allocated_bytes",
+                label: "Number of allocated bytes while this function was executing",
+                format: MarkerFieldFormat::Bytes,
+                searchable: false,
+            }),
+            MarkerSchemaField::Dynamic(MarkerDynamicField {
+                key: "deallocated_bytes",
+                label: "Number of deallocated bytes while this function was executing",
+                format: MarkerFieldFormat::Bytes,
+                searchable: false,
+            }),
+            MarkerSchemaField::Dynamic(MarkerDynamicField {
+                key: "reallocated_bytes",
+                label: "Number of reallocated bytes while this function was executing",
+                format: MarkerFieldFormat::Bytes,
+                searchable: false,
+            }),
         ];
 
         MarkerSchema {
@@ -317,13 +356,33 @@ impl<'a> ProfilerMarker for SpanMarker<'a> {
         let module_path = self.callsite.module_path.as_deref();
         let span_id = self.span.id;
         let thread_id = self.span.thread_id;
-        json!({
+
+        let mut value = json!({
             "type": Self::MARKER_TYPE_NAME,
             "filename": filename,
             "line": line,
             "module_path": module_path,
             "span_id": span_id,
             "thread_id": thread_id,
-        })
+        });
+
+        if let Some(MemoryStats {
+            allocations,
+            deallocations,
+            reallocations,
+            bytes_allocated,
+            bytes_deallocated,
+            bytes_reallocated,
+        }) = self.memory_delta
+        {
+            value["allocations"] = json!(allocations);
+            value["deallocations"] = json!(deallocations);
+            value["reallocations"] = json!(reallocations);
+            value["bytes_allocated"] = json!(bytes_allocated);
+            value["bytes_deallocated"] = json!(bytes_deallocated);
+            value["bytes_reallocated"] = json!(bytes_reallocated);
+        }
+
+        value
     }
 }