From ff9477678ad9cb29a2feadab60406d9726971805 Mon Sep 17 00:00:00 2001 From: Hendrik Makait Date: Wed, 18 May 2022 15:55:50 +0200 Subject: [PATCH] Restore install_signal_handlers due to downstream dependencies (#6366) --- distributed/cli/utils.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/distributed/cli/utils.py b/distributed/cli/utils.py index f25b7245c41..7b932e77364 100644 --- a/distributed/cli/utils.py +++ b/distributed/cli/utils.py @@ -5,6 +5,8 @@ import signal from typing import Any +from tornado.ioloop import IOLoop + logger = logging.getLogger(__name__) @@ -26,3 +28,32 @@ def handle_signal(signum, frame): old_handlers[sig] = signal.signal(sig, handle_signal) await event.wait() + + +def install_signal_handlers(loop=None, cleanup=None): + """ + Install global signal handlers to halt the Tornado IOLoop in case of + a SIGINT or SIGTERM. *cleanup* is an optional callback called, + before the loop stops, with a single signal number argument. + """ + import signal + + loop = loop or IOLoop.current() + + old_handlers = {} + + def handle_signal(sig, frame): + async def cleanup_and_stop(): + try: + if cleanup is not None: + await cleanup(sig) + finally: + loop.stop() + + loop.add_callback_from_signal(cleanup_and_stop) + # Restore old signal handler to allow for a quicker exit + # if the user sends the signal again. + signal.signal(sig, old_handlers[sig]) + + for sig in [signal.SIGINT, signal.SIGTERM]: + old_handlers[sig] = signal.signal(sig, handle_signal)