Enable Memory & CPU Profiling for df_engine

rust/otap-dataflow/Cargo.toml

@@ -41,6 +41,8 @@ serde_json.workspace = true
 clap.workspace = true
 mimalloc = { workspace = true, optional = true }
 sysinfo.workspace = true
+dhat = { workspace = true, optional = true }
+cfg-if.workspace = true
 
 [target.'cfg(not(windows))'.dependencies]
 tikv-jemallocator = { workspace = true, optional = true }
@@ -80,6 +82,7 @@ bitflags = "2.10"
 bytes = "1.11.1"
 bytemuck = "1.24"
 blake3 = "1.5.5"
+cfg-if = "1.0.4"
 chrono = { version = "0.4", features = ["serde"] }
 chrono-tz = { version = "0.10" }
 ciborium = "0.2.2"
@@ -88,6 +91,7 @@ core_affinity = "0.8.3"
 criterion = "0.8.0"
 data-encoding = "2.9.0"
 datafusion = { version = "53.0.0", default-features = false }
+dhat = "0.3.3"
 fluke-hpack = "0.3.1"
 flume = { version = "0.12.0", default-features = false, features = ["async"] }
 futures = "0.3.32"
@@ -239,6 +243,11 @@ aws = ["otap-df-otap/aws"]
 unsafe-optimizations = ["unchecked-index", "unchecked-arithmetic"]
 unchecked-index = []
 unchecked-arithmetic = []
+dhat-heap = [
+    "dep:dhat",
+]
+
+
 
 # Dev/test tooling (fake data generator).
 dev-tools = ["otap-df-core-nodes/dev-tools"]
@@ -322,6 +331,10 @@ missing_crate_level_docs = "deny"
 [profile.release]
 debug = "line-tables-only"  # minimum required for profiling
 
+[profile.profiling]
+inherits = "release"
+debug = 2
+
 # A more in-depth analysis is necessary to determine the optimal parameters for the release profile.
 #[profile.release]
 #lto = "thin"

rust/otap-dataflow/README.md

@@ -384,6 +384,74 @@
 docker build --build-context otel-arrow=../../ -f Dockerfile -t df_engine .
 ```
 
+## Profiling
+
+This section covers Memory & CPU Profiling for df_engine using dhat-rs and
+samply profilers respectively.
+
+**Requirements**:
+
+- dhat-rs <https://docs.rs/dhat/latest/dhat/>
+- samply <https://github.com/mstange/samply>
+
+**Installation**:
+
+```cmd/pwsh/bash
+cargo install --locked samply
+```
+
+Note: dhat-rs is a library crate designed for heap profiling in Rust programs.
+      Hence it does not require explicit installation.
+
+**Build**:
+
+**Build for both CPU & Memory profiling**:
+
+```cmd/pwsh/bash
+cargo build --profile profiling --no-default-features --features dhat-heap --workspace
+```
+
+**Build for only CPU profiling**:
+
+```cmd/pwsh/bash
+cargo build --profile profiling --workspace
+```
+
+**Run**:
+
+**Run with both Memory & CPU profiling enabled**:
+
+```cmd/pwsh
+samply record .\target\profiling\df_engine.exe --config .\configs\otap-noop.yaml
+```
+or
+
+```bash
+samply record ./target/profiling/df_engine --config ./configs/otap-noop.yaml
+```
+
+**Run with only Memory profiling enabled**:
+
+```cmd/pwsh
+.\target\profiling\df_engine.exe --config .\configs\otap-noop.yaml
+```
+or
+
+```bash
+./target/profiling/df_engine --config ./configs/otap-noop.yaml
+```
+
+**Result**:
+
+On successful termination of df_engine.exe, it will generate dhat-heap.json file
+for Memory profiling that need to be rendered by uploading it to:
+<https://nnethercote.github.io/dh_view/>.
+
+CPU profiling output will be automatically renderd on browser.
+
+Note: dhat needs a clean shutdown to generate dhat-heap.json file. In df_engine
+      this can be done manually with Ctrl-C.
+
 ## Contributing
 
 - [Contribution Guidelines](CONTRIBUTING.md)

rust/otap-dataflow/src/main.rs

@@ -23,19 +23,73 @@ use otap_df_core_nodes as _;
 use otap_df_otap::OTAP_PIPELINE_FACTORY;
 use std::path::PathBuf;
 use sysinfo::System;
+use cfg_if::cfg_if;
 
+// -----------------------------------------------------------------------------
+// Feature guard: jemalloc + mimalloc + dhat-heap any two together should fail (non-Windows only)
+// -----------------------------------------------------------------------------
 #[cfg(all(
-    not(windows),
-    feature = "jemalloc",
-    feature = "mimalloc",
     not(any(test, doc)),
-    not(clippy)
+    not(clippy),
+    any(
+        all(feature = "dhat-heap", feature = "mimalloc"),
+        all(feature = "dhat-heap", feature = "jemalloc"),
+        all(feature = "jemalloc", feature = "mimalloc"),
+    )
 ))]
 compile_error!(
-    "Features `jemalloc` and `mimalloc` are mutually exclusive. \
-     To build with mimalloc, use: cargo build --release --no-default-features --features mimalloc"
+    "Allocator features are mutually exclusive. Enable only one allocator: `dhat-heap`, `mimalloc`, `jemalloc`. \
+    Example: \
+        (mimalloc): cargo build --release --no-default-features --features mimalloc. \
+        (jemalloc): cargo build --release --no-default-features --features jemalloc. \
+        (dhat): cargo build --profile profiling --no-default-features --features dhat-heap."
 );
 
+#[cfg(feature = "dhat-heap")]
+use {
+    std::sync::{LazyLock, Mutex},
+    dhat::Profiler,
+};
+
+#[cfg(feature = "mimalloc")]
+use mimalloc::MiMalloc;
+
+#[cfg(all(not(windows), feature = "jemalloc"))]
+use tikv_jemallocator::Jemalloc;
+
+// -----------------------------------------------------------------------------
+// Global allocator selection.
+// -----------------------------------------------------------------------------
+cfg_if! {
+    // dhat (profiling) — wins everywhere when enabled
+    if #[cfg(feature = "dhat-heap")] {
+        #[global_allocator]
+        static GLOBAL: dhat::Alloc = dhat::Alloc;
+        static DHAT_PROFILER: LazyLock<Mutex<Option<Profiler>>> = LazyLock::new(|| Mutex::new(None));
+
+        fn dhat_start() {
+                let mut profiler = DHAT_PROFILER.lock().unwrap();
+                *profiler = Some(dhat::Profiler::new_heap());
+        }
+
+        fn dhat_finish() {
+        // Dropping the Profiler is what triggers dhat-heap.json generation. [1](https://outlook.office365.com/owa/?ItemID=AAMkADYwZTk3ZGM5LTU2ZjItNDM2NC1hZmQ0LWY5NjhmZmM5M2NiNwBGAAAAAAC7q57xFfmuSb3iOsuh%			2fVOABwAlpRXD5QvATJ55erMTPKHzAAAAH3TYAABzDoBtHPxlRIZv17%2bj%2fBx2AAhyk6imAAA%3d&exvsurl=1&viewmodel=ReadMessageItem)
+                let mut profiler = DHAT_PROFILER.lock().unwrap();
+                let _ = profiler.take();
+	}
+
+    // Windows default: mimalloc
+    } else if #[cfg(feature = "mimalloc")] {
+        #[global_allocator]
+        static GLOBAL: MiMalloc = MiMalloc;
+
+    // Linux default: jemalloc
+    } else if #[cfg(all(not(windows), feature = "jemalloc"))] {
+        #[global_allocator]
+        static GLOBAL: Jemalloc = Jemalloc;
+    }
+}
+
 // Crypto provider features are mutually exclusive.
 // The `not(any(test, doc))` and `not(clippy)` guards mirror the jemalloc/mimalloc
 // pattern so that `cargo test --all-features` (used in CI) does not fail.
@@ -72,19 +126,6 @@ compile_error!(
      Use --no-default-features to disable the default crypto provider, then enable exactly one."
 );
 
-#[cfg(feature = "mimalloc")]
-use mimalloc::MiMalloc;
-
-#[cfg(all(not(windows), feature = "jemalloc", not(feature = "mimalloc")))]
-use tikv_jemallocator::Jemalloc;
-
-#[cfg(feature = "mimalloc")]
-#[global_allocator]
-static GLOBAL: MiMalloc = MiMalloc;
-
-#[cfg(all(not(windows), feature = "jemalloc", not(feature = "mimalloc")))]
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
 
 #[derive(Parser, Debug)]
 #[command(
@@ -291,6 +332,10 @@ fn validate_engine_components(
 }
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
+    #[cfg(feature = "dhat-heap")]
+    {
+        dhat_start()
+    }
     // Install the rustls crypto provider selected by the crypto-* feature flag.
     // This must happen before any TLS connections (reqwest, tonic, etc.).
     otap_df_otap::crypto::install_crypto_provider()
@@ -318,6 +363,11 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let controller = Controller::new(&OTAP_PIPELINE_FACTORY);
     let result = controller.run_forever(engine_cfg);
+    #[cfg(feature = "dhat-heap")]
+    {
+        dhat_finish()
+    }
+
     match result {
         Ok(_) => {
             println!("Pipeline run successfully");