Logging the memory usage over time

tracing-trace/src/entry.rs

@@ -1,4 +1,5 @@
 use std::borrow::Cow;
+use std::ops::Sub;
 
 use serde::{Deserialize, Serialize};
 use tracing::span::Id as TracingId;
@@ -64,12 +65,14 @@ pub struct NewThread {
 pub struct SpanEnter {
     pub id: SpanId,
     pub time: std::time::Duration,
+    pub memory: Option<MemoryStats>,
 }
 
 #[derive(Clone, Copy, Debug, Serialize, Deserialize)]
 pub struct SpanExit {
     pub id: SpanId,
     pub time: std::time::Duration,
+    pub memory: Option<MemoryStats>,
 }
 
 #[derive(Clone, Copy, Debug, Serialize, Deserialize)]
@@ -86,6 +89,55 @@ pub struct SpanClose {
     pub time: std::time::Duration,
 }
 
+/// A struct with a lot of memory allocation stats akin
+/// to the `stats_alloc::Stats` one but implements the
+/// `Serialize/Deserialize` serde traits.
+#[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)]
+pub struct MemoryStats {
+    pub allocations: usize,
+    pub deallocations: usize,
+    pub reallocations: usize,
+    pub bytes_allocated: usize,
+    pub bytes_deallocated: usize,
+    pub bytes_reallocated: isize,
+}
+
+impl From<stats_alloc::Stats> for MemoryStats {
+    fn from(stats: stats_alloc::Stats) -> Self {
+        let stats_alloc::Stats {
+            allocations,
+            deallocations,
+            reallocations,
+            bytes_allocated,
+            bytes_deallocated,
+            bytes_reallocated,
+        } = stats;
+        MemoryStats {
+            allocations,
+            deallocations,
+            reallocations,
+            bytes_allocated,
+            bytes_deallocated,
+            bytes_reallocated,
+        }
+    }
+}
+
+impl Sub for MemoryStats {
+    type Output = Self;
+
+    fn sub(self, other: Self) -> Self::Output {
+        Self {
+            allocations: self.allocations - other.allocations,
+            deallocations: self.deallocations - other.deallocations,
+            reallocations: self.reallocations - other.reallocations,
+            bytes_allocated: self.bytes_allocated - other.bytes_allocated,
+            bytes_deallocated: self.bytes_deallocated - other.bytes_deallocated,
+            bytes_reallocated: self.bytes_reallocated - other.bytes_reallocated,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub struct ResourceId(pub(crate) usize);
 

tracing-trace/src/layer.rs

@@ -1,9 +1,11 @@
+use std::alloc::{GlobalAlloc, System};
 use std::borrow::Cow;
 use std::collections::HashMap;
 use std::io::Write;
 use std::ops::ControlFlow;
 use std::sync::RwLock;
 
+use stats_alloc::StatsAlloc;
 use tracing::span::{Attributes, Id as TracingId};
 use tracing::{Metadata, Subscriber};
 use tracing_subscriber::layer::Context;
@@ -15,21 +17,37 @@ use crate::entry::{
 use crate::{Error, Trace};
 
 /// Layer that measures the time spent in spans.
-pub struct TraceLayer {
+pub struct TraceLayer<A: GlobalAlloc + 'static = System> {
     sender: std::sync::mpsc::Sender<Entry>,
     callsites: RwLock<HashMap<OpaqueIdentifier, ResourceId>>,
     start_time: std::time::Instant,
-    // TODO: kero add handle to allocator stats here
+    memory_allocator: Option<&'static StatsAlloc<A>>,
 }
 
 impl<W: Write> Trace<W> {
-    pub fn new(writer: W) -> (Self, TraceLayer) {
+    pub fn new(writer: W) -> (Self, TraceLayer<System>) {
         let (sender, receiver) = std::sync::mpsc::channel();
         let trace = Trace { writer, receiver };
         let layer = TraceLayer {
             sender,
             callsites: Default::default(),
             start_time: std::time::Instant::now(),
+            memory_allocator: None,
+        };
+        (trace, layer)
+    }
+
+    pub fn with_stats_alloc<A: GlobalAlloc>(
+        writer: W,
+        stats_alloc: &'static StatsAlloc<A>,
+    ) -> (Self, TraceLayer<A>) {
+        let (sender, receiver) = std::sync::mpsc::channel();
+        let trace = Trace { writer, receiver };
+        let layer = TraceLayer {
+            sender,
+            callsites: Default::default(),
+            start_time: std::time::Instant::now(),
+            memory_allocator: Some(stats_alloc),
         };
         (trace, layer)
     }
@@ -137,13 +155,19 @@ where
     }
 
     fn on_enter(&self, id: &TracingId, _ctx: Context<'_, S>) {
-        // TODO kero: add memory here
-        self.send(Entry::SpanEnter(SpanEnter { id: id.into(), time: self.elapsed() }))
+        self.send(Entry::SpanEnter(SpanEnter {
+            id: id.into(),
+            time: self.elapsed(),
+            memory: self.memory_allocator.map(|ma| ma.stats().into()),
+        }))
     }
 
     fn on_exit(&self, id: &TracingId, _ctx: Context<'_, S>) {
-        // TODO kero: add memory here
-        self.send(Entry::SpanExit(SpanExit { id: id.into(), time: self.elapsed() }))
+        self.send(Entry::SpanExit(SpanExit {
+            id: id.into(),
+            time: self.elapsed(),
+            memory: self.memory_allocator.map(|ma| ma.stats().into()),
+        }))
     }
 
     fn on_close(&self, id: TracingId, _ctx: Context<'_, S>) {

tracing-trace/src/processor/firefox_profiler.rs

@@ -5,10 +5,11 @@ use fxprof_processed_profile::{
     MarkerFieldFormat, MarkerLocation, MarkerSchema, MarkerSchemaField, Profile, ProfilerMarker,
     ReferenceTimestamp, SamplingInterval, StringHandle, Timestamp,
 };
+use once_cell::unsync::Lazy;
 use serde_json::json;
 
 use crate::entry::{
-    Entry, NewCallsite, NewSpan, ResourceId, SpanClose, SpanEnter, SpanExit, SpanId,
+    Entry, MemoryStats, NewCallsite, NewSpan, ResourceId, SpanClose, SpanEnter, SpanExit, SpanId,
 };
 use crate::{Error, TraceReader};
 
@@ -33,6 +34,16 @@ pub fn to_firefox_profile<R: std::io::Read>(
     let subcategory = profile.add_subcategory(category, "subcategory");
 
     // TODO kero: add counters profile.add_counters + last_memory_value
+    let mut current_memory = MemoryStats::default();
+    let mut allocations_counter = Lazy::new(|| {
+        profile.add_counter(main, "mimmalloc", "Memory", "Amount of allocation calls")
+    });
+    let mut deallocations_counter = Lazy::new(|| {
+        profile.add_counter(main, "mimmalloc", "Memory", "Amount of deallocation calls")
+    });
+    let mut reallocations_counter = Lazy::new(|| {
+        profile.add_counter(main, "mimmalloc", "Memory", "Amount of reallocation calls")
+    });
 
     for entry in trace {
         let entry = entry?;
@@ -56,7 +67,7 @@ pub fn to_firefox_profile<R: std::io::Read>(
             Entry::NewSpan(span) => {
                 spans.insert(span.id, (span, SpanStatus::Outside));
             }
-            Entry::SpanEnter(SpanEnter { id, time }) => {
+            Entry::SpanEnter(SpanEnter { id, time, memory }) => {
                 let (_span, status) = spans.get_mut(&id).unwrap();
 
                 let SpanStatus::Outside = status else {
@@ -67,16 +78,41 @@ pub fn to_firefox_profile<R: std::io::Read>(
 
                 last_timestamp = Timestamp::from_nanos_since_reference(time.as_nanos() as u64);
 
-                /* TODO kero: compute delta and update them
-                profile.add_counter_sample(
-                    counter,
-                    timestamp,
-                    value_delta,
-                    number_of_operations_delta,
-                )
-                */
+                if let Some(stats) = memory {
+                    let MemoryStats {
+                        allocations,
+                        deallocations,
+                        reallocations,
+                        bytes_allocated,
+                        bytes_deallocated,
+                        bytes_reallocated,
+                    } = current_memory - stats;
+
+                    profile.add_counter_sample(
+                        *allocations_counter,
+                        last_timestamp,
+                        bytes_allocated as f64,
+                        allocations.try_into().unwrap(),
+                    );
+
+                    profile.add_counter_sample(
+                        *deallocations_counter,
+                        last_timestamp,
+                        bytes_deallocated as f64,
+                        deallocations.try_into().unwrap(),
+                    );
+
+                    profile.add_counter_sample(
+                        *reallocations_counter,
+                        last_timestamp,
+                        bytes_reallocated as f64,
+                        reallocations.try_into().unwrap(),
+                    );
+
+                    current_memory = stats;
+                }
             }
-            Entry::SpanExit(SpanExit { id, time }) => {
+            Entry::SpanExit(SpanExit { id, time, memory }) => {
                 let (span, status) = spans.get_mut(&id).unwrap();
 
                 let SpanStatus::Inside(begin) = status else {
@@ -108,14 +144,39 @@ pub fn to_firefox_profile<R: std::io::Read>(
                     1,
                 );
 
-                /* TODO kero: compute delta and update them
-                profile.add_counter_sample(
-                    counter,
-                    timestamp,
-                    value_delta,
-                    number_of_operations_delta,
-                )
-                */
+                if let Some(stats) = memory {
+                    let MemoryStats {
+                        allocations,
+                        deallocations,
+                        reallocations,
+                        bytes_allocated,
+                        bytes_deallocated,
+                        bytes_reallocated,
+                    } = current_memory - stats;
+
+                    profile.add_counter_sample(
+                        *allocations_counter,
+                        last_timestamp,
+                        bytes_allocated as f64,
+                        allocations.try_into().unwrap(),
+                    );
+
+                    profile.add_counter_sample(
+                        *deallocations_counter,
+                        last_timestamp,
+                        bytes_deallocated as f64,
+                        deallocations.try_into().unwrap(),
+                    );
+
+                    profile.add_counter_sample(
+                        *reallocations_counter,
+                        last_timestamp,
+                        bytes_reallocated as f64,
+                        reallocations.try_into().unwrap(),
+                    );
+
+                    current_memory = stats;
+                }
 
                 let (callsite, _) = calls.get(&span.call_id).unwrap();
 

tracing-trace/src/processor/fmt.rs

@@ -2,7 +2,8 @@ use std::collections::HashMap;
 use std::io::Read;
 
 use crate::entry::{
-    Entry, NewCallsite, NewSpan, NewThread, ResourceId, SpanClose, SpanEnter, SpanExit, SpanId,
+    Entry, MemoryStats, NewCallsite, NewSpan, NewThread, ResourceId, SpanClose, SpanEnter,
+    SpanExit, SpanId,
 };
 use crate::{Error, TraceReader};
 
@@ -28,7 +29,7 @@ pub fn print_trace<R: Read>(trace: TraceReader<R>) -> Result<(), Error> {
             Entry::NewSpan(span) => {
                 spans.insert(span.id, (span, SpanStatus::Outside));
             }
-            Entry::SpanEnter(SpanEnter { id, time }) => {
+            Entry::SpanEnter(SpanEnter { id, time, memory }) => {
                 let (span, status) = spans.get_mut(&id).unwrap();
 
                 let SpanStatus::Outside = status else {
@@ -39,14 +40,23 @@ pub fn print_trace<R: Read>(trace: TraceReader<R>) -> Result<(), Error> {
 
                 let span = *span;
 
-                println!(
-                    "[{}]{}::{} <-",
-                    print_thread(&threads, span.thread_id),
-                    print_backtrace(&spans, &calls, &span),
-                    print_span(&calls, &span)
-                );
+                match memory {
+                    Some(stats) => println!(
+                        "[{}]{}::{} ({}) <-",
+                        print_thread(&threads, span.thread_id),
+                        print_backtrace(&spans, &calls, &span),
+                        print_span(&calls, &span),
+                        print_memory(stats),
+                    ),
+                    None => println!(
+                        "[{}]{}::{} <-",
+                        print_thread(&threads, span.thread_id),
+                        print_backtrace(&spans, &calls, &span),
+                        print_span(&calls, &span),
+                    ),
+                }
             }
-            Entry::SpanExit(SpanExit { id, time }) => {
+            Entry::SpanExit(SpanExit { id, time, memory }) => {
                 let (span, status) = spans.get_mut(&id).unwrap();
 
                 let SpanStatus::Inside(begin) = status else {
@@ -58,13 +68,23 @@ pub fn print_trace<R: Read>(trace: TraceReader<R>) -> Result<(), Error> {
 
                 let span = *span;
 
-                println!(
-                    "[{}]{}::{} -> {}",
-                    print_thread(&threads, span.thread_id),
-                    print_backtrace(&spans, &calls, &span),
-                    print_span(&calls, &span),
-                    print_duration(time - begin),
-                )
+                match memory {
+                    Some(stats) => println!(
+                        "[{}]{}::{} ({}) -> {}",
+                        print_thread(&threads, span.thread_id),
+                        print_backtrace(&spans, &calls, &span),
+                        print_span(&calls, &span),
+                        print_memory(stats),
+                        print_duration(time - begin),
+                    ),
+                    None => println!(
+                        "[{}]{}::{} -> {}",
+                        print_thread(&threads, span.thread_id),
+                        print_backtrace(&spans, &calls, &span),
+                        print_span(&calls, &span),
+                        print_duration(time - begin),
+                    ),
+                }
             }
             Entry::SpanClose(SpanClose { id, time: _ }) => {
                 spans.remove(&id);
@@ -126,3 +146,7 @@ fn print_duration(duration: std::time::Duration) -> String {
         format!("{}d", duration.as_secs_f64() / 3600.0 / 24.0)
     }
 }
+
+fn print_memory(memory: MemoryStats) -> String {
+    // Format only the total allocations in GiB, MiB, KiB, Bytes
+}