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 = "0.3.3"
+ctrlc = "3"
+once_cell = "1"
 
 [target.'cfg(not(windows))'.dependencies]
 tikv-jemallocator = { workspace = true, optional = true }
@@ -217,6 +220,9 @@ aws = ["otap-df-otap/aws"]
 unsafe-optimizations = ["unchecked-index", "unchecked-arithmetic"]
 unchecked-index = []
 unchecked-arithmetic = []
+dhat-heap = []    # if you are doing heap profiling
+dhat-ad-hoc = []  # if you are doing ad hoc profiling
+
 
 # Dev/test tooling (fake data generator).
 dev-tools = ["otap-df-core-nodes/dev-tools"]
@@ -300,6 +306,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,49 @@
 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,6 +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 {
+    once_cell::sync::Lazy,
+    std::sync::Mutex,
+    dhat::Profiler,
+};
 
 #[cfg(all(
     not(windows),
@@ -288,7 +295,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: Lazy<Mutex<Option<Profiler>>> = Lazy::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()