fix: SIGINT worker processes when process_manager starts (#362)

Author: sminnee

taskiq/cli/worker/process_manager.py

@@ -1,4 +1,5 @@
 import logging
+import os
 import signal
 import sys
 from contextlib import suppress
@@ -265,7 +266,10 @@ def start(self) -> Optional[int]:  # noqa: C901
                     action.handle(self.workers, self.args, self.worker_function)
                     reloaded_workers.add(action.worker_num)
                 elif isinstance(action, ShutdownAction):
-                    logger.debug("Process manager closed.")
+                    logger.debug("Process manager closed, killing workers.")
+                    for worker in self.workers:
+                        if worker.pid:
+                            os.kill(worker.pid, signal.SIGINT)
                     return None
 
             for worker_num, worker in enumerate(self.workers):

taskiq/receiver/receiver.py

@@ -1,5 +1,6 @@
 import asyncio
 import inspect
+import signal
 from concurrent.futures import Executor
 from logging import getLogger
 from time import time
@@ -334,6 +335,12 @@ async def listen(self) -> None:  # pragma: no cover
             gr.start_soon(self.prefetcher, queue)
             gr.start_soon(self.runner, queue)
 
+            # Propagate cancellation to the prefetcher & runner
+            def _cancel(*_: Any) -> None:
+                gr.cancel_scope.cancel()
+
+            signal.signal(signal.SIGINT, _cancel)
+
         if self.on_exit is not None:
             self.on_exit(self)
 
@@ -361,9 +368,7 @@ async def prefetcher(
                 message = await iterator.__anext__()
                 fetched_tasks += 1
                 await queue.put(message)
-            except asyncio.CancelledError:
-                break
-            except StopAsyncIteration:
+            except (asyncio.CancelledError, StopAsyncIteration):
                 break
 
         await queue.put(QUEUE_DONE)
@@ -394,31 +399,35 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
                 self.sem.release()
 
         while True:
-            # Waits for semaphore to be released.
-            if self.sem is not None:
-                await self.sem.acquire()
-
-            self.sem_prefetch.release()
-            message = await queue.get()
-            if message is QUEUE_DONE:
-                # asyncio.wait will throw an error if there is nothing to wait for
-                if tasks:
-                    logger.info("Waiting for running tasks to complete.")
-                    await asyncio.wait(tasks, timeout=self.wait_tasks_timeout)
-                break
+            try:
+                # Waits for semaphore to be released.
+                if self.sem is not None:
+                    await self.sem.acquire()
+
+                self.sem_prefetch.release()
+                message = await queue.get()
+                if message is QUEUE_DONE:
+                    # asyncio.wait will throw an error if there is nothing to wait for
+                    if tasks:
+                        logger.info("Waiting for running tasks to complete.")
+                        await asyncio.wait(tasks, timeout=self.wait_tasks_timeout)
+                    break
 
-            task = asyncio.create_task(
-                self.callback(message=message, raise_err=False),
-            )
-            tasks.add(task)
-
-            # We want the task to remove itself from the set when it's done.
-            #
-            # Because if we won't save it anywhere,
-            # python's GC can silently cancel task
-            # and this behaviour considered to be a Hisenbug.
-            # https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/
-            task.add_done_callback(task_cb)
+                task = asyncio.create_task(
+                    self.callback(message=message, raise_err=False),
+                )
+                tasks.add(task)
+
+                # We want the task to remove itself from the set when it's done.
+                #
+                # Because if we won't save it anywhere,
+                # python's GC can silently cancel task
+                # and this behaviour considered to be a Hisenbug.
+                # https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/
+                task.add_done_callback(task_cb)
+
+            except asyncio.CancelledError:
+                break
 
     def _prepare_task(self, name: str, handler: Callable[..., Any]) -> None:
         """