From 6270125b94f85dc41d76cad9b15e38e93c5d634d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Berkay=20=C3=9Cr=C3=BCn?=
 <berkay.ueruen@aisec.fraunhofer.de>
Date: Sat, 9 Sep 2023 17:23:45 +0200
Subject: [PATCH] Add signal handler for graceful exit

---
 controller.py | 35 ++++++++++++++++++++++++++++++++++-
 hdf5logger.py | 14 ++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/controller.py b/controller.py
index 1518183..6bfe89f 100755
--- a/controller.py
+++ b/controller.py
@@ -21,6 +21,7 @@
 import logging
 from multiprocessing import Manager, Process
 from pathlib import Path
+import signal
 import subprocess
 import sys
 import tables
@@ -47,6 +48,24 @@
 
 clogger = logging.getLogger(__name__)
 
+stop_signal_received = False
+
+
+def signal_handler(signum, frame):
+    global stop_signal_received
+    stop_signal_received = True
+
+
+def register_signal_handlers(clogger, queue_output, p_logger):
+    signal.signal(
+        signal.SIGTERM,
+        signal_handler,
+    )
+    signal.signal(
+        signal.SIGINT,
+        signal_handler,
+    )
+
 
 def build_ranges_dict(fault_dict):
     """
@@ -606,9 +625,23 @@ def controller(
             continue
         goldenrun_data[keyword] = pd.DataFrame(goldenrun_data[keyword])
 
-    pbar = tqdm(total=len(faultlist), desc="Simulating faults", disable=not len(faultlist))
+    # Handlers are used for a graceful exit, in case of a signal
+    register_signal_handlers(clogger, queue_output, p_logger)
+
+    pbar = tqdm(
+        total=len(faultlist), desc="Simulating faults", disable=not len(faultlist)
+    )
     itter = 0
     while 1:
+        if stop_signal_received:
+            clogger.info("Stopping the execution")
+            queue_output.put({"index": -4})
+
+            for p in p_list:
+                p["process"].kill()
+
+            break
+
         if len(p_list) == 0 and itter == len(faultlist):
             clogger.debug("Done inserting qemu jobs")
             break
diff --git a/hdf5logger.py b/hdf5logger.py
index d4b337c..bd137b0 100644
--- a/hdf5logger.py
+++ b/hdf5logger.py
@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import signal
 import logging
 import time
 
@@ -25,6 +26,14 @@
 logger = logging.getLogger(__name__)
 
 
+def register_signal_handlers():
+    """
+    Ignore signals, they will be handled by the controller.py anyway
+    """
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+
 # Tables for storing the elements from queue
 class translation_block_exec_table(tables.IsDescription):
     tb = tables.UInt64Col()
@@ -469,6 +478,8 @@ def hdf5collector(
     log_config=False,
     overwrite_faults=False,
 ):
+    register_signal_handlers()
+
     prctl.set_name("logger")
     prctl.set_proctitle("logger")
     f = tables.open_file(hdf5path, mode, max_group_width=65536)
@@ -540,6 +551,9 @@ def hdf5collector(
             process_backup(f, exp_group, exp, myfilter)
             log_config = False
             continue
+        elif exp["index"] == -4:
+            f.close()
+            return
         else:
             continue