diff --git a/Cargo.lock b/Cargo.lock index 67920b4bc..1d149733c 100644 --- a/Cargo.lock +++ b/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", diff --git a/meilisearch/Cargo.toml b/meilisearch/Cargo.toml index 98588b6c7..39e120bfb 100644 --- a/meilisearch/Cargo.toml +++ b/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", diff --git a/meilisearch/src/main.rs b/meilisearch/src/main.rs index 5a329eee2..fd3879b36 100644 --- a/meilisearch/src/main.rs +++ b/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 = stats_alloc::StatsAlloc::new(MiMalloc); /// does all the setup before meilisearch is launched fn setup(opt: &Opt) -> anyhow::Result<()> { diff --git a/tracing-trace/Cargo.toml b/tracing-trace/Cargo.toml index 9215fdfd0..8cd863368 100644 --- a/tracing-trace/Cargo.toml +++ b/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" diff --git a/tracing-trace/src/entry.rs b/tracing-trace/src/entry.rs index dd91a2a61..06ec05258 100644 --- a/tracing-trace/src/entry.rs +++ b/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, } #[derive(Clone, Copy, Debug, Serialize, Deserialize)] pub struct SpanExit { pub id: SpanId, pub time: std::time::Duration, + pub memory: Option, } #[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 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); diff --git a/tracing-trace/src/layer.rs b/tracing-trace/src/layer.rs index cbc5cf6b2..f5464b699 100644 --- a/tracing-trace/src/layer.rs +++ b/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 { sender: std::sync::mpsc::Sender, callsites: RwLock>, start_time: std::time::Instant, - // TODO: kero add handle to allocator stats here + memory_allocator: Option<&'static StatsAlloc>, } impl Trace { - pub fn new(writer: W) -> (Self, TraceLayer) { + pub fn new(writer: W) -> (Self, TraceLayer) { 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( + writer: W, + stats_alloc: &'static StatsAlloc, + ) -> (Self, TraceLayer) { + 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>) { diff --git a/tracing-trace/src/processor/firefox_profiler.rs b/tracing-trace/src/processor/firefox_profiler.rs index d3ac495d8..ce04d8bcf 100644 --- a/tracing-trace/src/processor/firefox_profiler.rs +++ b/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( 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( 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( 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( 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(); diff --git a/tracing-trace/src/processor/fmt.rs b/tracing-trace/src/processor/fmt.rs index a9356ba26..8e99752cb 100644 --- a/tracing-trace/src/processor/fmt.rs +++ b/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(trace: TraceReader) -> 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(trace: TraceReader) -> 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(trace: TraceReader) -> 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 +}