From 5b4be35dbac38c5be858bc26a1d050307269035e Mon Sep 17 00:00:00 2001
From: Marc Mance <mmance@komputernerds.com>
Date: Thu, 1 Aug 2024 12:04:28 -0400
Subject: [PATCH] wrapped collect in try catch for keyboard interupt

---
 collect.py | 63 ++++++++++++++++++++++++++++++++++++------------------
 1 file changed, 42 insertions(+), 21 deletions(-)

diff --git a/collect.py b/collect.py
index 5ef484e..06651f7 100644
--- a/collect.py
+++ b/collect.py
@@ -259,12 +259,6 @@ def zpoolIostat():
 
 
 def collect_data(minutes):
-    if is_freebsd():
-        gstat_command = ["gstat", "-C", "-s", "-d", "-o", "-p", "-I", "5s"]
-        with open("gstat.csv", "a") as output_file:
-            # Create a Popen object with stdout redirected to the file
-            process = subprocess.Popen(gstat_command, stdout=output_file)
-
     for i in range(minutes):
         if i == 0:
             print("Minute:", end="", flush=True)
@@ -290,10 +284,6 @@ def collect_data(minutes):
 
         time.sleep(45)
 
-    # kill gstat if freebsd
-    if is_freebsd():
-        process.kill()
-
 
 def run_debug():
     print("Taking new debug.")
@@ -380,20 +370,51 @@ def main():
     minutes = int(input("Enter the duration in minutes: "))
     minutesToWait = int(input("Enter the delay before capture in minutes: "))
 
-    if not minutesToWait:
-        minutesToWait = 0
+    try:
+        if not minutesToWait:
+            minutesToWait = 0
 
-    if minutesToWait:
-        print(f"Delaying capture by {minutesToWait} minutes...")
-        for i in range(minutesToWait):
-            print(f"{i} ", end="", flush=True)
-            time.sleep(minutesToWait * 60)
-        print("")
+        if minutesToWait:
+            print(f"Delaying capture by {minutesToWait} minutes...")
+            for i in range(minutesToWait):
+                print(f"{i} ", end="", flush=True)
+                time.sleep(minutesToWait * 60)
+                print("")
 
-    print("Starting Collection")
+        print("Starting Collection")
 
-    # Collect data
-    collect_data(minutes)
+        # Collect data
+        if is_freebsd():
+            timeout = f"{minutes}m"
+            gstat_command = [
+                "timeout",
+                timeout,
+                "gstat",
+                "-C",
+                "-s",
+                "-d",
+                "-o",
+                "-p",
+                "-I",
+                "5s",
+            ]
+
+        with open("gstat.csv", "a") as output_file:
+            # Create a Popen object with stdout redirected to the file
+            process = subprocess.Popen(gstat_command, stdout=output_file)
+
+        collect_data(minutes)
+
+        # kill gstat if freebsd
+        if is_freebsd():
+            process.kill()
+
+    except KeyboardInterrupt:
+        print("Caught Ctrl-C,  cancelling collection...")
+        # kill gstat if freebsd
+        if is_freebsd():
+            process.kill()
+        exit()
 
     # Copy data files to /var/log (replace with appropriate copying function)
     collect_csv()