From cc79cd0b049e123e56f20c0c5f49e73d070d0b7c Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Tue, 23 Jan 2024 13:03:30 +0100
Subject: [PATCH] Switch to a single view indicating current usage

---
 tracing-trace/src/entry.rs                    |  8 +++
 .../src/processor/firefox_profiler.rs         | 58 ++++---------------
 2 files changed, 20 insertions(+), 46 deletions(-)

diff --git a/tracing-trace/src/entry.rs b/tracing-trace/src/entry.rs
index d92ff09b8..34b3aaa52 100644
--- a/tracing-trace/src/entry.rs
+++ b/tracing-trace/src/entry.rs
@@ -133,6 +133,14 @@ impl MemoryStats {
             bytes_reallocated: self.bytes_reallocated.checked_sub(other.bytes_reallocated)?,
         })
     }
+
+    pub fn usage(&self) -> isize {
+        (self.bytes_allocated - self.bytes_deallocated) as isize + self.bytes_reallocated
+    }
+
+    pub fn operations(&self) -> usize {
+        self.allocations + self.deallocations + self.reallocations
+    }
 }
 
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
diff --git a/tracing-trace/src/processor/firefox_profiler.rs b/tracing-trace/src/processor/firefox_profiler.rs
index eda9f6467..dd56939a4 100644
--- a/tracing-trace/src/processor/firefox_profiler.rs
+++ b/tracing-trace/src/processor/firefox_profiler.rs
@@ -148,20 +148,14 @@ pub fn to_firefox_profile<R: std::io::Read>(
 }
 
 struct MemoryCounterHandles {
-    allocations: CounterHandle,
-    deallocations: CounterHandle,
-    reallocations: CounterHandle,
+    usage: 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 }
+        let usage =
+            profile.add_counter(main, "mimmalloc", "Memory", "Amount of memory currently in use");
+        Self { usage }
     }
 }
 
@@ -171,51 +165,23 @@ fn add_memory_samples(
     memory: Option<MemoryStats>,
     last_timestamp: Timestamp,
     memory_counters: &mut Option<MemoryCounterHandles>,
-    current_memory: &mut MemoryStats,
+    last_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,
+        memory_counters.usage,
         last_timestamp,
-        bytes_allocated as f64,
-        allocations.try_into().unwrap(),
+        stats.usage() as f64 - last_memory.usage() as f64,
+        stats.operations().checked_sub(last_memory.operations()).unwrap_or_default() as u32,
     );
 
-    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;
+    *last_memory = stats;
 }
 
 fn to_timestamp(time: std::time::Duration) -> Timestamp {
@@ -378,9 +344,9 @@ impl<'a> ProfilerMarker for SpanMarker<'a> {
             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["allocated_bytes"] = json!(bytes_allocated);
+            value["deallocated_bytes"] = json!(bytes_deallocated);
+            value["reallocated_bytes"] = json!(bytes_reallocated);
         }
 
         value