Disable enable_heap_profiling and improve some codes (#380)

components/engine_panic/Cargo.toml

@@ -10,9 +10,9 @@ license = "Apache-2.0"
 testexport = []
 
 [dependencies]
+encryption = { workspace = true }
 engine_traits = { workspace = true }
 kvproto = { workspace = true }
-encryption = { workspace = true }
 raft = { workspace = true }
 tikv_alloc = { workspace = true }
 # FIXME: Remove this dep from the engine_traits interface

components/engine_rocks/Cargo.toml

@@ -62,6 +62,6 @@ package = "rocksdb"
 features = ["encryption"]
 
 [dev-dependencies]
+proptest = "1.0.0"
 rand = "0.8"
 toml = "0.5"
-proptest = "1.0.0"

components/engine_traits/Cargo.toml

@@ -11,6 +11,7 @@ testexport = []
 
 [dependencies]
 collections = { workspace = true }
+encryption = { workspace = true }
 error_code = { workspace = true }
 fail = "0.5"
 file_system = { workspace = true }
@@ -20,7 +21,6 @@ lazy_static = "1.0"
 log_wrappers = { workspace = true }
 protobuf = "2"
 raft = { workspace = true }
-encryption = { workspace = true }
 serde = "1.0"
 slog = { workspace = true }
 slog-global = { workspace = true }

components/engine_traits_tests/Cargo.toml

@@ -33,5 +33,5 @@ engine_traits = { workspace = true }
 kvproto = { workspace = true }
 panic_hook = { workspace = true }
 tempfile = "3.0"
-tikv_alloc = { workspace = true }
 test_util = { workspace = true }
+tikv_alloc = { workspace = true }

components/hybrid_engine/Cargo.toml

@@ -15,23 +15,23 @@ path = "tests/failpoints/mod.rs"
 required-features = ["failpoints"]
 
 [dependencies]
-engine_traits = { workspace = true }
-txn_types = { workspace = true }
-tikv_util = { workspace = true }
+crossbeam = { workspace = true }
 engine_rocks = { workspace = true }
+engine_traits = { workspace = true }
+fail = "0.5"
+kvproto = { workspace = true }
+lazy_static = "1.4.0"
 online_config = { workspace = true }
+prometheus = { version = "0.13", default-features = false, features = ["nightly"] }
+prometheus-static-metric = "0.5"
+raft = { workspace = true }
+raftstore = { workspace = true }
 range_cache_memory_engine = { workspace = true }
 slog = { workspace = true }
 slog-global = { workspace = true }
 tempfile = "3.0"
-prometheus = { version = "0.13", default-features = false, features = ["nightly"] }
-prometheus-static-metric = "0.5"
-lazy_static = "1.4.0"
-crossbeam = { workspace = true }
-fail = "0.5"
-raftstore = { workspace = true }
-raft = { workspace = true }
-kvproto = { workspace = true }
+tikv_util = { workspace = true }
+txn_types = { workspace = true }
 
 [dev-dependencies]
 tempfile = "3.0"

components/raft_log_engine/Cargo.toml

@@ -9,9 +9,9 @@ license = "Apache-2.0"
 failpoints = ["raft-engine/failpoints"]
 
 [dependencies]
+codec = { workspace = true }
 encryption = { workspace = true }
 engine_traits = { workspace = true }
-codec = { workspace = true }
 file_system = { workspace = true }
 kvproto = { workspace = true }
 lazy_static = "1.4.0"

components/range_cache_memory_engine/Cargo.toml

@@ -15,42 +15,42 @@ path = "tests/failpoints/mod.rs"
 required-features = ["failpoints"]
 
 [dependencies]
-engine_traits = { workspace = true }
-collections = { workspace = true }
-crossbeam-skiplist = { workspace = true }
 bytes = "1.0"
+collections = { workspace = true }
 crossbeam = { workspace = true }
+crossbeam-skiplist = { workspace = true }
+dashmap = "5.1"
+engine_rocks = { workspace = true }
+engine_traits = { workspace = true }
+fail = "0.5"
 futures = { version = "0.3", features = ["compat"] }
-tikv_util = { workspace = true }
-txn_types = { workspace = true }
+hex = "0.4"
+keys = { workspace = true }
 kvproto = { workspace = true }
+lazy_static = "1.4.0"
+libc = "0.2"
 log_wrappers = { workspace = true }
+online_config = { workspace = true }
+parking_lot = "0.12"
 pd_client = { workspace = true }
+prometheus = { version = "0.13", default-features = false, features = ["nightly"] }
+prometheus-static-metric = "0.5"
 raftstore = { workspace = true }
-dashmap = "5.1"
+rand = "0.8"
 security = { workspace = true }
 serde = "1.0"
 serde_derive = "1.0"
 serde_json = "1.0"
-slog-global = { workspace = true }
 slog = { workspace = true }
-engine_rocks = { workspace = true }
-fail = "0.5"
-yatp = { workspace = true }
-parking_lot = "0.12"
-keys = { workspace = true }
-prometheus = { version = "0.13", default-features = false, features = ["nightly"] }
-prometheus-static-metric = "0.5"
-lazy_static = "1.4.0"
-hex = "0.4"
+slog-global = { workspace = true }
 thiserror = "1.0"
-online_config = { workspace = true }
-libc = "0.2"
-rand = "0.8"
+tikv_util = { workspace = true }
+txn_types = { workspace = true }
+yatp = { workspace = true }
 
 [dev-dependencies]
+proptest = "1.0.0"
 tempfile = "3.0"
 test_pd = { workspace = true }
 test_pd_client = { workspace = true }
 test_util = { workspace = true }
-proptest = "1.0.0"

proxy_components/proxy_ffi/src/jemalloc_utils.rs

@@ -20,6 +20,8 @@ extern "C" {
     ) -> ::std::os::raw::c_int;
 }
 
+extern crate libc;
+
 #[allow(unused_variables)]
 #[allow(unused_mut)]
 #[allow(unused_unsafe)]
@@ -29,32 +31,42 @@ pub fn issue_mallctl_args(
     oldsize: *mut u64,
     newptr: *mut ::std::os::raw::c_void,
     newsize: u64,
-) {
+) -> ::std::os::raw::c_int {
     unsafe {
         let c_str = std::ffi::CString::new(command).unwrap();
         let c_ptr: *const ::std::os::raw::c_char = c_str.as_ptr() as *const ::std::os::raw::c_char;
         // See unprefixed_malloc_on_supported_platforms in tikv-jemalloc-sys.
         #[cfg(any(test, feature = "testexport"))]
         {
+            // Test part
             #[cfg(feature = "jemalloc")]
             {
                 // See NO_UNPREFIXED_MALLOC
                 #[cfg(any(target_os = "android", target_os = "dragonfly", target_os = "macos"))]
-                _rjem_mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
+                return _rjem_mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
                 #[cfg(not(any(
                     target_os = "android",
                     target_os = "dragonfly",
                     target_os = "macos"
                 )))]
-                mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
+                return mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
             }
+            0
         }
 
         #[cfg(not(any(test, feature = "testexport")))]
         {
-            // Must linked to tiflash.
+            // No test part
             #[cfg(feature = "external-jemalloc")]
-            let r = mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
+            {
+                // Must linked to tiflash.
+                return mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
+            }
+            #[cfg(not(feature = "external-jemalloc"))]
+            {
+                // Happens only with `raftstore-proxy-main`
+                return mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
+            }
         }
     }
 }

proxy_components/proxy_server/src/config.rs

@@ -279,6 +279,9 @@ pub struct ProxyConfig {
 
     #[online_config(skip)]
     pub engine_store: EngineStoreConfig,
+
+    #[online_config(submodule)]
+    pub memory: MemoryConfig,
 }
 
 /// We use custom default, in case of later non-ordinary config items.
@@ -296,6 +299,7 @@ impl Default for ProxyConfig {
             readpool: ReadPoolConfig::default(),
             import: ImportConfig::default(),
             engine_store: EngineStoreConfig::default(),
+            memory: MemoryConfig::default(),
         }
     }
 }
@@ -367,6 +371,15 @@ pub fn setup_default_tikv_config(default: &mut TikvConfig) {
     // Do not add here, try use `address_proxy_config`.
 }
 
+#[derive(Clone, Default, Serialize, Deserialize, PartialEq, Debug, OnlineConfig)]
+#[serde(default)]
+#[serde(rename_all = "kebab-case")]
+pub struct MemoryConfig {
+    // Whether enables the heap profiling which may have a bit performance overhead about 2% for
+    // the default sample rate.
+    pub enable_heap_profiling: bool,
+}
+
 /// This function changes TiKV's config according to ProxyConfig.
 /// Add a case in `test_config_proxy_default_no_config_item` to guard this
 /// logic.
@@ -433,6 +446,9 @@ pub fn address_proxy_config(config: &mut TikvConfig, proxy_config: &ProxyConfig)
     if proxy_config.engine_store.enable_fast_add_peer && !proxy_config.engine_store.enable_unips {
         fatal!("fast add peer can only work when using unips");
     }
+
+    // Currently, we do not support continous memory profiling.
+    config.memory.enable_heap_profiling = false;
 }
 
 pub fn validate_and_persist_config(config: &mut TikvConfig, persist: bool) {

proxy_components/proxy_server/src/status_server/mod.rs

@@ -1,7 +1,7 @@
 // Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
 
-mod profile;
-mod vendored_utils;
+pub mod profile;
+pub mod vendored_utils;
 
 use std::{
     error::Error as StdError,

proxy_components/proxy_server/src/status_server/profile.rs

@@ -234,7 +234,7 @@ pub fn read_file(path: &str) -> Result<Vec<u8>, String> {
 pub fn jeprof_heap_profile(path: &str) -> Result<Vec<u8>, String> {
     info!("using jeprof to process {}", path);
     let output = Command::new("./jeprof")
-        .args(["--show_bytes", "./bin/tikv-server", path, "--svg"])
+        .args(["--show_bytes", "./bin/tiflash/tiflash", path, "--svg"])
         .output()
         .map_err(|e| format!("jeprof: {}", e))?;
     if !output.status.success() {

proxy_components/proxy_server/src/status_server/vendored_utils.rs

@@ -5,6 +5,7 @@ use tikv_alloc::error::ProfResult;
 
 pub fn activate_prof() -> ProfResult<()> {
     {
+        tikv_util::debug!("activate_prof");
         let mut value: bool = true;
         let len = std::mem::size_of_val(&value) as u64;
         issue_mallctl_args(
@@ -18,8 +19,22 @@ pub fn activate_prof() -> ProfResult<()> {
     Ok(())
 }
 
+pub fn has_activate_prof() -> bool {
+    let mut value: bool = false;
+    let mut len = std::mem::size_of_val(&value) as u64;
+    issue_mallctl_args(
+        "prof.active",
+        &mut value as *mut _ as *mut _,
+        &mut len as *mut _ as *mut _,
+        std::ptr::null_mut(),
+        0,
+    );
+    value
+}
+
 pub fn deactivate_prof() -> ProfResult<()> {
     {
+        tikv_util::debug!("deactivate_prof");
         let mut value: bool = false;
         let len = std::mem::size_of_val(&value) as u64;
         issue_mallctl_args(
@@ -33,18 +48,64 @@ pub fn deactivate_prof() -> ProfResult<()> {
     Ok(())
 }
 
+extern crate libc;
+
 pub fn dump_prof(path: &str) -> tikv_alloc::error::ProfResult<()> {
     {
         let mut bytes = std::ffi::CString::new(path)?.into_bytes_with_nul();
         let mut ptr = bytes.as_mut_ptr() as *mut ::std::os::raw::c_char;
         let len = std::mem::size_of_val(&ptr) as u64;
-        issue_mallctl_args(
+        let r = issue_mallctl_args(
+            "prof.dump",
+            std::ptr::null_mut(),
+            std::ptr::null_mut(),
+            &mut ptr as *mut _ as *mut _,
+            len,
+        );
+        if r != 0 {
+            unsafe {
+                let err = *libc::__errno_location();
+                let err_msg = libc::strerror(err);
+                let c_str = std::ffi::CStr::from_ptr(err_msg);
+                let str_slice = c_str.to_str().unwrap_or("Unknown error");
+                tikv_util::warn!(
+                    "dump_prof returns non-zero {} error_code: {} error_message: {}",
+                    r,
+                    err,
+                    str_slice
+                );
+            }
+        }
+    }
+    Ok(())
+}
+
+pub fn adhoc_dump(path: &str) -> tikv_alloc::error::ProfResult<()> {
+    {
+        let mut bytes = std::ffi::CString::new(path)?.into_bytes_with_nul();
+        let mut ptr = bytes.as_mut_ptr() as *mut ::std::os::raw::c_char;
+        let len = std::mem::size_of_val(&ptr) as u64;
+        let r = issue_mallctl_args(
             "prof.dump",
             std::ptr::null_mut(),
             std::ptr::null_mut(),
             &mut ptr as *mut _ as *mut _,
             len,
         );
+        if r != 0 {
+            unsafe {
+                let err = *libc::__errno_location();
+                let err_msg = libc::strerror(err);
+                let c_str = std::ffi::CStr::from_ptr(err_msg);
+                let str_slice = c_str.to_str().unwrap_or("Unknown error");
+                tikv_util::warn!(
+                    "adhoc_dump returns non-zero {} error_code: {} error_message: {}",
+                    r,
+                    err,
+                    str_slice
+                );
+            }
+        }
     }
     Ok(())
 }

proxy_scripts/ci_check.sh

@@ -47,6 +47,7 @@ elif [[ $M == "testnew" ]]; then
     cargo check --package proxy_server --features="$ENABLE_FEATURES"
     # tests based on mock-engine-store, with compat for new proxy
     LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib cargo test --package proxy_tests --features="$ENABLE_FEATURES" --test proxy shared::jemalloc --features="jemalloc"
+    MALLOC_CONF="prof:true,prof_active:false" LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib cargo test --package proxy_tests --features="$ENABLE_FEATURES" --test proxy shared::pprof_jemalloc --features=mem-profiling --features="jemalloc"
     cargo test --package proxy_tests --features="$ENABLE_FEATURES" --test proxy shared::write
     cargo test --package proxy_tests --features="$ENABLE_FEATURES" --test proxy shared::snapshot
     cargo test --package proxy_tests --features="$ENABLE_FEATURES" --test proxy shared::config

proxy_scripts/make_env.sh

@@ -1,6 +1,6 @@
 wget https://github.com/jemalloc/jemalloc/releases/download/5.2.1/jemalloc-5.2.1.tar.bz2
 tar -xvf jemalloc-5.2.1.tar.bz2
 cd jemalloc-5.2.1
-./configure
+./configure --enable-prof
 make
 make install