[feature]: Use jemalloc in Unix & able to control dump heap profile using http interface

Cargo.toml

@@ -64,10 +64,6 @@ members = [
     "bin/gravity_node"
 ]
 
-
-
-
-
 [workspace.dependencies]
 aptos-consensus = { path = "./aptos-core/consensus" }
 
@@ -302,11 +298,9 @@ inferno = "0.11.14"
 internment = { version = "0.5.0", features = ["arc"] }
 ipnet = "2.5.0"
 itertools = "0.13"
-jemallocator = { version = "0.5.0", features = [
-    "profiling",
-    "unprefixed_malloc_on_supported_platforms",
-] }
-jemalloc-sys = "0.5.4"
+tikv-jemalloc-ctl = "0.6"
+tikv-jemallocator = { version = "0.6" }
+tikv-jemalloc-sys = { version = "0.6" }
 json-patch = "0.2.6"
 jsonwebtoken = "8.1"
 jwt = "0.16.0"

Makefile

@@ -20,15 +20,15 @@ CARGO_FEATURES := $(if $(FEATURE),--features $(FEATURE),)
 all: $(BINARY)
 
 gravity_node:
-	cd bin/gravity_node && cargo build $(CARGO_FLAGS) $(CARGO_FEATURES)
+	cargo build -p gravity_node $(CARGO_FLAGS) $(CARGO_FEATURES)
 
 bench:
-	cd bin/bench && cargo build $(CARGO_FLAGS) $(CARGO_FEATURES)
+	cargo build -p bench $(CARGO_FLAGS) $(CARGO_FEATURES)
 
 kvstore:
-	cd bin/kvstore && cargo build $(CARGO_FLAGS) $(CARGO_FEATURES)
+	cargo build -p kvstore $(CARGO_FLAGS) $(CARGO_FEATURES)
 
 clean:
 	for dir in $(BIN_PATHS); do \
-	    (cd $$dir && cargo clean); \
-	done
+		(cd $$dir && cargo clean); \
+	done

crates/api/Cargo.toml

@@ -65,7 +65,12 @@ rcgen = "0.9"
 tokio-test = "*"
 reqwest = { version = "0.12.9", features = ["rustls-tls", "json"] }
 coex-bridge = { path = "../coex-bridge" }
+tikv-jemallocator.workspace = true
+tikv-jemalloc-ctl.workspace = true
+tikv-jemalloc-sys.workspace = true
+once_cell = { workspace = true }
 
 [features]
 default = []
-failpoints = ["fail/failpoints", "aptos-consensus/failpoints", "aptos-mempool/failpoints", "aptos-config/failpoints"]
+failpoints = ["fail/failpoints", "aptos-consensus/failpoints", "aptos-mempool/failpoints", "aptos-config/failpoints"]
+jemalloc-profiling = ["tikv-jemallocator/profiling", "tikv-jemalloc-sys/profiling"]

crates/api/src/consensus_api.rs

@@ -1,4 +1,4 @@
-use std::{collections::BTreeMap, sync::Arc};
+use std::sync::Arc;
 
 use crate::{
     bootstrap::{
@@ -11,7 +11,8 @@ use crate::{
     network::{create_network_runtime, extract_network_configs},
 };
 use api_types::{
-    compute_res::ComputeRes, u256_define::BlockId, ConsensusApi, ExecError, ExecutionLayer, ExternalBlock, ExternalBlockMeta
+    compute_res::ComputeRes, u256_define::BlockId, ConsensusApi, ExecError, ExecutionLayer,
+    ExternalBlock, ExternalBlockMeta,
 };
 use aptos_build_info::build_information;
 use aptos_config::{config::NodeConfig, network_id::NetworkId};
@@ -20,14 +21,18 @@ use aptos_consensus::gravity_state_computer::ConsensusAdapterArgs;
 use aptos_event_notifications::EventNotificationSender;
 use aptos_logger::{info, warn};
 use aptos_network_builder::builder::NetworkBuilder;
-use aptos_storage_interface::{DbReader, DbReaderWriter};
+use aptos_storage_interface::DbReaderWriter;
 use aptos_telemetry::service::start_telemetry_service;
 use async_trait::async_trait;
 
 use aptos_types::chain_id::ChainId;
 use futures::channel::mpsc;
 use tokio::runtime::Runtime;
 
+#[cfg(unix)]
+#[global_allocator]
+static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 pub struct ConsensusEngine {
     address: String,
     execution_layer: ExecutionLayer,
@@ -56,7 +61,6 @@ fn fail_point_check(node_config: &NodeConfig) {
 }
 
 impl ConsensusEngine {
-
     pub fn init(
         node_config: NodeConfig,
         execution_layer: ExecutionLayer,
@@ -167,19 +171,25 @@ impl ConsensusEngine {
         );
         runtimes.push(consensus_runtime);
         // trigger this to make epoch manager invoke new epoch
-        if !node_config.https_cert_pem_path.to_str().unwrap().is_empty()
-            && !node_config.https_key_pem_path.to_str().unwrap().is_empty()
-        {
-            let args = HttpsServerArgs {
-                address: node_config.https_server_address,
-                execution_api: execution_layer.execution_api.clone(),
-                cert_pem: node_config.https_cert_pem_path,
-                key_pem: node_config.https_key_pem_path,
-            };
-            let runtime = aptos_runtimes::spawn_named_runtime("Http".into(), None);
-            runtime.spawn(async move { https_server(args) });
-            runtimes.push(runtime);
-        }
+        let args = HttpsServerArgs {
+            address: node_config.https_server_address,
+            execution_api: execution_layer.execution_api.clone(),
+            cert_pem: node_config
+                .https_cert_pem_path
+                .clone()
+                .to_str()
+                .filter(|s| !s.is_empty())
+                .map(|_| node_config.https_cert_pem_path),
+            key_pem: node_config
+                .https_key_pem_path
+                .clone()
+                .to_str()
+                .filter(|s| !s.is_empty())
+                .map(|_| node_config.https_key_pem_path),
+        };
+        let runtime = aptos_runtimes::spawn_named_runtime("Http".into(), None);
+        runtime.spawn(async move { https_server(args) });
+        runtimes.push(runtime);
         let arc_consensus_engine = Arc::new(Self {
             address: node_config.validator_network.as_ref().unwrap().listen_address.to_string(),
             execution_layer: execution_layer.clone(),
@@ -202,8 +212,8 @@ impl ConsensusApi for ConsensusEngine {
             .recv_ordered_block(BlockId(parent_id), ordered_block)
             .await
         {
-            Ok(_) => {},
-            Err(ExecError::DuplicateExecError) => {},
+            Ok(_) => {}
+            Err(ExecError::DuplicateExecError) => {}
             Err(_) => panic!("send_ordered_block should not fail"),
         }
     }

crates/api/src/https/heap_profiler.rs

@@ -0,0 +1,70 @@
+use aptos_logger::info;
+use aptos_logger::warn;
+use axum::response::IntoResponse;
+use axum::Json;
+use once_cell::sync::Lazy;
+use serde::Deserialize;
+use serde::Serialize;
+use std::env;
+use std::sync::{Arc, Mutex};
+use tikv_jemalloc_ctl::raw;
+
+pub struct HeapProfiler {
+    mutex: Arc<Mutex<()>>,
+}
+
+const PROF_ACTIVE: &[u8] = b"prof.active\0";
+const PROF_THREAD_ACTIVE_INIT: &[u8] = b"prof.thread_active_init\0";
+
+pub static PROFILER: Lazy<HeapProfiler> = Lazy::new(|| HeapProfiler::new());
+
+#[derive(Deserialize, Serialize)]
+pub struct ControlProfileRequest {
+    enable: bool,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct ControlProfileResponse {
+    pub response: String,
+}
+
+/// User should use binary with feature api/jemalloc-profiling enabled.
+/// This feature can be enabled by ```Cargo build --features api/jemalloc-profiling```
+pub async fn control_profiler(request: ControlProfileRequest) -> impl IntoResponse {
+    #[cfg(feature = "jemalloc-profiling")]
+    match PROFILER.set_prof_active(request.enable) {
+        Ok(_) => Json(ControlProfileResponse { response: "success".to_string() }),
+        Err(e) => Json(ControlProfileResponse { response: e }),
+    }
+    #[cfg(not(feature = "jemalloc-profiling"))]
+    Json(ControlProfileResponse {
+        response: "jemalloc profiling is not enabled".to_string(),
+    })
+}
+
+impl HeapProfiler {
+    pub fn new() -> Self {
+        Self { mutex: Arc::new(Mutex::new(())) }
+    }
+
+    pub fn set_prof_active(&self, prof: bool) -> Result<(), String> {
+        let _guard = self.mutex.lock().unwrap();
+        if let Err(err) = unsafe { raw::write(PROF_ACTIVE, prof) } {
+            let err = String::from(format!("jemalloc heap profiling active failed: {}", err));
+            warn!("{}", err);
+            return Err(err);
+        }
+        if let Err(err) = unsafe { raw::write(PROF_THREAD_ACTIVE_INIT, prof) } {
+            let err =
+                String::from(format!("jemalloc heap profiling thread_active_init failed: {}", err));
+            warn!("{}", err);
+            return Err(err);
+        }
+        if prof {
+            info!("jemalloc heap profiling started");
+        } else {
+            info!("jemalloc heap profiling stopped");
+        }
+        Ok(())
+    }
+}

crates/api/src/https/mod.rs

@@ -1,3 +1,4 @@
+pub mod heap_profiler;
 mod set_failpoints;
 mod tx;
 use std::{net::SocketAddr, path::PathBuf, sync::Arc};
@@ -15,14 +16,15 @@ use axum::{
     Json, Router,
 };
 use axum_server::tls_rustls::RustlsConfig;
+use heap_profiler::control_profiler;
 use set_failpoints::{set_failpoint, FailpointConf};
 use tx::{get_tx_by_hash, submit_tx, TxRequest};
 
 pub struct HttpsServerArgs {
     pub address: String,
     pub execution_api: Arc<dyn ExecutionChannel>,
-    pub cert_pem: PathBuf,
-    pub key_pem: PathBuf,
+    pub cert_pem: Option<PathBuf>,
+    pub key_pem: Option<PathBuf>,
 }
 
 async fn ensure_https(req: Request<Body>, next: Next) -> Response {
@@ -47,25 +49,44 @@ pub async fn https_server(args: HttpsServerArgs) {
     let set_fail_point_lambda =
         |Json(request): Json<FailpointConf>| async move { set_failpoint(request).await };
 
+    let control_profiler_lambda = |Json(request): Json<heap_profiler::ControlProfileRequest>| async move {
+        control_profiler(request).await
+    };
+
     let https_app = Router::new()
         .route("/tx/submit_tx", post(submit_tx_lambda))
         .route("/tx/get_tx_by_hash/:hash_value", get(get_tx_by_hash_lambda))
         .layer(middleware::from_fn(ensure_https));
-    let http_app = Router::new().route("/set_failpoint", post(set_fail_point_lambda));
+    let http_app = Router::new()
+        .route("/set_failpoint", post(set_fail_point_lambda))
+        .route("/mem_prof", post(control_profiler_lambda));
     let app = Router::new().merge(https_app).merge(http_app);
-    // configure certificate and private key used by https
-    let config = RustlsConfig::from_pem_file(args.cert_pem.clone(), args.key_pem.clone())
-        .await
-        .unwrap_or_else(|e| {
-            panic!("error {:?}, cert {:?}, key {:?} doesn't work", e, args.cert_pem, args.key_pem)
-        });
     let addr: SocketAddr = args.address.parse().unwrap();
-    info!("https server listen address {}", addr);
-    axum_server::bind_rustls(addr, config).serve(app.into_make_service()).await.unwrap_or_else(
-        |e| {
-            panic!("failed to bind rustls due to {:?}", e);
-        },
-    );
+    match (args.cert_pem.clone(), args.key_pem.clone()) {
+        (Some(cert_path), Some(key_path)) => {
+            // configure certificate and private key used by https
+            let config =
+                RustlsConfig::from_pem_file(cert_path, key_path).await.unwrap_or_else(|e| {
+                    panic!(
+                        "error {:?}, cert {:?}, key {:?} doesn't work",
+                        e, args.cert_pem, args.key_pem
+                    )
+                });
+            info!("https server listen address {}", addr);
+            axum_server::bind_rustls(addr, config)
+                .serve(app.into_make_service())
+                .await
+                .unwrap_or_else(|e| {
+                    panic!("failed to bind rustls due to {:?}", e);
+                });
+        }
+        _ => {
+            info!("http server listen address {}", addr);
+            axum_server::bind(addr).serve(app.into_make_service()).await.unwrap_or_else(|e| {
+                panic!("failed to bind http due to {:?}", e);
+            });
+        }
+    }
 }
 
 #[cfg(test)]
@@ -96,15 +117,15 @@ mod test {
         let cert_pem = cert.serialize_pem().unwrap();
         let key_pem = cert.serialize_private_key_pem();
         let dir = env!("CARGO_MANIFEST_DIR").to_owned();
-        fs::create_dir(dir.clone() + "/src/https/test");
-        fs::write(dir.clone() + "/src/https/test/cert.pem", cert_pem);
-        fs::write(dir.clone() + "/src/https/test/key.pem", key_pem);
+        fs::create_dir(dir.clone() + "/src/https/test").unwrap();
+        fs::write(dir.clone() + "/src/https/test/cert.pem", cert_pem).unwrap();
+        fs::write(dir.clone() + "/src/https/test/key.pem", key_pem).unwrap();
 
         let args = HttpsServerArgs {
             address: "127.0.0.1:5425".to_owned(),
             execution_api: Arc::new(MockExecutionApi {}),
-            cert_pem: PathBuf::from(dir.clone() + "/src/https/test/cert.pem"),
-            key_pem: PathBuf::from(dir.clone() + "/src/https/test/key.pem"),
+            cert_pem: Some(PathBuf::from(dir.clone() + "/src/https/test/cert.pem")),
+            key_pem: Some(PathBuf::from(dir.clone() + "/src/https/test/key.pem")),
         };
         let _handler = tokio::spawn(https_server(args));
         tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;