Merge pull request #22 from stackrox/feature/add-duration

Add benchmark duration
stackrox · May 7, 2024 · 87ad0d8 · 87ad0d8
2 parents 16ff022 + c5cd995
commit 87ad0d8
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 5 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -24,6 +24,10 @@ pub struct WorkloadConfig {
 
     /// Custom workload configuration.
     pub workload: Workload,
+
+    /// For how long to run the worker. Default value is zero, meaning no limit.
+    #[serde(default = "default_duration")]
+    pub duration: u64,
 }
 
 fn default_workers() -> usize {
@@ -34,6 +38,10 @@ fn default_per_core() -> bool {
     true
 }
 
+fn default_duration() -> u64 {
+    0
+}
+
 /// Workload specific configuration, contains one enum value for each
 /// workload type.
 #[derive(Debug, Copy, Clone, Deserialize)]

diff --git a/src/main.rs b/src/main.rs
@@ -21,16 +21,19 @@ use config::Config;
 use core_affinity::CoreId;
 use fork::{fork, Fork};
 use itertools::iproduct;
+use nix::sys::signal::{kill, Signal};
 use nix::sys::wait::waitpid;
 use nix::unistd::Pid;
-use std::env;
+use std::time::SystemTime;
+use std::{env, thread, time};
 
 use berserker::{worker::new_worker, WorkloadConfig};
 
 fn main() {
     let args: Vec<String> = env::args().collect();
     let default_config = String::from("workload.toml");
     let config_path = &args.get(1).unwrap_or(&default_config);
+    let duration_timer = SystemTime::now();
 
     let config = Config::builder()
         // Add in `./Settings.toml`
@@ -94,8 +97,38 @@ fn main() {
 
     info!("In total: {}", upper);
 
-    for handle in handles.into_iter().flatten() {
-        info!("waitpid: {}", handle);
-        waitpid(Pid::from_raw(handle), None).unwrap();
-    }
+    let processes = &handles.clone();
+
+    thread::scope(|s| {
+        if config.duration != 0 {
+            // Spin a watcher thread
+            s.spawn(move || loop {
+                thread::sleep(time::Duration::from_secs(1));
+                let elapsed = duration_timer.elapsed().unwrap().as_secs();
+
+                if elapsed > config.duration {
+                    for handle in processes.into_iter().flatten() {
+                        info!("Terminating: {}", *handle);
+                        match kill(Pid::from_raw(*handle), Signal::SIGTERM) {
+                            Ok(()) => {
+                                continue;
+                            }
+                            Err(e) => {
+                                continue;
+                            }
+                        }
+                    }
+
+                    break;
+                }
+            });
+        }
+
+        s.spawn(move || {
+            for handle in processes.into_iter().flatten() {
+                info!("waitpid: {}", *handle);
+                waitpid(Pid::from_raw(*handle), None).unwrap();
+            }
+        });
+    });
 }
diff --git a/src/worker/endpoints.rs b/src/worker/endpoints.rs
@@ -29,6 +29,7 @@ impl EndpointWorker {
             workload: _,
             per_core: _,
             workers: _,
+            duration: _,
         } = workload;
 
         EndpointWorker {