From 450fe91e938c1e15ce01ccaf23806cd118d6820a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Bourdeauducq?= Date: Tue, 27 Feb 2024 15:46:23 +0800 Subject: [PATCH] artiq_client: handle Ctrl-C gracefully --- artiq/frontend/artiq_client.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/artiq/frontend/artiq_client.py b/artiq/frontend/artiq_client.py index a83cfd774..00f2db602 100755 --- a/artiq/frontend/artiq_client.py +++ b/artiq/frontend/artiq_client.py @@ -22,6 +22,7 @@ from sipyco.pc_rpc import Client from sipyco.sync_struct import Subscriber from sipyco.broadcast import Receiver from sipyco import common_args, pyon +from sipyco.asyncio_tools import SignalHandler from artiq.tools import (scale_from_metadata, short_format, parse_arguments, parse_devarg_override) @@ -278,13 +279,20 @@ def _run_subscriber(host, port, subscriber): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: - loop.run_until_complete(subscriber.connect(host, port)) + signal_handler = SignalHandler() + signal_handler.setup() try: - loop.run_until_complete(asyncio.wait_for(subscriber.receive_task, - None)) - print("Connection to master lost") + loop.run_until_complete(subscriber.connect(host, port)) + try: + _, pending = loop.run_until_complete(asyncio.wait( + [signal_handler.wait_terminate(), subscriber.receive_task], + return_when=asyncio.FIRST_COMPLETED)) + for task in pending: + task.cancel() + finally: + loop.run_until_complete(subscriber.close()) finally: - loop.run_until_complete(subscriber.close()) + signal_handler.teardown() finally: loop.close()