Logging the memory usage over time

Cargo.lock

@@ -3669,6 +3669,7 @@ dependencies = [
  "siphasher 1.0.0",
  "slice-group-by",
  "static-files",
+ "stats_alloc",
  "sysinfo",
  "tar",
  "temp-env",
@@ -5229,6 +5230,11 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
 
+[[package]]
+name = "stats_alloc"
+version = "0.1.10"
+source = "git+https://github.com/Kerollmops/stats_alloc?branch=stable-const-fn-trait#6f83c52160c7d0550fdf770e1f73d239b0ff9a97"
+
 [[package]]
 name = "strsim"
 version = "0.10.0"
@@ -5706,8 +5712,10 @@ version = "0.1.0"
 dependencies = [
  "color-spantrace",
  "fxprof-processed-profile",
+ "once_cell",
  "serde",
  "serde_json",
+ "stats_alloc",
  "tracing",
  "tracing-error",
  "tracing-subscriber",

meilisearch/Cargo.toml

@@ -108,6 +108,7 @@ url = { version = "2.5.0", features = ["serde"] }
 tracing = "0.1.40"
 tracing-subscriber = "0.3.18"
 tracing-trace = { version = "0.1.0", path = "../tracing-trace" }
+stats_alloc = { git = "https://github.com/Kerollmops/stats_alloc", branch = "stable-const-fn-trait", optional = true }
 
 [dev-dependencies]
 actix-rt = "2.9.0"
@@ -136,7 +137,7 @@ vergen = { version = "7.5.1", default-features = false, features = ["git"] }
 zip = { version = "0.6.6", optional = true }
 
 [features]
-default = ["analytics", "meilisearch-types/all-tokenizations", "mini-dashboard"]
+default = ["stats_alloc", "analytics", "meilisearch-types/all-tokenizations", "mini-dashboard"]
 analytics = ["segment"]
 mini-dashboard = [
     "actix-web-static-files",

meilisearch/src/main.rs

@@ -14,12 +14,18 @@ use is_terminal::IsTerminal;
 use meilisearch::analytics::Analytics;
 use meilisearch::{analytics, create_app, prototype_name, setup_meilisearch, Opt};
 use meilisearch_auth::{generate_master_key, AuthController, MASTER_KEY_MIN_SIZE};
+use mimalloc::MiMalloc;
 use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
 use tracing_subscriber::layer::SubscriberExt as _;
 use tracing_subscriber::Layer;
 
+#[cfg(not(feature = "stats_alloc"))]
 #[global_allocator]
-static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;
+static ALLOC: MiMalloc = MiMalloc;
+
+#[cfg(feature = "stats_alloc")]
+#[global_allocator]
+static ALLOC: stats_alloc::StatsAlloc<MiMalloc> = stats_alloc::StatsAlloc::new(MiMalloc);
 
 /// does all the setup before meilisearch is launched
 fn setup(opt: &Opt) -> anyhow::Result<()> {

tracing-trace/Cargo.toml

@@ -13,3 +13,5 @@ serde_json = "1.0.111"
 tracing = "0.1.40"
 tracing-error = "0.2.0"
 tracing-subscriber = "0.3.18"
+stats_alloc = { git = "https://github.com/Kerollmops/stats_alloc", branch = "stable-const-fn-trait" }
+once_cell = "1.19.0"

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 {
-        self.send(Entry::SpanEnter(SpanEnter { id: id.into(), time: self.elapsed() }))
+            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 {
-        self.send(Entry::SpanExit(SpanExit { id: id.into(), time: self.elapsed() }))
+            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
+                if let Some(stats) = memory {
+                    let MemoryStats {
+                        allocations,
+                        deallocations,
+                        reallocations,
+                        bytes_allocated,
+                        bytes_deallocated,
+                        bytes_reallocated,
+                    } = current_memory - stats;
+
                     profile.add_counter_sample(
-                    counter,
+                        *allocations_counter,
-                    timestamp,
+                        last_timestamp,
-                    value_delta,
+                        bytes_allocated as f64,
-                    number_of_operations_delta,
+                        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
+                if let Some(stats) = memory {
+                    let MemoryStats {
+                        allocations,
+                        deallocations,
+                        reallocations,
+                        bytes_allocated,
+                        bytes_deallocated,
+                        bytes_reallocated,
+                    } = current_memory - stats;
+
                     profile.add_counter_sample(
-                    counter,
+                        *allocations_counter,
-                    timestamp,
+                        last_timestamp,
-                    value_delta,
+                        bytes_allocated as f64,
-                    number_of_operations_delta,
+                        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!(
+                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)
+                        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!(
+                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
+}