Heap profiler with leak tracking

quickwit/Cargo.toml

@@ -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/quickwit-cli/Cargo.toml

@@ -81,6 +81,10 @@ quickwit-storage = { workspace = true, features = ["testsuite"] }
 
 [features]
 jemalloc = ["dep:tikv-jemalloc-ctl", "dep:tikv-jemallocator"]
+jemalloc-profiled = [
+  "quickwit-common/jemalloc-profiled",
+  "quickwit-serve/jemalloc-profiled"
+]
 ci-test = []
 pprof = ["quickwit-serve/pprof"]
 openssl-support = ["openssl-probe"]
@@ -127,6 +131,11 @@ release-macos-feature-vendored-set = [
   "quickwit-metastore/postgres",
   "quickwit-doc-mapper/multilang",
 ]
+release-heap-profiled = [
+  "release-feature-set",
+  "jemalloc-profiled",
+  "tokio/tracing",
+]
 
 [package.metadata.cargo-machete]
 # used to enable the `multilang` feature

quickwit/quickwit-cli/src/jemalloc.rs

@@ -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/logger.rs

@@ -27,6 +27,7 @@ use time::format_description::BorrowedFormatItem;
 use tracing::{Event, Level, Subscriber};
 use tracing_subscriber::EnvFilter;
 use tracing_subscriber::field::RecordFields;
+use tracing_subscriber::filter::filter_fn;
 use tracing_subscriber::fmt::FmtContext;
 use tracing_subscriber::fmt::format::{
     DefaultFields, Format, FormatEvent, FormatFields, Full, Json, JsonFields, Writer,
@@ -40,6 +41,14 @@ use crate::QW_ENABLE_OPENTELEMETRY_OTLP_EXPORTER_ENV_KEY;
 #[cfg(feature = "tokio-console")]
 use crate::QW_ENABLE_TOKIO_CONSOLE_ENV_KEY;
 
+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 +61,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 +96,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,17 +109,44 @@ pub fn setup_logging_and_tracing(
     } else {
         let event_format = EventFormat::get_from_env();
         let fmt_fields = event_format.format_fields();
-
-        registry
+        #[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),
+        );
+        // the heap profiler disables the env filter reloading because it seemed
+        // to overwrite the profiling filter
+        #[cfg(feature = "jemalloc-profiled")]
+        let registry = registry
+            .with(
+                tracing_subscriber::fmt::layer()
+                    .event_format(EventFormat::get_from_env())
+                    .fmt_fields(EventFormat::get_from_env().format_fields())
+                    .with_ansi(ansi_colors)
+                    .with_filter(filter_fn(|metadata| {
+                        metadata.is_span()
+                            || (metadata.is_event()
+                                && metadata.level() == &Level::TRACE
+                                && metadata
+                                    .target()
+                                    .starts_with("quickwit_common::jemalloc_profiled"))
+                    })),
+            )
             .with(
                 tracing_subscriber::fmt::layer()
                     .event_format(event_format)
                     .fmt_fields(fmt_fields)
-                    .with_ansi(ansi_colors),
-            )
+                    .with_ansi(ansi_colors)
+                    .with_filter(startup_env_filter(level)?),
+            );
+
+        registry
             .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)?;

quickwit/quickwit-common/Cargo.toml

@@ -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,15 @@ tracing = { workspace = true }
 [features]
 testsuite = []
 named_tasks = ["tokio/tracing"]
+jemalloc-profiled = [
+    "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"] }