Heap profiler with leak tracking (#5763)

* Heap profiler with leak tracking

* Improve wording and types

* Better rationalization around nested allocations

* Pad the atomic flags to the cache line size

* Add padding

* Add tokio trace

* Move backtrace to trace event format

* Improve loc for spawned task when using tokio/tracing

* Move mutex from tracker

* Add docker ci

* Refactor logger for readability

* Fix typos

* Fix typo

Author: rdettai

.github/workflows/publish_docker_images.yml

@@ -62,6 +62,11 @@ jobs:
           echo "QW_COMMIT_DATE=$(TZ=UTC0 git log -1 --format=%cd --date=format-local:%Y-%m-%dT%H:%M:%SZ)" >> $GITHUB_ENV
           echo "QW_COMMIT_HASH=$(git rev-parse HEAD)" >> $GITHUB_ENV
           echo "QW_COMMIT_TAGS=$(git tag --points-at HEAD | tr '\n' ',')" >> $GITHUB_ENV
+          if [[ "${{ github.event_name }}" == "push" && "${{ github.ref_type }}" == "tag" && "${GITHUB_REF#refs/tags/}" == *"jemprof"* ]]; then
+            echo "CARGO_FEATURES=release-jemalloc-profiled" >> $GITHUB_ENV
+          else
+            echo "CARGO_FEATURES=release-feature-set" >> $GITHUB_ENV
+          fi
 
       - name: Build and push image
         uses: docker/build-push-action@v6
@@ -73,6 +78,7 @@ jobs:
             QW_COMMIT_DATE=${{ env.QW_COMMIT_DATE }}
             QW_COMMIT_HASH=${{ env.QW_COMMIT_HASH }}
             QW_COMMIT_TAGS=${{ env.QW_COMMIT_TAGS }}
+            CARGO_FEATURES=${{ env.CARGO_FEATURES }}
           labels: ${{ steps.meta.outputs.labels }}
           outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true
 

quickwit/Cargo.lock

@@ -85,6 +85,7 @@ assert-json-diff = "2"
 async-compression = { version = "0.4", features = ["tokio", "gzip"] }
 async-speed-limit = "0.4"
 async-trait = "0.1"
+backtrace = "0.3"
 base64 = "0.22"
 binggan = { version = "0.14" }
 bytes = { version = "1", features = ["serde"] }

quickwit/Cargo.toml

@@ -22,6 +22,7 @@ path = "src/generate_markdown.rs"
 
 [dependencies]
 anyhow = { workspace = true }
+backtrace = { workspace = true, optional = true }
 bytesize = { workspace = true }
 chrono = { workspace = true }
 clap = { workspace = true }
@@ -81,6 +82,11 @@ quickwit-storage = { workspace = true, features = ["testsuite"] }
 
 [features]
 jemalloc = ["dep:tikv-jemalloc-ctl", "dep:tikv-jemallocator"]
+jemalloc-profiled = [
+  "dep:backtrace",
+  "quickwit-common/jemalloc-profiled",
+  "quickwit-serve/jemalloc-profiled"
+]
 ci-test = []
 pprof = ["quickwit-serve/pprof"]
 openssl-support = ["openssl-probe"]
@@ -127,6 +133,10 @@ release-macos-feature-vendored-set = [
   "quickwit-metastore/postgres",
   "quickwit-doc-mapper/multilang",
 ]
+release-jemalloc-profiled = [
+  "release-feature-set",
+  "jemalloc-profiled",
+]
 
 [package.metadata.cargo-machete]
 # used to enable the `multilang` feature

quickwit/quickwit-cli/Cargo.toml

@@ -19,6 +19,10 @@ use tikv_jemallocator::Jemalloc;
 use tracing::error;
 
 #[global_allocator]
+#[cfg(feature = "jemalloc-profiled")]
+pub static GLOBAL: quickwit_common::jemalloc_profiled::JemallocProfiled =
+    quickwit_common::jemalloc_profiled::JemallocProfiled(Jemalloc);
+#[cfg(not(feature = "jemalloc-profiled"))]
 pub static GLOBAL: Jemalloc = Jemalloc;
 
 const JEMALLOC_METRICS_POLLING_INTERVAL: Duration = Duration::from_secs(1);

quickwit/quickwit-cli/src/jemalloc.rs

@@ -40,6 +40,16 @@ use crate::QW_ENABLE_OPENTELEMETRY_OTLP_EXPORTER_ENV_KEY;
 #[cfg(feature = "tokio-console")]
 use crate::QW_ENABLE_TOKIO_CONSOLE_ENV_KEY;
 
+/// Load the default logging filter from the environment. The filter can later
+/// be updated using the result callback of [setup_logging_and_tracing].
+fn startup_env_filter(level: Level) -> anyhow::Result<EnvFilter> {
+    let env_filter = env::var("RUST_LOG")
+        .map(|_| EnvFilter::from_default_env())
+        .or_else(|_| EnvFilter::try_new(format!("quickwit={level},tantivy=WARN")))
+        .context("failed to set up tracing env filter")?;
+    Ok(env_filter)
+}
+
 pub fn setup_logging_and_tracing(
     level: Level,
     ansi_colors: bool,
@@ -52,13 +62,10 @@ pub fn setup_logging_and_tracing(
             return Ok(quickwit_serve::do_nothing_env_filter_reload_fn());
         }
     }
-    let env_filter = env::var("RUST_LOG")
-        .map(|_| EnvFilter::from_default_env())
-        .or_else(|_| EnvFilter::try_new(format!("quickwit={level},tantivy=WARN")))
-        .context("failed to set up tracing env filter")?;
     global::set_text_map_propagator(TraceContextPropagator::new());
-    let (reloadable_env_filter, reload_handle) = tracing_subscriber::reload::Layer::new(env_filter);
-    let registry = tracing_subscriber::registry().with(reloadable_env_filter);
+    let (reloadable_env_filter, reload_handle) =
+        tracing_subscriber::reload::Layer::new(startup_env_filter(level)?);
+    let registry = tracing_subscriber::registry();
     // Note on disabling ANSI characters: setting the ansi boolean on event format is insufficient.
     // It is thus set on layers, see https://github.com/tokio-rs/tracing/issues/1817
     if get_bool_from_env(QW_ENABLE_OPENTELEMETRY_OTLP_EXPORTER_ENV_KEY, false) {
@@ -90,6 +97,7 @@ pub fn setup_logging_and_tracing(
         let fmt_fields = event_format.format_fields();
 
         registry
+            .with(reloadable_env_filter)
             .with(telemetry_layer)
             .with(
                 tracing_subscriber::fmt::layer()
@@ -102,24 +110,44 @@ pub fn setup_logging_and_tracing(
     } else {
         let event_format = EventFormat::get_from_env();
         let fmt_fields = event_format.format_fields();
+        #[cfg(not(feature = "jemalloc-profiled"))]
+        let registry = registry.with(reloadable_env_filter).with(
+            tracing_subscriber::fmt::layer()
+                .event_format(event_format)
+                .fmt_fields(fmt_fields)
+                .with_ansi(ansi_colors),
+        );
+        #[cfg(feature = "jemalloc-profiled")]
+        let registry = jemalloc_profiled::configure_registry(
+            registry,
+            event_format,
+            fmt_fields,
+            ansi_colors,
+            level,
+        )?;
 
         registry
-            .with(
-                tracing_subscriber::fmt::layer()
-                    .event_format(event_format)
-                    .fmt_fields(fmt_fields)
-                    .with_ansi(ansi_colors),
-            )
             .try_init()
             .context("failed to register tracing subscriber")?;
     }
+
     Ok(Arc::new(move |env_filter_def: &str| {
         let new_env_filter = EnvFilter::try_new(env_filter_def)?;
         reload_handle.reload(new_env_filter)?;
         Ok(())
     }))
 }
 
+/// We do not rely on the RFC3339 implementation, because it has a nanosecond precision.
+/// See discussion here: https://github.com/time-rs/time/discussions/418
+fn time_formatter() -> UtcTime<Vec<BorrowedFormatItem<'static>>> {
+    let time_format = time::format_description::parse(
+        "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z",
+    )
+    .expect("time format description should be valid");
+    UtcTime::new(time_format)
+}
+
 enum EventFormat<'a> {
     Full(Format<Full, UtcTime<Vec<BorrowedFormatItem<'a>>>>),
     Json(Format<Json>),
@@ -136,17 +164,9 @@ impl EventFormat<'_> {
             let json_format = tracing_subscriber::fmt::format().json();
             EventFormat::Json(json_format)
         } else {
-            // We do not rely on the RFC3339 implementation, because it has a nanosecond precision.
-            // See discussion here: https://github.com/time-rs/time/discussions/418
-            let timer_format = time::format_description::parse(
-                "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z",
-            )
-            .expect("time format description should be valid");
-            let timer = UtcTime::new(timer_format);
-
             let full_format = tracing_subscriber::fmt::format()
                 .with_target(true)
-                .with_timer(timer);
+                .with_timer(time_formatter());
 
             EventFormat::Full(full_format)
         }
@@ -191,3 +211,129 @@ impl FormatFields<'_> for FieldFormat {
         }
     }
 }
+
+/// Logger configurations specific to the jemalloc profiler.
+///
+/// A custom event formatter is used to print the backtrace of the
+/// profiling events.
+#[cfg(feature = "jemalloc-profiled")]
+pub(super) mod jemalloc_profiled {
+    use std::fmt;
+
+    use quickwit_common::jemalloc_profiled::JEMALLOC_PROFILER_TARGET;
+    use time::format_description::BorrowedFormatItem;
+    use tracing::{Event, Level, Metadata, Subscriber};
+    use tracing_subscriber::Layer;
+    use tracing_subscriber::filter::filter_fn;
+    use tracing_subscriber::fmt::format::{DefaultFields, Writer};
+    use tracing_subscriber::fmt::time::{FormatTime, UtcTime};
+    use tracing_subscriber::fmt::{FmtContext, FormatEvent, FormatFields, FormattedFields};
+    use tracing_subscriber::layer::SubscriberExt;
+    use tracing_subscriber::registry::LookupSpan;
+
+    use super::{EventFormat, FieldFormat, startup_env_filter, time_formatter};
+
+    /// An event formatter specific to the memory profiler output.
+    ///
+    /// Also displays a backtrace after spans and the fields of the tracing
+    /// event (into separate lines).
+    struct ProfilingFormat {
+        time_formatter: UtcTime<Vec<BorrowedFormatItem<'static>>>,
+    }
+
+    impl Default for ProfilingFormat {
+        fn default() -> Self {
+            Self {
+                time_formatter: time_formatter(),
+            }
+        }
+    }
+
+    impl<S, N> FormatEvent<S, N> for ProfilingFormat
+    where
+        S: Subscriber + for<'a> LookupSpan<'a>,
+        N: for<'a> FormatFields<'a> + 'static,
+    {
+        fn format_event(
+            &self,
+            ctx: &FmtContext<'_, S, N>,
+            mut writer: Writer<'_>,
+            event: &Event<'_>,
+        ) -> fmt::Result {
+            self.time_formatter.format_time(&mut writer)?;
+            write!(writer, " {JEMALLOC_PROFILER_TARGET} ")?;
+            if let Some(scope) = ctx.event_scope() {
+                let mut seen = false;
+
+                for span in scope.from_root() {
+                    write!(writer, "{}", span.metadata().name())?;
+                    seen = true;
+
+                    let ext = span.extensions();
+                    if let Some(fields) = &ext.get::<FormattedFields<N>>() {
+                        if !fields.is_empty() {
+                            write!(writer, "{{{}}}:", fields)?;
+                        }
+                    }
+                }
+
+                if seen {
+                    writer.write_char(' ')?;
+                }
+            };
+
+            ctx.format_fields(writer.by_ref(), event)?;
+            writeln!(writer)?;
+
+            // Print a backtrace to help idenify the callsite
+            backtrace::trace(|frame| {
+                backtrace::resolve_frame(frame, |symbol| {
+                    if let Some(symbole_name) = symbol.name() {
+                        let _ = writeln!(writer, "{}", symbole_name);
+                    } else {
+                        let _ = writeln!(writer, "symb failed");
+                    }
+                });
+                true
+            });
+            Ok(())
+        }
+    }
+
+    fn profiler_tracing_filter(metadata: &Metadata) -> bool {
+        metadata.is_span() || (metadata.is_event() && metadata.target() == JEMALLOC_PROFILER_TARGET)
+    }
+
+    /// Configures the regular logging layer and a specific layer that gathers
+    /// extra debug information for the jemalloc profiler.
+    ///
+    /// The the jemalloc profiler formatter disables the env filter reloading
+    /// because the [tracing_subscriber::reload::Layer] seems to overwrite the
+    /// TRACE level span filter even though it's applied to a separate layer.
+    pub(super) fn configure_registry<S>(
+        registry: S,
+        event_format: EventFormat<'static>,
+        fmt_fields: FieldFormat,
+        ansi_colors: bool,
+        level: Level,
+    ) -> anyhow::Result<impl Subscriber + for<'span> LookupSpan<'span>>
+    where
+        S: Subscriber + for<'span> LookupSpan<'span>,
+    {
+        Ok(registry
+            .with(
+                tracing_subscriber::fmt::layer()
+                    .event_format(ProfilingFormat::default())
+                    .fmt_fields(DefaultFields::new())
+                    .with_ansi(ansi_colors)
+                    .with_filter(filter_fn(profiler_tracing_filter)),
+            )
+            .with(
+                tracing_subscriber::fmt::layer()
+                    .event_format(event_format)
+                    .fmt_fields(fmt_fields)
+                    .with_ansi(ansi_colors)
+                    .with_filter(startup_env_filter(level)?),
+            ))
+    }
+}

quickwit/quickwit-cli/src/logger.rs

@@ -14,6 +14,7 @@ license.workspace = true
 anyhow = { workspace = true }
 async-speed-limit = { workspace = true }
 async-trait = { workspace = true }
+backtrace = { workspace = true, optional = true }
 bytesize = { workspace = true }
 coarsetime = { workspace = true }
 dyn-clone = { workspace = true }
@@ -37,6 +38,8 @@ siphasher = { workspace = true }
 sysinfo = { workspace = true }
 tempfile = { workspace = true }
 thiserror = { workspace = true }
+tikv-jemallocator = { workspace = true, optional = true }
+tikv-jemalloc-ctl = { workspace = true, optional = true }
 tokio = { workspace = true }
 tokio-metrics = { workspace = true }
 tokio-stream = { workspace = true }
@@ -47,9 +50,16 @@ tracing = { workspace = true }
 [features]
 testsuite = []
 named_tasks = ["tokio/tracing"]
+jemalloc-profiled = [
+    "named_tasks",
+    "dep:backtrace",
+    "dep:tikv-jemallocator",
+    "dep:tikv-jemalloc-ctl"
+]
 
 [dev-dependencies]
 serde_json = { workspace = true }
 tempfile = { workspace = true }
 proptest = { workspace = true }
+serial_test = { workspace = true }
 tokio = { workspace = true, features = ["test-util"] }