Agent: Wait for master to terminate on windows CTRL_CLOSE_EVENT

monkey/infection_monkey/utils/signal_handler.py

@@ -3,7 +3,6 @@ import signal
 
 from infection_monkey.i_master import IMaster
 from infection_monkey.utils.environment import is_windows_os
-from infection_monkey.utils.exceptions.planned_shutdown_error import PlannedShutdownError
 
 logger = logging.getLogger(__name__)
 
@@ -12,28 +11,29 @@ class StopSignalHandler:
     def __init__(self, master: IMaster):
         self._master = master
 
-    def handle_posix_signals(self, signum, _):
-        self._handle_signal(signum)
-        # Windows signal handlers must return boolean. Only raising this exception for POSIX
-        # signals.
-        raise PlannedShutdownError("Monkey Agent got an interrupt signal")
+    def handle_posix_signals(self, signum: int, _):
+        self._handle_signal(signum, False)
 
-    def handle_windows_signals(self, signum):
+    def handle_windows_signals(self, signum: int):
         import win32con
 
-        # TODO: This signal handler gets called for a CTRL_CLOSE_EVENT, but the system immediately
-        #       kills the process after the handler returns. After the master is implemented and the
-        #       setup/teardown of the Agent is fully refactored, revisit this signal handler and
-        #       modify as necessary to more gracefully handle CTRL_CLOSE_EVENT signals.
-        if signum in {win32con.CTRL_C_EVENT, win32con.CTRL_BREAK_EVENT, win32con.CTRL_CLOSE_EVENT}:
-            self._handle_signal(signum)
+        if signum in {win32con.CTRL_C_EVENT, win32con.CTRL_BREAK_EVENT}:
+            self._handle_signal(signum, False)
+            return True
+
+        if signum == win32con.CTRL_CLOSE_EVENT:
+            # After the signal handler returns True, the OS will forcefully kill the process.
+            # Calling self._handle_signal() with block=True to give the master a chance to
+            # gracefully shut down. Note that the OS has a timeout that will forcefully kill the
+            # process if this handler hasn't returned in time.
+            self._handle_signal(signum, True)
             return True
 
         return False
 
-    def _handle_signal(self, signum):
+    def _handle_signal(self, signum: int, block: bool):
         logger.info(f"The Monkey Agent received signal {signum}")
-        self._master.terminate()
+        self._master.terminate(block)
 
 
 def register_signal_handlers(master: IMaster):