Enable heap profiling in pyroscope (#8371)

* Start heap profiling if MALLOC_CONF set * Update doc
2026-07-23 11:11:00 -05:00 · 2026-03-13 11:45:30 +01:00
parent a1525d0ee6
commit c0e45c7987
4 changed files with 177 additions and 29 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3897,6 +3897,23 @@ version = "1.0.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2"

+[[package]]
+name = "jemalloc_pprof"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a0d44c349cfe2654897fadcb9de4f0bfbf48288ec344f700b2bd59f152dd209"
+dependencies = [
+ "anyhow",
+ "libc",
+ "mappings",
+ "once_cell",
+ "pprof_util",
+ "tempfile",
+ "tikv-jemalloc-ctl",
+ "tokio",
+ "tracing",
+]
+
 [[package]]
 name = "jieba-macros"
 version = "0.8.1"
@@ -4255,6 +4272,19 @@ dependencies = [
 "syn 2.0.117",
 ]

+[[package]]
+name = "mappings"
+version = "0.7.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8bab1e61a4b76757edb59cd81fcaa7f3ba9018d43b527d9abfad877b4c6c60f2"
+dependencies = [
+ "anyhow",
+ "libc",
+ "once_cell",
+ "pprof_util",
+ "tracing",
+]
+
 [[package]]
 name = "matchers"
 version = "0.2.0"
@@ -4554,6 +4584,20 @@ dependencies = [
 "windows-sys 0.59.0",
 ]

+[[package]]
+name = "num"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23"
+dependencies = [
+ "num-bigint",
+ "num-complex",
+ "num-integer",
+ "num-iter",
+ "num-rational",
+ "num-traits",
+]
+
 [[package]]
 name = "num-bigint"
 version = "0.4.6"
@@ -4642,6 +4686,17 @@ dependencies = [
 "num-traits",
 ]

+[[package]]
+name = "num-rational"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824"
+dependencies = [
+ "num-bigint",
+ "num-integer",
+ "num-traits",
+]
+
 [[package]]
 name = "num-traits"
 version = "0.2.19"
@@ -5255,6 +5310,20 @@ dependencies = [
 "thiserror 2.0.18",
 ]

+[[package]]
+name = "pprof_util"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eea0cc524de808a6d98d192a3d99fe95617031ad4a52ec0a0f987ef4432e8fe1"
+dependencies = [
+ "anyhow",
+ "backtrace",
+ "flate2",
+ "num",
+ "paste",
+ "prost 0.14.3",
+]
+
 [[package]]
 name = "ppv-lite86"
 version = "0.2.21"
@@ -5717,6 +5786,7 @@ version = "2.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "85eebd4bcbf45db75f67d2ba20ea0207bd111d2029c07a7db3229289173d4387"
 dependencies = [
+ "jemalloc_pprof",
 "lazy_static",
 "libc",
 "libflate",
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -151,7 +151,7 @@ actix-web-extras = "0.1.0"

 [target.'cfg(target_os = "linux")'.dependencies]
 procfs = { version = "0.18.0", default-features = false }
-pyroscope = { version = "2.0.0", features = ["backend-pprof-rs"] }
+pyroscope = { version = "2.0.0", features = ["backend-pprof-rs", "backend-jemalloc"] }
 # Backtrace
 rstack-self = { version = "0.3.0", optional = true }

@@ -161,6 +161,7 @@ tikv-jemallocator = { version = "0.6.1", features = [
    "stats",
    "unprefixed_malloc_on_supported_platforms",
    "background_threads",
+    "profiling",
 ] }
 tikv-jemalloc-ctl = { version = "0.6.1", features = ["stats"] }

--- a/docs/DEVELOPMENT.md
+++ b/docs/DEVELOPMENT.md
@@ -309,6 +309,14 @@ curl -X PATCH http://localhost:6333/debugger \

 View flame graphs at `http://localhost:4040` — select application `qdrant`, filter by `identifier = qdrant-local` (`process_cpu:cpu:nanoseconds:cpu:nanoseconds{service_name="qdrant", identifier="qdrant-local"}`).

+### Pyroscope (continuous Heap profiling)
+
+If Jemalloc profiler is enabled, you can also use Pyroscope for continuous heap profiling.
+To enable the profiler start Qdrant with `MALLOC_CONF="prof:true,prof_active:true"` environment variable.
+Then when Pyroscope is activated via `POST /debugger` it will also start collecting heap stats and send them to the server.
+You can view them in the Pyroscope UI by selecting `memory:inuse_space:bytes:space:bytes{service_name="qdrant", identifier="your-identifier"}`.
+
+
 ## API changes

 ### REST
--- a/src/common/pyroscope_state.rs
+++ b/src/common/pyroscope_state.rs
@@ -1,6 +1,7 @@
 #[cfg(target_os = "linux")]
 pub mod pyro {

+    use pyroscope::backend::jemalloc::jemalloc_backend;
    use pyroscope::backend::{BackendConfig, PprofConfig, pprof_backend};
    use pyroscope::pyroscope::{PyroscopeAgentBuilder, PyroscopeAgentRunning};
    use pyroscope::{PyroscopeAgent, PyroscopeError};
@@ -9,11 +10,12 @@ pub mod pyro {

    pub struct PyroscopeState {
        pub config: PyroscopeConfig,
-        pub agent: Option<PyroscopeAgent<PyroscopeAgentRunning>>,
+        pub cpu_agent: Option<PyroscopeAgent<PyroscopeAgentRunning>>,
+        pub heap_agent: Option<PyroscopeAgent<PyroscopeAgentRunning>>,
    }

    impl PyroscopeState {
-        fn build_agent(
+        fn build_cpu_agent(
            config: &PyroscopeConfig,
        ) -> Result<PyroscopeAgent<PyroscopeAgentRunning>, PyroscopeError> {
            // pprof uses tempfile::NamedTempFile which respects TMPDIR.
@@ -33,7 +35,7 @@ pub mod pyro {
            let backend_impl = pprof_backend(pprof_config, BackendConfig::default());

            log::info!(
-                "Starting pyroscope agent with identifier {}",
+                "Starting pyroscope CPU agent with identifier {}",
                &config.identifier
            );
            // TODO: Add more tags like peerId and peerUrl
@@ -52,42 +54,109 @@ pub mod pyro {
            Ok(running_agent)
        }

+        fn build_heap_agent(
+            config: &PyroscopeConfig,
+        ) -> Result<PyroscopeAgent<PyroscopeAgentRunning>, PyroscopeError> {
+            let backend_impl = jemalloc_backend();
+
+            log::info!(
+                "Starting pyroscope heap agent with identifier {}",
+                &config.identifier
+            );
+            let agent = PyroscopeAgentBuilder::new(
+                config.url.clone(),
+                "qdrant",
+                0,
+                "pyroscope-rs",
+                env!("CARGO_PKG_VERSION"),
+                backend_impl,
+            )
+            .tags(vec![("app", "Qdrant"), ("identifier", &config.identifier)])
+            .build()?;
+            let running_agent = agent.start()?;
+
+            Ok(running_agent)
+        }
+
+        fn is_jemalloc_profiling_enabled() -> bool {
+            std::env::var("MALLOC_CONF")
+                .map(|conf| conf.split(',').any(|opt| opt.trim() == "prof:true"))
+                .unwrap_or(false)
+        }
+
        pub fn from_config(config: Option<PyroscopeConfig>) -> Option<Self> {
            match config {
                Some(pyro_config) => {
-                    let agent = PyroscopeState::build_agent(&pyro_config);
-                    match agent {
-                        Ok(agent) => Some(PyroscopeState {
-                            config: pyro_config,
-                            agent: Some(agent),
-                        }),
+                    let cpu_agent = match PyroscopeState::build_cpu_agent(&pyro_config) {
+                        Ok(agent) => Some(agent),
                        Err(err) => {
-                            log::warn!("Pyroscope agent failed to start {err}");
-                            None
+                            log::warn!("Pyroscope CPU agent failed to start: {err}");
+                            return None;
                        }
-                    }
+                    };
+
+                    let heap_agent = if Self::is_jemalloc_profiling_enabled() {
+                        match PyroscopeState::build_heap_agent(&pyro_config) {
+                            Ok(agent) => {
+                                log::info!(
+                                    "Pyroscope heap agent started (MALLOC_CONF has prof:true)"
+                                );
+                                Some(agent)
+                            }
+                            Err(err) => {
+                                log::warn!("Pyroscope heap agent failed to start: {err}");
+                                None
+                            }
+                        }
+                    } else {
+                        log::info!(
+                            "Jemalloc profiling not enabled, skipping heap agent. Set MALLOC_CONF=\"prof:true,prof_active:true\" to enable."
+                        );
+                        None
+                    };
+
+                    Some(PyroscopeState {
+                        config: pyro_config,
+                        cpu_agent,
+                        heap_agent,
+                    })
                }
                None => None,
            }
        }

-        pub fn stop_agent(&mut self) -> bool {
-            log::info!("Stopping pyroscope agent");
-            if let Some(agent) = self.agent.take() {
-                return match agent.stop() {
-                    Ok(stopped_agent) => {
-                        log::info!("Stopped pyroscope agent. Shutting it down");
-                        stopped_agent.shutdown();
-                        log::info!("Pyroscope agent shut down completed.");
-                        true
-                    }
-                    Err(err) => {
-                        log::warn!("Pyroscope agent failed to stop {err}");
-                        false
-                    }
-                };
+        fn stop_single_agent(agent: PyroscopeAgent<PyroscopeAgentRunning>, name: &str) -> bool {
+            match agent.stop() {
+                Ok(stopped_agent) => {
+                    log::info!("Stopped pyroscope {name} agent. Shutting it down");
+                    stopped_agent.shutdown();
+                    log::info!("Pyroscope {name} agent shut down completed.");
+                    true
+                }
+                Err(err) => {
+                    log::warn!("Pyroscope {name} agent failed to stop: {err}");
+                    false
+                }
            }
-            true
+        }
+
+        pub fn stop_agent(&mut self) -> bool {
+            log::info!("Stopping pyroscope agents");
+            let mut success = true;
+
+            if let Some(agent) = self.cpu_agent.take()
+                && !Self::stop_single_agent(agent, "CPU")
+            {
+                success = false;
+            }
+
+            if let Some(agent) = self.heap_agent.take()
+                && !Self::stop_single_agent(agent, "heap")
+            {
+                success = false;
+            }
+
+            success
        }
    }