add new flag "loop"

Author: zhangyi

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "py-spy-for-datakit"
-version = "0.3.15"
+version = "0.3.16"
 authors = ["Yi Zhang <[email protected]>"]
 repository = "https://github.com/GuanceCloud/py-spy-for-datakit"
 homepage = "https://github.com/GuanceCloud/py-spy-for-datakit"

Cargo.toml

@@ -59,6 +59,7 @@ pub struct Config {
     pub service: String,
     pub env: String,
     pub version: String,
+    pub loop_duration: u64,
 }
 
 #[allow(non_camel_case_types)]
@@ -128,7 +129,8 @@ impl Default for Config {
                port: 9529,
                service: String::from("unnamed-service"),
                env: String::from("unnamed-env"),
-               version: String::from("unnamed-version")
+               version: String::from("unnamed-version"),
+               loop_duration: 0,
         }
     }
 }
@@ -238,6 +240,14 @@ impl Config {
                 .help("The number of seconds to sample for")
                 .default_value("60")
                 .takes_value(true))
+            .arg(Arg::new("loop")
+                .short('L')
+                .long("loop")
+                .value_name("loop")
+                .help("continuously run profiler in a loop within the specified seconds, 0 represents infinity")
+                .default_value("0")
+                .takes_value(true)
+            )
             .arg(rate.clone())
             .arg(subprocesses.clone())
             .arg(gil.clone())
@@ -381,7 +391,10 @@ impl Config {
                         None => RecordDuration::Seconds(60),
                         Some(seconds) => RecordDuration::Seconds(seconds.parse().expect("invalid duration"))
                     };
-
+                    config.loop_duration = match matches.value_of("loop") {
+                        None => 0,
+                        Some(seconds) => seconds.parse().expect("invalid loop parameter")
+                    };
                     config.host = match matches.value_of("host") {
                         Some(host) => host.to_owned(),
                         None => String::from("127.0.0.1")

src/config.rs

@@ -240,6 +240,11 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error>
         }
     };
 
+    let max_total_intervals = match &config.loop_duration {
+        0 => None,
+        _ => Some(config.loop_duration * config.sampling_rate)
+    };
+
     use indicatif::ProgressBar;
     let progress = match (config.hide_progress, &config.duration) {
         (true, _) => ProgressBar::hidden(),
@@ -257,7 +262,9 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error>
 
     let mut errors = 0;
     let mut intervals = 0;
+    let mut total_intervals = 0;
     let mut samples = 0;
+    let mut loop_complete = false;
     println!();
 
     let running = Arc::new(AtomicBool::new(true));
@@ -294,6 +301,7 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error>
         }
 
         intervals += 1;
+        total_intervals += 1;
         if let Some(max_intervals) = max_intervals {
             if intervals >= max_intervals {
                 exit_message = "";
@@ -314,7 +322,15 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error>
 
                     output = new_output(config)?;
                     intervals = 0;
-                    profile_start_time = Utc::now()
+                    profile_start_time = Utc::now();
+
+                    if let Some(total) = max_total_intervals {
+                        if total_intervals >= total {
+                            loop_complete = true;
+                            break;
+                        }
+                    }
+
                 } else {
                     break;
                 }
@@ -378,13 +394,15 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error>
     if config.command.as_str() == cmd_datakit {
         let mut buf:Vec<u8> = Vec::with_capacity(128);
         output.write(&mut buf)?;
-        if let Some(tx) = channel_sender {
-            if let Err(e) = tx.send(Sample{
-                start: profile_start_time,
-                end:Utc::now(),
-                payload: buf,
-            }) {
-                eprintln!("fail to send data to channel: {:?}", e)
+        if !loop_complete {
+            if let Some(tx) = channel_sender {
+                if let Err(e) = tx.send(Sample{
+                    start: profile_start_time,
+                    end:Utc::now(),
+                    payload: buf,
+                }) {
+                    eprintln!("fail to send data to channel: {:?}", e)
+                }
             }
         }