dump to prof.prefix by default

risingwavelabs · Sep 9, 2023 · 6cf7404 · 6cf7404
1 parent a4793d7
commit 6cf7404
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 14 deletions.
diff --git a/src/common/src/config.rs b/src/common/src/config.rs
@@ -658,18 +658,19 @@ impl AsyncStackTraceOption {
 
 #[derive(Clone, Debug, Serialize, Deserialize, DefaultFromSerde)]
 pub struct AutoDumpHeapProfileConfig {
+    /// Enable to auto dump heap profile when memory usage is high
+    #[serde(default = "default::auto_dump_heap_profile::enabled")]
+    pub enabled: bool,
+
+    /// The directory to dump heap profile. If empty, the prefix in `MALLOC_CONF` will be used
     #[serde(default = "default::auto_dump_heap_profile::dir")]
     pub dir: String,
+
+    /// The proportion (number between 0 and 1) of memory usage to trigger heap profile dump
     #[serde(default = "default::auto_dump_heap_profile::threshold")]
     pub threshold: f32,
 }
 
-impl AutoDumpHeapProfileConfig {
-    pub fn enabled(&self) -> bool {
-        !self.dir.is_empty()
-    }
-}
-
 serde_with::with_prefix!(streaming_prefix "stream_");
 serde_with::with_prefix!(batch_prefix "batch_");
 
@@ -1126,8 +1127,12 @@ pub mod default {
     }
 
     pub mod auto_dump_heap_profile {
+        pub fn enabled() -> bool {
+            true
+        }
+
         pub fn dir() -> String {
-            ".".to_string() // current directory
+            "".to_string()
         }
 
         pub fn threshold() -> f32 {

diff --git a/src/compute/src/memory_management/memory_manager.rs b/src/compute/src/memory_management/memory_manager.rs
@@ -54,7 +54,7 @@ impl GlobalMemoryManager {
                 .unwrap();
         tracing::info!("memory control policy: {:?}", &memory_control_policy);
 
-        if auto_dump_heap_profile_config.enabled() {
+        if auto_dump_heap_profile_config.enabled {
             fs::create_dir_all(&auto_dump_heap_profile_config.dir).unwrap();
         }
         Arc::new(Self {

diff --git a/src/compute/src/memory_management/policy.rs b/src/compute/src/memory_management/policy.rs
@@ -102,7 +102,7 @@ impl JemallocMemoryControl {
     }
 
     fn dump_heap_prof(&self, cur_used_memory_bytes: usize, prev_used_memory_bytes: usize) {
-        if !self.auto_dump_heap_profile_config.enabled() {
+        if !self.auto_dump_heap_profile_config.enabled {
             return;
         }
 
@@ -120,11 +120,21 @@ impl JemallocMemoryControl {
                 "{}.exceed-threshold-aggressive-heap-prof.compute.dump.{}\0",
                 time_prefix, self.dump_seq,
             );
-            let file_path = Path::new(&self.auto_dump_heap_profile_config.dir)
-                .join(Path::new(&file_name))
-                .to_str()
-                .unwrap()
-                .to_string();
+
+            let file_path = if !self.auto_dump_heap_profile_config.dir.is_empty() {
+                Path::new(&self.auto_dump_heap_profile_config.dir)
+                    .join(Path::new(&file_name))
+                    .to_str()
+                    .unwrap()
+                    .to_string()
+            } else {
+                let prof_prefix_mib = jemalloc_prof::prefix::mib().unwrap();
+                let prof_prefix = prof_prefix_mib.read().unwrap();
+                let mut file_path = prof_prefix.to_string_lossy().to_string();
+                file_path.push_str(&file_name);
+                file_path
+            };
+
             let file_path_str = Box::leak(file_path.into_boxed_str());
             let file_path_bytes = unsafe { file_path_str.as_bytes_mut() };
             let file_path_ptr = file_path_bytes.as_mut_ptr();

diff --git a/src/config/example.toml b/src/config/example.toml
@@ -7,6 +7,7 @@ metrics_level = "Info"
 telemetry_enabled = true
 
 [server.auto_dump_heap_profile]
+enabled = true
 dir = ""
 threshold = 0.8999999761581421