pick memory-profiling to release-7.5

proxy_components/engine_store_ffi/Cargo.toml

@@ -27,6 +27,8 @@ cloud-aws = ["sst_importer/cloud-aws"]
 cloud-gcp = ["sst_importer/cloud-gcp"]
 cloud-azure = ["sst_importer/cloud-azure"]
 
+external-jemalloc = ["proxy_ffi/external-jemalloc"]
+
 [dependencies]
 batch-system = { workspace = true, default-features = false }
 bitflags = "1.0.1"

proxy_components/proxy_ffi/Cargo.toml

@@ -23,6 +23,8 @@ test-engines-panic = [
   "engine_test/test-engines-panic",
 ]
 
+external-jemalloc = []
+
 [dependencies]
 encryption = { workspace = true, default-features = false }
 engine_rocks = { workspace = true, default-features = false }

proxy_components/proxy_ffi/src/jemalloc_utils.rs

@@ -0,0 +1,213 @@
+// Copyright 2024 TiKV Project Authors. Licensed under Apache-2.0.
+use std::sync::Mutex;
+
+extern "C" {
+    // External jemalloc
+    pub fn mallctl(
+        name: *const ::std::os::raw::c_char,
+        oldp: *mut ::std::os::raw::c_void,
+        oldlenp: *mut u64,
+        newp: *mut ::std::os::raw::c_void,
+        newlen: u64,
+    ) -> ::std::os::raw::c_int;
+
+    // Embedded jemalloc
+    pub fn _rjem_mallctl(
+        name: *const ::std::os::raw::c_char,
+        oldp: *mut ::std::os::raw::c_void,
+        oldlenp: *mut u64,
+        newp: *mut ::std::os::raw::c_void,
+        newlen: u64,
+    ) -> ::std::os::raw::c_int;
+
+    pub fn malloc_stats_print(
+        write_cb: Option<unsafe extern "C" fn(*mut c_void, *const ::std::os::raw::c_char)>,
+        cbopaque: *mut c_void,
+        opts: *const ::std::os::raw::c_char,
+    );
+}
+
+#[allow(unused_variables)]
+#[allow(unused_mut)]
+#[allow(unused_unsafe)]
+#[allow(unreachable_code)]
+pub fn issue_mallctl_args(
+    command: &str,
+    oldptr: *mut ::std::os::raw::c_void,
+    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"))]
+                return _rjem_mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
+                #[cfg(not(any(
+                    target_os = "android",
+                    target_os = "dragonfly",
+                    target_os = "macos"
+                )))]
+                return mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
+            }
+            0
+        }
+
+        #[cfg(not(any(test, feature = "testexport")))]
+        {
+            // No test part
+            #[cfg(feature = "external-jemalloc")]
+            {
+                // Must linked to tiflash.
+                return mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
+            }
+            #[cfg(not(feature = "external-jemalloc"))]
+            {
+                // Happens only with `raftstore-proxy-main`
+                #[cfg(not(any(
+                    target_os = "android",
+                    target_os = "dragonfly",
+                    target_os = "macos"
+                )))]
+                {
+                    return mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
+                }
+                0
+            }
+        }
+    }
+}
+
+#[allow(unused_variables)]
+#[allow(unused_mut)]
+#[allow(unused_unsafe)]
+pub fn issue_mallctl(command: &str) -> u64 {
+    type PtrUnderlying = u64;
+    let mut ptr: PtrUnderlying = 0;
+    let mut size = std::mem::size_of::<PtrUnderlying>() as u64;
+    issue_mallctl_args(
+        command,
+        &mut ptr as *mut _ as *mut ::std::os::raw::c_void,
+        &mut size as *mut u64,
+        std::ptr::null_mut(),
+        0,
+    );
+    ptr
+}
+
+pub fn get_allocatep_on_thread_start() -> u64 {
+    issue_mallctl("thread.allocatedp")
+}
+
+pub fn get_deallocatep_on_thread_start() -> u64 {
+    issue_mallctl("thread.deallocatedp")
+}
+
+pub fn get_allocate() -> u64 {
+    issue_mallctl("thread.allocated")
+}
+
+pub fn get_deallocate() -> u64 {
+    issue_mallctl("thread.deallocated")
+}
+
+use std::ffi::{c_char, c_void, CStr};
+struct CaptureContext {
+    buffer: Mutex<String>,
+}
+
+#[allow(dead_code)]
+extern "C" fn write_to_string(ctx: *mut c_void, message: *const c_char) {
+    if ctx.is_null() || message.is_null() {
+        return;
+    }
+
+    let context = unsafe { &*(ctx as *mut CaptureContext) };
+
+    let c_str = unsafe { CStr::from_ptr(message) };
+    if let Ok(str_slice) = c_str.to_str() {
+        let mut buffer = context.buffer.lock().unwrap();
+        buffer.push_str(str_slice);
+    }
+}
+
+#[allow(unused_variables)]
+#[allow(unused_mut)]
+#[allow(unused_unsafe)]
+pub fn get_malloc_stats() -> String {
+    let context = CaptureContext {
+        buffer: Mutex::new(String::new()),
+    };
+
+    // Use Json format.
+    let c_str = std::ffi::CString::new("J").unwrap();
+    let ops_str = c_str.as_ptr() as *const std::os::raw::c_char;
+
+    unsafe {
+        // 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_malloc_stats_print(
+                    Some(write_to_string),
+                    &context as *const _ as *mut c_void,
+                    ops_str,
+                );
+                #[cfg(not(any(
+                    target_os = "android",
+                    target_os = "dragonfly",
+                    target_os = "macos"
+                )))]
+                malloc_stats_print(
+                    Some(write_to_string),
+                    &context as *const _ as *mut c_void,
+                    ops_str,
+                );
+            }
+        }
+
+        #[cfg(not(any(test, feature = "testexport")))]
+        {
+            // No test part
+            #[cfg(feature = "external-jemalloc")]
+            {
+                // Must linked to tiflash.
+                malloc_stats_print(
+                    Some(write_to_string),
+                    &context as *const _ as *mut c_void,
+                    ops_str,
+                );
+            }
+            #[cfg(not(feature = "external-jemalloc"))]
+            {
+                // Happens only with `raftstore-proxy-main`
+                #[cfg(not(any(
+                    target_os = "android",
+                    target_os = "dragonfly",
+                    target_os = "macos"
+                )))]
+                {
+                    malloc_stats_print(
+                        Some(write_to_string),
+                        &context as *const _ as *mut c_void,
+                        ops_str,
+                    );
+                }
+            }
+        }
+    }
+
+    let buffer = context.buffer.lock().unwrap();
+    buffer.clone()
+}

proxy_components/proxy_ffi/src/lib.rs

@@ -21,6 +21,7 @@ pub mod raftstore_proxy;
 pub mod raftstore_proxy_helper_impls;
 pub mod read_index_helper;
 // FFI releated with reading from SST/RocksDB files.
+pub mod jemalloc_utils;
 pub mod snapshot_reader_impls;
 pub mod utils;
 

proxy_components/proxy_server/Cargo.toml

@@ -35,9 +35,12 @@ backup-stream-debug = ["backup-stream/backup-stream-debug"]
 testexport = ["engine_tiflash/testexport", "engine_store_ffi/testexport", "tikv/testexport"]
 pprof-fp = ["tikv/pprof-fp"]
 
+external-jemalloc = ["engine_store_ffi/external-jemalloc"]
+
 [dependencies]
 api_version = { workspace = true }
 async-stream = "0.2"
+backtrace = "0.3"
 backup = { workspace = true, default-features = false }
 backup-stream = { workspace = true, default-features = false }
 causal_ts = { workspace = true }
@@ -80,6 +83,7 @@ pin-project = "1.0"
 pprof = { version = "0.11", default-features = false, features = ["flamegraph", "protobuf-codec", "cpp"] }
 prometheus = { version = "0.13", features = ["nightly"] }
 protobuf = { version = "2.8", features = ["bytes"] }
+proxy_ffi = { workspace = true, default-features = false }
 raft = { version = "0.7.0", default-features = false, features = ["protobuf-codec"] }
 raft_log_engine = { workspace = true, default-features = false }
 raftstore = { workspace = true, default-features = false }

proxy_components/proxy_server/src/run.rs

@@ -613,6 +613,11 @@ impl<ER: RaftEngine, F: KvFormat> TiKvServer<ER, F> {
             Arc::clone(&security_mgr),
         );
 
+        #[cfg(feature = "external-jemalloc")]
+        info!("linked with external jemalloc");
+        #[cfg(not(feature = "external-jemalloc"))]
+        info!("linked without external jemalloc");
+
         // Initialize and check config
         info!("using proxy config"; "config" => ?proxy_config);
         crate::config::address_proxy_config(&mut config, &proxy_config);