From 437d36e0ac6826a0d3e1e14bf514098d01c60d8c Mon Sep 17 00:00:00 2001
From: Yuhao Su <31772373+yuhao-su@users.noreply.github.com>
Date: Sun, 8 Oct 2023 12:50:31 +0800
Subject: [PATCH] fix(heap_profiling): free dumped file name str (#12649)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
---
 src/compute/src/memory_management/policy.rs   | 19 +++++++-------
 .../src/rpc/service/monitor_service.rs        | 25 +++++++++----------
 2 files changed, 21 insertions(+), 23 deletions(-)

diff --git a/src/compute/src/memory_management/policy.rs b/src/compute/src/memory_management/policy.rs
index 60efc57cd8af5..cea094b0b15d5 100644
--- a/src/compute/src/memory_management/policy.rs
+++ b/src/compute/src/memory_management/policy.rs
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::ffi::CStr;
+use std::ffi::CString;
 use std::path::Path;
 use std::sync::atomic::AtomicU64;
 use std::sync::Arc;
@@ -111,21 +111,20 @@ impl JemallocMemoryControl {
                 return;
             }
 
-            let time_prefix = chrono::Local::now().format("%Y-%m-%d-%H-%M-%S").to_string();
-            let file_name = format!("{}.{}\0", time_prefix, AUTO_DUMP_SUFFIX);
+            let time_prefix = chrono::Local::now().format("%Y-%m-%d-%H-%M-%S");
+            let file_name = format!("{}.{}", time_prefix, AUTO_DUMP_SUFFIX);
 
             let file_path = Path::new(&self.heap_profiling_config.dir)
-                .join(Path::new(&file_name))
+                .join(&file_name)
                 .to_str()
-                .unwrap()
-                .to_string();
+                .expect("file path is not valid utf8")
+                .to_owned();
+            let file_path_c = CString::new(file_path).expect("0 byte in file path");
 
-            // `file_path_str` is leaked because `jemalloc_dump_mib.write` requires static lifetime
-            let file_path_str = Box::leak(file_path.into_boxed_str());
-            let file_path_bytes = file_path_str.as_bytes();
+            // FIXME(yuhao): `unsafe` here because `jemalloc_dump_mib.write` requires static lifetime
             if let Err(e) = self
                 .jemalloc_dump_mib
-                .write(CStr::from_bytes_with_nul(file_path_bytes).unwrap())
+                .write(unsafe { &*(file_path_c.as_c_str() as *const _) })
             {
                 tracing::warn!("Auto Jemalloc dump heap file failed! {:?}", e);
             } else {
diff --git a/src/compute/src/rpc/service/monitor_service.rs b/src/compute/src/rpc/service/monitor_service.rs
index 2ab2c2968597a..ea95a87574572 100644
--- a/src/compute/src/rpc/service/monitor_service.rs
+++ b/src/compute/src/rpc/service/monitor_service.rs
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::ffi::CString;
 use std::fs;
 use std::path::Path;
 use std::sync::Arc;
@@ -121,7 +122,6 @@ impl MonitorService for MonitorServiceImpl {
         &self,
         request: Request<HeapProfilingRequest>,
     ) -> Result<Response<HeapProfilingResponse>, Status> {
-        use std::ffi::CStr;
         use std::fs::create_dir_all;
         use std::path::PathBuf;
 
@@ -139,9 +139,9 @@ impl MonitorService for MonitorServiceImpl {
             ));
         }
 
-        let time_prefix = chrono::Local::now().format("%Y-%m-%d-%H-%M-%S").to_string();
-        let file_name = format!("{}.{}\0", time_prefix, MANUALLY_DUMP_SUFFIX);
-        let arg_dir = request.into_inner().get_dir().clone();
+        let time_prefix = chrono::Local::now().format("%Y-%m-%d-%H-%M-%S");
+        let file_name = format!("{}.{}", time_prefix, MANUALLY_DUMP_SUFFIX);
+        let arg_dir = request.into_inner().dir;
         let dir = PathBuf::from(if arg_dir.is_empty() {
             &self.server_config.heap_profiling.dir
         } else {
@@ -153,20 +153,19 @@ impl MonitorService for MonitorServiceImpl {
         let file_path = file_path_buf
             .to_str()
             .ok_or_else(|| Status::internal("The file dir is not a UTF-8 String"))?;
+        let file_path_c =
+            CString::new(file_path).map_err(|_| Status::internal("0 byte in file path"))?;
 
-        // `file_path_str` is leaked because `prof::dump::write` requires static lifetime
-        let file_path_str = Box::leak(file_path.to_string().into_boxed_str());
-        let file_path_bytes = file_path_str.as_bytes();
-        let response = if let Err(e) = tikv_jemalloc_ctl::prof::dump::write(
-            CStr::from_bytes_with_nul(file_path_bytes).unwrap(),
-        ) {
+        // FIXME(yuhao): `unsafe` here because `jemalloc_dump_mib.write` requires static lifetime
+        if let Err(e) =
+            tikv_jemalloc_ctl::prof::dump::write(unsafe { &*(file_path_c.as_c_str() as *const _) })
+        {
             tracing::warn!("Manually Jemalloc dump heap file failed! {:?}", e);
             Err(Status::internal(e.to_string()))
         } else {
-            tracing::info!("Manually Jemalloc dump heap file created: {}", &file_path);
+            tracing::info!("Manually Jemalloc dump heap file created: {}", file_path);
             Ok(Response::new(HeapProfilingResponse {}))
-        };
-        response
+        }
     }
 
     #[cfg_attr(coverage, no_coverage)]