asprof flamegraphs saved in tmp folder, duplicate jvm arg fix

aws · Jul 24, 2024 · 7105bec · 7105bec
1 parent 820288a
commit 7105bec
Show file tree

Hide file tree

Showing 9 changed files with 146 additions and 25 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -51,3 +51,4 @@ inferno = "0.11.19"
 indexmap = "2.1.0"
 cfg-if = "1.0"
 tempfile = "3"
+serial_test = "3.1.1"
diff --git a/src/data.rs b/src/data.rs
@@ -46,6 +46,7 @@ use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::fs::{File, OpenOptions};
 use std::ops::Sub;
+use std::path::PathBuf;
 use sysctldata::SysctlData;
 use systeminfo::SystemInfo;
 use vmstat::{Vmstat, VmstatRaw};
@@ -54,21 +55,23 @@ use vmstat::{Vmstat, VmstatRaw};
 pub struct CollectorParams {
     pub collection_time: u64,
     pub elapsed_time: u64,
-    pub data_file_path: String,
-    pub data_dir: String,
+    pub data_file_path: PathBuf,
+    pub data_dir: PathBuf,
     pub run_name: String,
     pub profile: HashMap<String, String>,
+    pub tmp_dir: PathBuf,
 }
 
 impl CollectorParams {
     fn new() -> Self {
         CollectorParams {
             collection_time: 0,
             elapsed_time: 0,
-            data_file_path: String::new(),
-            data_dir: String::new(),
+            data_file_path: PathBuf::new(),
+            data_dir: PathBuf::new(),
             run_name: String::new(),
             profile: HashMap::new(),
+            tmp_dir: PathBuf::new(),
         }
     }
 }
@@ -115,14 +118,15 @@ impl DataType {
         let full_path = format!("{}/{}", param.dir_name, name);
 
         self.file_name = name;
-        self.full_path = full_path;
+        self.full_path = full_path.clone();
         self.dir_name = param.dir_name.clone();
         self.collector_params.run_name = param.dir_name.clone();
         self.collector_params.collection_time = param.period;
         self.collector_params.elapsed_time = 0;
-        self.collector_params.data_file_path = self.full_path.clone();
-        self.collector_params.data_dir = param.dir_name.clone();
+        self.collector_params.data_file_path = PathBuf::from(full_path);
+        self.collector_params.data_dir = PathBuf::from(param.dir_name);
         self.collector_params.profile = param.profile.clone();
+        self.collector_params.tmp_dir = PathBuf::from(param.tmp_dir);
 
         self.file_handle = Some(
             OpenOptions::new()

diff --git a/src/data/flamegraphs.rs b/src/data/flamegraphs.rs
@@ -50,8 +50,10 @@ impl CollectData for FlamegraphRaw {
     fn after_data_collection(&mut self, params: CollectorParams) -> Result<()> {
         let data_dir = PathBuf::from(&params.data_dir);
 
-        let file_pathbuf =
-            data_dir.join(get_file_name(params.data_dir, "perf_profile".to_string())?);
+        let file_pathbuf = data_dir.join(get_file_name(
+            params.data_dir.display().to_string(),
+            "perf_profile".to_string(),
+        )?);
 
         let perf_jit_loc = data_dir.join("perf.data.jit");
 

diff --git a/src/data/java_profile.rs b/src/data/java_profile.rs
@@ -9,7 +9,6 @@ use log::{debug, error, trace};
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::io::Write;
-use std::path::PathBuf;
 use std::process::{Child, Command};
 use std::sync::Mutex;
 use std::{fs, fs::File};
@@ -39,17 +38,19 @@ impl JavaProfileRaw {
     }
 
     fn launch_asprof(&self, jids: Vec<String>, params: CollectorParams) -> Result<()> {
-        let data_dir = PathBuf::from(params.data_dir.clone());
         for jid in &jids {
-            let mut html_loc = data_dir.clone();
-            html_loc.push(format!("{}-java-flamegraph-{}.html", params.run_name, jid));
-
             match Command::new("asprof")
                 .args([
                     "-d",
                     &(params.collection_time - params.elapsed_time).to_string(),
                     "-f",
-                    html_loc.to_str().unwrap(),
+                    format!(
+                        "{}/{}-java-flamegraph-{}.html",
+                        params.tmp_dir.display(),
+                        params.run_name,
+                        jid
+                    )
+                    .as_str(),
                     jid.as_str(),
                 ])
                 .spawn()
@@ -79,7 +80,7 @@ impl JavaProfileRaw {
     fn get_jids(&mut self, arg: &str) -> Vec<String> {
         let mut jids: Vec<String> = Vec::new();
         for (key, value) in self.process_map.clone().into_iter() {
-            if arg == value[0] {
+            if arg == value[0] || arg == key {
                 jids.push(key);
             }
         }
@@ -163,12 +164,13 @@ impl CollectData for JavaProfileRaw {
                             error!("No JVM with name/PID '{}'.", arg);
                             continue;
                         }
-                        jids = self.get_jids(arg);
+                        jids.append(&mut self.get_jids(arg));
                     }
                 }
             }
         }
-
+        jids.sort();
+        jids.dedup();
         self.launch_asprof(jids, params.clone())
     }
 
@@ -210,7 +212,22 @@ impl CollectData for JavaProfileRaw {
             }
         }
 
-        let data_dir = PathBuf::from(params.data_dir.clone());
+        let data_dir = params.data_dir.clone();
+        for key in self.process_map.keys() {
+            let mut html_path = data_dir.clone();
+            html_path.push(format!("{}-java-flamegraph-{}.html", params.run_name, key));
+
+            let html_loc = html_path.to_str().unwrap();
+            let tmp_loc = format!(
+                "{}/{}-java-flamegraph-{}.html",
+                params.tmp_dir.display(),
+                params.run_name,
+                key
+            );
+
+            fs::rename(tmp_loc.clone(), html_loc).ok();
+        }
+
         let mut jps_map = File::create(
             data_dir
                 .clone()

diff --git a/src/data/perf_profile.rs b/src/data/perf_profile.rs
@@ -7,10 +7,10 @@ use anyhow::Result;
 use ctor::ctor;
 use log::{error, trace};
 use serde::{Deserialize, Serialize};
+use std::fs;
 use std::io::Write;
 use std::process::{Child, Command, Stdio};
 use std::sync::Mutex;
-use std::{fs, path::PathBuf};
 
 pub static PERF_PROFILE_FILE_NAME: &str = "perf_profile";
 pub static PERF_TOP_FUNCTIONS_FILE_NAME: &str = "top_functions";
@@ -48,7 +48,7 @@ impl CollectData for PerfProfileRaw {
                 "-e",
                 "cpu-clock:pppH",
                 "-o",
-                &params.data_file_path,
+                &params.data_file_path.display().to_string(),
                 "--",
                 "sleep",
                 &params.collection_time.to_string(),
@@ -88,7 +88,7 @@ impl CollectData for PerfProfileRaw {
             Ok(_) => trace!("'perf record' executed successfully."),
         }
         let mut top_functions_file =
-            fs::File::create(PathBuf::from(params.data_dir).join(PERF_TOP_FUNCTIONS_FILE_NAME))?;
+            fs::File::create(params.data_dir.join(PERF_TOP_FUNCTIONS_FILE_NAME))?;
 
         let out = Command::new("perf")
             .args([
@@ -97,7 +97,7 @@ impl CollectData for PerfProfileRaw {
                 "--percent-limit",
                 "1",
                 "-i",
-                &params.data_file_path,
+                &params.data_file_path.display().to_string(),
             ])
             .output();
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -478,6 +478,7 @@ pub struct InitParams {
     pub run_name: String,
     pub collector_version: String,
     pub commit_sha_short: String,
+    pub tmp_dir: String,
 }
 
 impl InitParams {
@@ -514,6 +515,7 @@ impl InitParams {
             run_name,
             collector_version,
             commit_sha_short,
+            tmp_dir: String::new(),
         }
     }
 }