fix(core): Let sanic manage the qworker process

coco/core.py

@@ -12,7 +12,7 @@
 import os
 import socket
 import time
-from multiprocessing import Process, set_start_method
+from multiprocessing import set_start_method
 from pathlib import Path
 
 import click
@@ -83,12 +83,12 @@ def __init__(
         if full_reset:
             reset = False
 
-        # In case constructor crashes before this gets assigned, so that destructor
-        # doesn't fail.
-        self.qworker = None
         self.state = None
         self.redis_sync = None
 
+        # Set to True when a coco_shutdown is happening
+        self.coco_shutdown = False
+
         # Load the config
         self._load_config(conf, testing)
 
@@ -180,52 +180,12 @@ def __init__(
         else:
             sock = None
 
-        # Start the worker process
-        self.qworker = Process(
-            target=worker.main_loop,
-            args=(
-                self.endpoints,
-                self.forwarder,
-                self.config["port"],
-                self.config["metrics_port"],
-                int(self.config["redis_port"]),
-                self.frontend_timeout,
-            ),
-        )
-        self.qworker.daemon = True
-        try:
-            self.qworker.start()
-        except RuntimeError:
-            self.qworker.join()
-
-        self._call_endpoints_on_start()
-
         # This blocks until cocod terminates
         self._start_server(sock)
 
+        # Coco is done, delete the redis async connection
         self.redis_async = None
 
-    def __del__(self):
-        """
-        Destruct :class:`Core`.
-
-        Join the worker process.
-        """
-        if self.redis_sync:
-            logger.info("Joining worker process...")
-            try:
-                self.redis_sync.rpush("queue", "coco_shutdown")
-            except RuntimeError as e:
-                logger.error(
-                    "Failed sending shutdown command to worker "
-                    f"(have to kill it): {type(e)}: {e}"
-                )
-            self._kill_worker()
-
-    def _kill_worker(self):
-        if self.qworker:
-            self.qworker.kill()
-
     def _call_endpoints_on_start(self):
         for endpoint in self.endpoints.values():
             # Initialise request counter
@@ -265,6 +225,41 @@ def _start_server(self, sock: socket.socket | None = None):
         self.sanic_app.config.REQUEST_TIMEOUT = self.frontend_timeout
         self.sanic_app.config.RESPONSE_TIMEOUT = self.frontend_timeout
 
+        def start_qworker(app):
+            """Start the qworker."""
+            app.manager.manage(
+                "qworker",
+                worker.main_loop,
+                {
+                    "app": app,
+                    "endpoints": self.endpoints,
+                    "forwarder": self.forwarder,
+                    "coco_port": self.config["port"],
+                    "metrics_port": self.config["metrics_port"],
+                    "redis_port": int(self.config["redis_port"]),
+                    "frontend_timeout": self.frontend_timeout,
+                },
+                transient=False,
+                restartable=True,
+                auto_start=True,
+            )
+
+            self._call_endpoints_on_start()
+
+        def signal_coco_shutdown(app):
+            """Tell the qworker to shutdown via redis."""
+            if self.redis_sync and not self.coco_shutdown:
+                self.coco_shutdown = True
+                logger.info("Signalling coco_shutdown...")
+                try:
+                    self.redis_sync.rpush("queue", "coco_shutdown")
+                except RuntimeError as e:
+                    logger.error(f"queueing coco_shutdown in redis failed: {e}")
+
+        # Get Sanic to start/stop the qworker process when it starts/terminates
+        self.sanic_app.register_listener(start_qworker, "main_process_ready")
+        self.sanic_app.register_listener(signal_coco_shutdown, "before_server_stop")
+
         # Create the Redis connection pool, use sanic to start it so that it
         # ends up in the same event loop
         async def init_redis_async(*_):

coco/worker.py

@@ -55,7 +55,7 @@ async def _open_redis_connection(redis_port):
 
 
 def main_loop(
-    endpoints, forwarder, coco_port, metrics_port, redis_port, frontend_timeout
+    app, endpoints, forwarder, coco_port, metrics_port, redis_port, frontend_timeout
 ):
     """
     Wait for tasks and run them.
@@ -64,6 +64,8 @@ def main_loop(
 
     Parameters
     ----------
+    app : Sanic.app
+        The controlling Sanic app.
     endpoints : dict
         A dict with keys being endpoint names and values being of type
         :class:`Endpoint`.
@@ -212,16 +214,25 @@ async def go():
         # optionally close connection
         await conn.close()
 
-    # NOTE: need to create a new event loop here otherwise macOS seems to have
-    # issues involving the asyncio event loop and the Process fork
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-
-    # Start up slack logging for the worker
-    slack.start(loop)
-
-    scheduler = Scheduler(endpoints, "127.0.0.1", coco_port, frontend_timeout)
-    loop.run_until_complete(asyncio.gather(go(), scheduler.start()))
-
-    # Cleanup
-    loop.run_until_complete(slack.stop())
+    try:
+        # NOTE: need to create a new event loop here otherwise macOS seems to have
+        # issues involving the asyncio event loop and the Process fork
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        # Start up slack logging for the worker
+        slack.start(loop)
+
+        scheduler = Scheduler(endpoints, "127.0.0.1", coco_port, frontend_timeout)
+        loop.run_until_complete(asyncio.gather(go(), scheduler.start()))
+
+        # Cleanup
+        loop.run_until_complete(slack.stop())
+    except KeyboardInterrupt:
+        # Normal termination
+        logger.info("qworker shutdown")
+        pass
+    except (RuntimeError, OSError) as e:
+        # Shutdown Sanic on error
+        logger.error(f"qworker encountered an error: {e}")
+        app.manager.terminate()