Enable Memory & CPU Profiling for df_engine

rust/otap-dataflow/Cargo.toml

@@ -41,6 +41,9 @@ serde_json.workspace = true
 clap.workspace = true
 mimalloc = { workspace = true, optional = true }
 sysinfo.workspace = true
+dhat = { workspace = true, optional = true }
+ctrlc = { workspace = true, optional = true }
+
 
 [target.'cfg(not(windows))'.dependencies]
 tikv-jemallocator = { workspace = true, optional = true }
@@ -86,8 +89,10 @@ ciborium = "0.2.2"
 clap = { version = "4.5.42", features = ["derive"] }
 core_affinity = "0.8.3"
 criterion = "0.8.0"
+ctrlc = "3"
 data-encoding = "2.9.0"
 datafusion = { version = "51.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"
@@ -217,6 +222,12 @@ aws = ["otap-df-otap/aws"]
 unsafe-optimizations = ["unchecked-index", "unchecked-arithmetic"]
 unchecked-index = []
 unchecked-arithmetic = []
+dhat-heap = [
+    "dep:dhat",
+    "dep:ctrlc",
+]
+
+
 
 # Dev/test tooling (fake data generator).
 dev-tools = ["otap-df-core-nodes/dev-tools"]
@@ -300,6 +311,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,50 @@ cargo run --example <example_name>
 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**:
+
+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**:
+cargo build --profile profiling --features dhat-heap --workspace
+
+**Build for only CPU profiling**:
+cargo build --profile profiling --workspace
+
+**Run**:
+
+**Run with both Memory & CPU profiling enabled**:
+samply record ./target/profiling/df_engine.exe --config .\configs\otap-noop.yaml
+
+**Run with only Memory profiling enabled**:
+./target/profiling/df_engine.exe --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,14 +23,13 @@ use otap_df_core_nodes as _;
 use otap_df_otap::OTAP_PIPELINE_FACTORY;
 use std::path::PathBuf;
 use sysinfo::System;
+use ctrlc;
+#[cfg(feature = "dhat-heap")]
+use {
+    std::sync::{LazyLock, Mutex},
+    dhat::Profiler,
+};
 
-#[cfg(all(
-    not(windows),
-    feature = "jemalloc",
-    feature = "mimalloc",
-    not(any(test, doc)),
-    not(clippy)
-))]
 compile_error!(
     "Features `jemalloc` and `mimalloc` are mutually exclusive. \
      To build with mimalloc, use: cargo build --release --no-default-features --features mimalloc"
@@ -78,11 +77,37 @@ use mimalloc::MiMalloc;
 #[cfg(all(not(windows), feature = "jemalloc", not(feature = "mimalloc")))]
 use tikv_jemallocator::Jemalloc;
 
-#[cfg(feature = "mimalloc")]
+// =======================
+// dhat (profiling) — wins everywhere when enabled
+// =======================
+#[cfg(all(
+    feature = "dhat-heap",
+    not(any(test, doc, clippy))
+))]
+#[global_allocator]
+static GLOBAL: dhat::Alloc = dhat::Alloc;
+
+// =======================
+// Windows default: mimalloc (only if dhat-heap is OFF)
+// =======================
+#[cfg(all(
+    feature = "mimalloc",
+    not(feature = "dhat-heap"),
+    not(any(test, doc, clippy))
+))]
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;
 
-#[cfg(all(not(windows), feature = "jemalloc", not(feature = "mimalloc")))]
+// =======================
+// Linux default: jemalloc (only if dhat-heap is OFF)
+// =======================
+#[cfg(all(
+    feature = "jemalloc",
+    not(windows),
+    not(feature = "dhat-heap"),
+    not(feature = "mimalloc"),
+    not(any(test, doc, clippy))
+))]
 #[global_allocator]
 static GLOBAL: Jemalloc = Jemalloc;
 
@@ -290,7 +315,25 @@ fn validate_engine_components(
     Ok(())
 }
 
+#[cfg(feature = "dhat-heap")]
+#[global_allocator]
+static ALLOC: dhat::Alloc = dhat::Alloc; // 1. The actual allocator
+#[cfg(feature = "dhat-heap")]
+static DHAT_PROFILER: LazyLock<Mutex<Option<Profiler>>> = LazyLock::new(|| Mutex::new(None));
 fn main() -> Result<(), Box<dyn std::error::Error>> {
+    #[cfg(feature = "dhat-heap")]
+    ctrlc::set_handler(|| {
+        if let Ok(mut guard) = DHAT_PROFILER.lock() {
+            // This triggers the drop and the file write
+            guard.take(); 
+        }
+        std::process::exit(130);
+    })?;
+    #[cfg(feature = "dhat-heap")]
+    {
+        let mut profiler = DHAT_PROFILER.lock().unwrap();
+        *profiler = Some(dhat::Profiler::new_heap());
+    }
     // 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()