From 132a0ee50146f8fea74e1a944439c590a70995d1 Mon Sep 17 00:00:00 2001
From: Gregory Ballantine <gballantine@metaunix.net>
Date: Wed, 17 Aug 2022 15:28:11 -0400
Subject: [PATCH] Added a runtime parameter to the CPU stress test to
 automatically limit how long it runs

---
 src/main.rs       |  6 ++++--
 src/stress/cpu.rs | 38 +++++++++++++++++++++-----------------
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 36264ab..9306575 100644
--- a/src/main.rs
+++ b/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 {
diff --git a/src/stress/cpu.rs b/src/stress/cpu.rs
index b0c2e13..b6c137b 100644
--- a/src/stress/cpu.rs
+++ b/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() {