refactor asyncio loop

Author: ddkohler

yaqd-core/yaqd_core/_is_daemon.py

@@ -123,17 +123,6 @@ def _traits(cls) -> List[str]:
 
     @classmethod
     def main(cls):
-        """Run the event loop."""
-        loop = asyncio.get_event_loop()
-        if sys.platform.startswith("win"):
-            signals = ()
-        else:
-            signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
-        for s in signals:
-            loop.add_signal_handler(
-                s, lambda s=s: asyncio.create_task(cls.shutdown_all(s, loop))
-            )
-
         parser = argparse.ArgumentParser()
         parser.add_argument(
             "--config",
@@ -199,19 +188,28 @@ def main(cls):
         with open(config_filepath, "rb") as f:
             config_file = tomli.load(f)
 
-        loop.create_task(cls._main(config_filepath, config_file, args))
+        # Run the event loop
         try:
-            loop.run_forever()
+            asyncio.run(
+                cls._main(config_filepath, config_file, args)
+            )
         except asyncio.CancelledError:
             pass
-        finally:
-            loop.close()
 
     @classmethod
     async def _main(cls, config_filepath, config_file, args=None):
         """Parse command line arguments, start event loop tasks."""
         loop = asyncio.get_running_loop()
-        cls.__servers = []
+        if sys.platform.startswith("win"):
+            signals = ()
+        else:
+            signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
+        for s in signals:
+            loop.add_signal_handler(
+                s, lambda s=s: asyncio.create_task(cls.shutdown_all(s, loop))
+            )
+
+        cls.__servers = set()
         for section in config_file:
             if section == "shared-settings":
                 continue
@@ -225,7 +223,7 @@ async def _main(cls, config_filepath, config_file, args=None):
 
         while cls.__servers:
             awaiting = cls.__servers
-            cls.__servers = []
+            cls.__servers = set()
             await asyncio.wait(awaiting)
             await asyncio.sleep(1)
         loop.stop()
@@ -252,7 +250,11 @@ def server(daemon):
             server(daemon), config.get("host", ""), config.get("port", None)
         )
         daemon._server = ser
-        cls.__servers.append(asyncio.create_task(ser.serve_forever()))
+        # cls.__servers.add(asyncio.create_task())
+        task = asyncio.create_task(ser.serve_forever())
+        cls.__servers.add(task)
+        # Add a done callback to remove the task from the set when it completes
+        task.add_done_callback(cls.__servers.discard)
 
     @classmethod
     def _parse_config(cls, config_file, section, args=None):
@@ -300,13 +302,16 @@ async def shutdown_all(cls, sig, loop):
         tasks = [
             t for t in asyncio.all_tasks() 
             if (
-                t is not asyncio.current_task() 
-                and "serve_forever" not in t.get_coro().__repr__()
+                t is not asyncio.current_task()
+                # and "serve_forever" not in t.get_coro().__repr__()
             )
         ]
+        logger.info("gathering")
         for task in tasks:
             logger.info(task.get_coro())
         await asyncio.gather(*tasks, return_exceptions=True)
+        logger.info("gathered")
+        logger.info([task.get_coro() for task in asyncio.all_tasks()])
         [d._save_state() for d in cls._daemons]
         if hasattr(signal, "SIGHUP") and sig == signal.SIGHUP:
             config_filepath = [d._config_filepath for d in cls._daemons][0]

yaqd-core/yaqd_core/_protocol.py

@@ -28,7 +28,7 @@ def connection_made(self, transport):
         self.transport = transport
         self.unpacker = avrorpc.Unpacker(self._avro_protocol)
         self._daemon._connection_made(peername)
-        self.task = asyncio.get_event_loop().create_task(self.process_requests())
+        self.task = asyncio.get_running_loop().create_task(self.process_requests())
 
     def data_received(self, data):
         """Process an incomming request."""