Restore install_signal_handlers due to downstream dependencies

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)