From 5f77d4f5fa4e9657a5ab0e2fea4fe906d33f8e52 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Thu, 12 Jan 2023 12:35:02 +0800
Subject: [PATCH] applets: fix asyncio loop management

---
 artiq/frontend/artiq_browser.py   |  6 +++---
 artiq/frontend/artiq_dashboard.py |  3 ++-
 artiq/gui/applets.py              | 14 ++++++++------
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/artiq/frontend/artiq_browser.py b/artiq/frontend/artiq_browser.py
index 49517f9cb..853b8ed31 100755
--- a/artiq/frontend/artiq_browser.py
+++ b/artiq/frontend/artiq_browser.py
@@ -49,7 +49,7 @@ def get_argparser():
 
 class Browser(QtWidgets.QMainWindow):
     def __init__(self, smgr, datasets_sub, browse_root,
-                 master_host, master_port):
+                 master_host, master_port, *, loop=None):
         QtWidgets.QMainWindow.__init__(self)
         smgr.register(self)
 
@@ -81,7 +81,7 @@ class Browser(QtWidgets.QMainWindow):
         self.files.dataset_changed.connect(
             self.experiments.dataset_changed)
 
-        self.applets = applets.AppletsDock(self, datasets_sub)
+        self.applets = applets.AppletsDock(self, datasets_sub, loop=loop)
         smgr.register(self.applets)
         atexit_register_coroutine(self.applets.stop, loop=loop)
 
@@ -152,7 +152,7 @@ def main():
     smgr = state.StateManager(args.db_file)
 
     browser = Browser(smgr, datasets_sub, args.browse_root,
-                      args.server, args.port)
+                      args.server, args.port, loop=loop)
     widget_log_handler.callback = browser.log.model.append
 
     if os.name == "nt":
diff --git a/artiq/frontend/artiq_dashboard.py b/artiq/frontend/artiq_dashboard.py
index 31060fb81..e6ea6ad0a 100755
--- a/artiq/frontend/artiq_dashboard.py
+++ b/artiq/frontend/artiq_dashboard.py
@@ -195,7 +195,8 @@ def main():
                                                "server": args.server,
                                                "port_notify": args.port_notify,
                                                "port_control": args.port_control,
-                                           })
+                                           },
+                                           loop=loop)
     atexit_register_coroutine(d_applets.stop, loop=loop)
     smgr.register(d_applets)
     broadcast_clients["ccb"].notify_cbs.append(d_applets.ccb_notify)
diff --git a/artiq/gui/applets.py b/artiq/gui/applets.py
index b119d14bc..275e61ba7 100644
--- a/artiq/gui/applets.py
+++ b/artiq/gui/applets.py
@@ -92,9 +92,9 @@ class AppletIPCServer(AsyncioParentComm):
         finally:
             self.datasets_sub.notify_cbs.remove(self._on_mod)
 
-    def start_server(self, embed_cb, fix_initial_size_cb):
+    def start_server(self, embed_cb, fix_initial_size_cb, *, loop=None):
         self.server_task = asyncio.ensure_future(
-            self.serve(embed_cb, fix_initial_size_cb))
+            self.serve(embed_cb, fix_initial_size_cb), loop=loop)
 
     async def stop_server(self):
         if hasattr(self, "server_task"):
@@ -327,7 +327,7 @@ class _CompleterDelegate(QtWidgets.QStyledItemDelegate):
 
 
 class AppletsDock(QtWidgets.QDockWidget):
-    def __init__(self, main_window, datasets_sub, extra_substitutes={}):
+    def __init__(self, main_window, datasets_sub, extra_substitutes={}, *, loop=None):
         """
         :param extra_substitutes: Map of extra ``${strings}`` to substitute in applet
             commands to their respective values.
@@ -342,6 +342,8 @@ class AppletsDock(QtWidgets.QDockWidget):
         self.extra_substitutes = extra_substitutes
         self.applet_uids = set()
 
+        self._loop = loop
+
         self.table = QtWidgets.QTreeWidget()
         self.table.setColumnCount(2)
         self.table.setHeaderLabels(["Name", "Command"])
@@ -441,7 +443,7 @@ class AppletsDock(QtWidgets.QDockWidget):
         dock = _AppletDock(self.datasets_sub, item.applet_uid, name, spec, self.extra_substitutes)
         self.main_window.addDockWidget(QtCore.Qt.RightDockWidgetArea, dock)
         dock.setFloating(True)
-        asyncio.ensure_future(dock.start())
+        asyncio.ensure_future(dock.start(), loop=self._loop)
         dock.sigClosed.connect(partial(self.on_dock_closed, item, dock))
         return dock
 
@@ -480,7 +482,7 @@ class AppletsDock(QtWidgets.QDockWidget):
 
     def on_dock_closed(self, item, dock):
         item.applet_geometry = dock.saveGeometry()
-        asyncio.ensure_future(dock.terminate())
+        asyncio.ensure_future(dock.terminate(), loop=self._loop)
         item.setCheckState(0, QtCore.Qt.Unchecked)
 
     def get_untitled(self):
@@ -569,7 +571,7 @@ class AppletsDock(QtWidgets.QDockWidget):
                 if wi.ty == "applet":
                     dock = wi.applet_dock
                     if dock is not None:
-                        asyncio.ensure_future(dock.restart())
+                        asyncio.ensure_future(dock.restart(), loop=self._loop)
                 elif wi.ty == "group":
                     for i in range(wi.childCount()):
                         walk(wi.child(i))