Added automatic worker restarts.

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

taskiq/__main__.py

@@ -1,13 +1,11 @@
-import asyncio
-
 from taskiq.cli.args import TaskiqArgs
 from taskiq.cli.worker import run_worker
 
 
 def main() -> None:
     """Main entrypoint for CLI."""
     args = TaskiqArgs.from_cli()
-    asyncio.run(run_worker(args))
+    run_worker(args)
 
 
 if __name__ == "__main__":

taskiq/cli/args.py

@@ -1,6 +1,7 @@
 import enum
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
 from dataclasses import dataclass
+from typing import List, Optional
 
 
 class LogLevel(str, enum.Enum):  # noqa: WPS600
@@ -24,13 +25,15 @@ class TaskiqArgs:
     log_level: str
     workers: int
     log_collector_format: str
+    max_threadpool_threads: int
     no_parse: bool
 
     @classmethod
-    def from_cli(cls) -> "TaskiqArgs":
+    def from_cli(cls, args: Optional[List[str]] = None) -> "TaskiqArgs":  # noqa: WPS213
         """
         Construct TaskiqArgs instanc from CLI arguments.
 
+        :param args: list of args as for cli.
         :return: TaskiqArgs instance.
         """
         parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
@@ -98,6 +101,14 @@ def from_cli(cls) -> "TaskiqArgs":
                 " with pydantic."
             ),
         )
+        parser.add_argument(
+            "--max-threadpool-threads",
+            type=int,
+            help="Maximum number of threads for executing sync functions.",
+        )
 
-        namespace = parser.parse_args()
+        if args is None:
+            namespace = parser.parse_args(args)
+        else:
+            namespace = parser.parse_args()
         return TaskiqArgs(**namespace.__dict__)

taskiq/cli/task_runner.py renamed to taskiq/cli/async_task_runner.py

@@ -1,9 +1,12 @@
 import asyncio
 import inspect
 import io
+import signal
+import sys
+from concurrent.futures import Executor, ThreadPoolExecutor
 from logging import getLogger
 from time import time
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, NoReturn, Optional
 
 from pydantic import parse_obj_as
 
@@ -99,6 +102,7 @@ async def run_task(  # noqa: WPS210
     signature: Optional[inspect.Signature],
     message: TaskiqMessage,
     cli_args: TaskiqArgs,
+    executor: Optional[Executor] = None,
 ) -> TaskiqResult[Any]:
     """
     This function actually executes functions.
@@ -118,6 +122,7 @@ async def run_task(  # noqa: WPS210
     :param signature: signature of an original function.
     :param message: received message.
     :param cli_args: CLI arguments for worker.
+    :param executor: executor to run sync tasks.
     :return: result of execution.
     """
     loop = asyncio.get_running_loop()
@@ -133,7 +138,7 @@ async def run_task(  # noqa: WPS210
                 returned = await target(*message.args, **message.kwargs)
             else:
                 returned = await loop.run_in_executor(
-                    None,
+                    executor,
                     run_sync,
                     target,
                     message,
@@ -156,7 +161,64 @@ async def run_task(  # noqa: WPS210
     )
 
 
-async def async_listen_messages(  # noqa: C901, WPS210
+def exit_process(task: asyncio.Task[Any]) -> NoReturn:
+    """
+    This function exits from the current process.
+
+    It receives asyncio Task of broker.shutdown().
+    We check if there were an exception or returned value.
+
+    If the function raised an exception, we print it with stack trace.
+    If it returned a value, we log it.
+
+    After this, we cancel all current tasks in the loop
+    and exits.
+
+    :param task: broker.shutdown task.
+    """
+    exitcode = 0
+    try:
+        result = task.result()
+        if result is not None:
+            logger.info("Broker returned value on shutdown: '%s'" % str(result))
+    except Exception as exc:
+        logger.warning("Exception was found while shutting down!")
+        logger.warning(exc, exc_info=True)
+        exitcode = 1
+
+    loop = asyncio.get_event_loop()
+    for running_task in asyncio.all_tasks(loop):
+        running_task.cancel()
+
+    logger.info("Killing worker process.")
+    sys.exit(exitcode)
+
+
+def signal_handler(broker: AsyncBroker) -> None:
+    """
+    Exit signal handler.
+
+    This signal handler
+    calls _close_broker and after
+    the task is done it exits.
+
+    :param broker: current broker.
+    """
+    if getattr(broker, "_is_shutting_down", False):
+        # We're already shutting down the broker.
+        return
+
+    # We set this flag to not call this method twice.
+    # Since we add an asynchronous task in loop
+    # It can wait for execution for some time.
+    # We want to execute shutdown only once. Otherwise
+    # it would give us Undefined Behaviour.
+    broker._is_shutting_down = True  # type: ignore  # noqa: WPS437
+    task = asyncio.create_task(broker.shutdown())
+    task.add_done_callback(exit_process)
+
+
+async def async_listen_messages(  # noqa: C901, WPS210, WPS213
     broker: AsyncBroker,
     cli_args: TaskiqArgs,
 ) -> None:
@@ -169,8 +231,23 @@ async def async_listen_messages(  # noqa: C901, WPS210
     :param broker: broker to listen to.
     :param cli_args: CLI arguments for worker.
     """
+    loop = asyncio.get_event_loop()
+    loop.add_signal_handler(
+        signal.SIGTERM,
+        signal_handler,
+        broker,
+    )
+    loop.add_signal_handler(
+        signal.SIGINT,
+        signal_handler,
+        broker,
+    )
+
     logger.info("Runing startup event.")
     await broker.startup()
+    executor = ThreadPoolExecutor(
+        max_workers=cli_args.max_threadpool_threads,
+    )
     logger.info("Listening started.")
     task_registry: Dict[str, Callable[..., Any]] = {}
     task_signatures: Dict[str, inspect.Signature] = {}
@@ -197,6 +274,7 @@ async def async_listen_messages(  # noqa: C901, WPS210
             task_signatures.get(message.task_name),
             message,
             cli_args,
+            executor,
         )
         try:
             await broker.result_backend.set_result(message.task_id, result)

taskiq/cli/worker.py

@@ -1,13 +1,15 @@
 import asyncio
+import signal
 from importlib import import_module
 from logging import basicConfig, getLevelName, getLogger
 from multiprocessing import Process
 from pathlib import Path
-from typing import Any
+from time import sleep
+from typing import Any, List
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.cli.args import TaskiqArgs
-from taskiq.cli.task_runner import async_listen_messages
+from taskiq.cli.async_task_runner import async_listen_messages
 
 logger = getLogger("taskiq.worker")
 
@@ -86,7 +88,7 @@ def start_listen(args: TaskiqArgs) -> None:
         raise ValueError("Unknown broker type. Please use AsyncBroker instance.")
 
 
-async def run_worker(args: TaskiqArgs) -> None:
+def run_worker(args: TaskiqArgs) -> None:  # noqa: C901, WPS210, WPS213
     """
     This function starts worker processes.
 
@@ -100,22 +102,68 @@ async def run_worker(args: TaskiqArgs) -> None:
         format=("[%(asctime)s][%(levelname)-7s][%(processName)s] %(message)s"),
     )
     logger.info("Starting %s worker processes." % args.workers)
-    worker_processes = []
-    for worker_num in range(args.workers):
+    worker_processes: List[Process] = []
+    for process in range(args.workers):
         work_proc = Process(
             target=start_listen,
             kwargs={"args": args},
-            name=f"worker-{worker_num}",
+            name=f"worker-{process}",
         )
         work_proc.start()
         logger.debug(
             "Started process worker-%d with pid %s "
             % (
-                worker_num,
+                process,
                 work_proc.pid,
             ),
         )
         worker_processes.append(work_proc)
 
-    for wp in worker_processes:
-        wp.join()
+    # This flag signalizes that we do need to restart processes.
+    do_restarts = True
+
+    def signal_handler(_signal: int, _frame: Any) -> None:
+        """
+        This handler is used only by main process.
+
+        If the OS sent you SIGINT or SIGTERM,
+        we should kill all spawned processes.
+
+        :param _signal: incoming signal.
+        :param _frame: current execution frame.
+        """
+        nonlocal do_restarts  # noqa: WPS420
+        nonlocal worker_processes  # noqa: WPS420
+
+        do_restarts = False  # noqa: WPS442
+        for process in worker_processes:  # noqa: WPS442
+            # This is how we send SIGTERM to child
+            # processes.
+            process.terminate()
+            process.join()
+
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+
+    while worker_processes and do_restarts:
+        # List of processes to remove.
+        sleep(1)
+        process_to_remove = []
+        for worker_id, worker in enumerate(worker_processes):
+            if worker.is_alive():
+                continue
+            if worker.exitcode is not None and worker.exitcode > 0 and do_restarts:
+                logger.info("Trying to restart the worker-%s" % worker_id)
+                worker_processes[worker_id] = Process(
+                    target=start_listen,
+                    kwargs={"args": args},
+                    name=f"worker-{worker_id}",
+                )
+                worker_processes[worker_id].start()
+            else:
+                logger.info("Worker-%s has finished." % worker_id)
+                worker.join()
+                process_to_remove.append(worker)
+
+        for dead_process in process_to_remove:
+            worker_processes.remove(dead_process)