diff --git a/distributed/_signals.py b/distributed/_signals.py index f25b7245c4..730e14d0fc 100644 --- a/distributed/_signals.py +++ b/distributed/_signals.py @@ -25,4 +25,8 @@ def handle_signal(signum, frame): for sig in signals: old_handlers[sig] = signal.signal(sig, handle_signal) - await event.wait() + try: + await event.wait() + finally: + for sig in signals: + signal.signal(sig, old_handlers[sig]) diff --git a/distributed/cli/tests/test_dask_scheduler.py b/distributed/cli/tests/test_dask_scheduler.py index 1b01d1ad35..e6ea59867a 100644 --- a/distributed/cli/tests/test_dask_scheduler.py +++ b/distributed/cli/tests/test_dask_scheduler.py @@ -413,7 +413,7 @@ def test_version_option(): @pytest.mark.slow -def test_idle_timeout(loop): +def test_idle_timeout(): start = time() runner = CliRunner() result = runner.invoke( @@ -424,6 +424,23 @@ def test_idle_timeout(loop): assert result.exit_code == 0 +@pytest.mark.slow +def test_restores_signal_handler(): + # another test could have altered the signal handler, so use a new function + # that both has sensible sigint behaviour *and* can be used as a sentinel + def raise_ki(): + raise KeyboardInterrupt + + original_handler = signal.signal(signal.SIGINT, raise_ki) + try: + CliRunner().invoke( + distributed.cli.dask_scheduler.main, ["--idle-timeout", "1s"] + ) + assert signal.getsignal(signal.SIGINT) is raise_ki + finally: + signal.signal(signal.SIGINT, original_handler) + + def test_multiple_workers_2(loop): text = """ def dask_setup(worker):