Skip to content

Commit

Permalink
Restore install_signal_handlers due to downstream dependencies (#6366)
Browse files Browse the repository at this point in the history
  • Loading branch information
hendrikmakait authored May 18, 2022
1 parent 36e9946 commit ff94776
Showing 1 changed file with 31 additions and 0 deletions.
31 changes: 31 additions & 0 deletions distributed/cli/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
import signal
from typing import Any

from tornado.ioloop import IOLoop

logger = logging.getLogger(__name__)


Expand All @@ -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)

0 comments on commit ff94776

Please sign in to comment.