diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index 71f199e39..f011f5ca5 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -47,7 +47,10 @@ def start_background_process(): # Daemon is set to True so that the process kills itself when the main process dies background_process = Process( - target=AikidoBackgroundProcess, args=(comms.address, comms.key), daemon=True + target=AikidoBackgroundProcess, + args=(comms.address, comms.key), + name="zen-agent-process", + daemon=True, ) background_process.start() diff --git a/aikido_zen/background_process/aikido_background_process.py b/aikido_zen/background_process/aikido_background_process.py index dc0162110..881e03b9b 100644 --- a/aikido_zen/background_process/aikido_background_process.py +++ b/aikido_zen/background_process/aikido_background_process.py @@ -3,6 +3,7 @@ """ import multiprocessing.connection as con +import signal import time import sched import traceback @@ -40,12 +41,14 @@ def __init__(self, address, key): self.queue = Queue() self.connection_manager = None # Start reporting thread : - Thread(target=self.reporting_thread).start() + Thread(target=self.reporting_thread, daemon=True).start() logger.debug( "Background process started successfully, with UDS File : %s", address ) + add_exit_handlers() + while True: conn = listener.accept() while True: @@ -105,3 +108,17 @@ def send_to_connection_manager(self, event_scheduler): blocked=queue_attack_item[2], stack=queue_attack_item[3], ) + + +def add_exit_handlers(): + """ + We add graceful exit handlers here since the process keeps hanging otherwise. + """ + + def exit_gracefully(sig, frame): + sys.exit(0) + + signal.signal(signal.SIGINT, exit_gracefully) + signal.signal(signal.SIGTERM, exit_gracefully) + signal.signal(signal.SIGQUIT, exit_gracefully) + signal.signal(signal.SIGHUP, exit_gracefully) diff --git a/aikido_zen/decorators/gunicorn.py b/aikido_zen/decorators/gunicorn.py index 9eb64e606..599c70aaf 100644 --- a/aikido_zen/decorators/gunicorn.py +++ b/aikido_zen/decorators/gunicorn.py @@ -5,10 +5,6 @@ import aikido_zen -# Run our background process as a child of gunicorn (exits safely) -aikido_zen.protect("daemon_only") - - def post_fork(prev_func): """ Aikido decorator for gunicorn config @@ -16,7 +12,7 @@ def post_fork(prev_func): """ def aik_post_fork(server, worker): - aikido_zen.protect("daemon_disabled") + aikido_zen.protect() prev_func(server, worker) return aik_post_fork