Added a runtime parameter to the CPU stress test to automatically limit how long it runs

src/main.rs

@@ -39,7 +39,9 @@ enum CpuCommands {
     // CPU stress test subcommand
     #[clap(name = "stress", about = "Stress test the CPU with math!")]
     StressTest {
-        #[clap(short = 't', long, default_value_t = 0, help = "Number of threads to use; defaults to CPU's max thread count.")]
+        #[clap(short = 'r', long, default_value_t = 5, help = "Length of time (in minutes) to run the stress test. Defaults to 5")]
+        runtime: u16,
+        #[clap(short = 't', long, default_value_t = 0, help = "Number of threads to use; defaults to CPU's max thread count. Defaults to 0 (automatic)")]
         threads: usize,
     },
 }
@@ -122,7 +124,7 @@ fn main() {
     // map subcommands back to the main command
     match &cli.command {
         Commands::Cpu(args) => match &args.cpu_commands {
-            CpuCommands::StressTest { threads } => stress::cpu::cpu_stress_math(*threads),
+            CpuCommands::StressTest { runtime, threads } => stress::cpu::cpu_stress_math(*runtime, *threads),
         },
 
         Commands::Disk(args) => match &args.disk_commands {

src/stress/cpu.rs

@@ -1,28 +1,32 @@
-use std::thread;
+use std::{thread, time};
+use std::process::exit;
 use sysinfo::{System,SystemExt};
 
-pub fn cpu_stress_math(threads: usize) {
-	// fetch system information
-	let mut sys = System::new_all();
-	sys.refresh_all();
-	let num_cpus = sys.cpus().len();
+pub fn cpu_stress_math(runtime: u16, threads: usize) {
+    // fetch system information
+    let mut sys = System::new_all();
+    sys.refresh_all();
+    let num_cpus = sys.cpus().len();
 
-	let mut num_threads = threads;
-	if num_threads == 0 {
-		println!("Number of threads not specified, defaulting to CPU's thread count of {}.", num_cpus);
-		num_threads = num_cpus;
-	} else {
-		println!("Using specified thread count of {}", num_threads);
-	}
+    let mut num_threads = threads;
+    if num_threads == 0 {
+        println!("Number of threads not specified, defaulting to CPU's thread count of {}.", num_cpus);
+        num_threads = num_cpus;
+    } else {
+        println!("Using specified thread count of {}", num_threads);
+    }
 
-	for i in 1..num_threads {
-        println!("Spawning thread number {}", i);
+    for i in 0..num_threads {
+        println!("Spawning thread number {}", i + 1);
         thread::spawn (|| {
             worker();
         });
     }
-    println!("Using main as last thread");
-    worker();
+
+    println!("Sleeping main thread for the allotted runtime of {} minute(s).", runtime);
+    let duration = time::Duration::from_secs((runtime * 60).into());
+    thread::sleep(duration);
+    exit(0);
 }
 
 fn worker() {