Added max-restarts option. (#190)

Author: s3rius

poetry.lock

@@ -8,7 +8,7 @@
 from taskiq.abc.cmd import TaskiqCMD
 
 
-def main() -> None:  # noqa: WPS210  # pragma: no cover
+def main() -> None:  # noqa: C901, WPS210  # pragma: no cover
     """
     Main entrypoint of the taskiq.
 
@@ -70,7 +70,9 @@ def main() -> None:  # noqa: WPS210  # pragma: no cover
 
     command = subcommands[args.subcommand]
     sys.argv.pop(0)
-    command.exec(sys.argv[1:])
+    status = command.exec(sys.argv[1:])
+    if status is not None:
+        exit(status)  # noqa: WPS421
 
 
 if __name__ == "__main__":  # pragma: no cover

taskiq/__main__.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Sequence
+from typing import Optional, Sequence
 
 
 class TaskiqCMD(ABC):  # pragma: no cover
@@ -8,7 +8,7 @@ class TaskiqCMD(ABC):  # pragma: no cover
     short_help = ""
 
     @abstractmethod
-    def exec(self, args: Sequence[str]) -> None:
+    def exec(self, args: Sequence[str]) -> Optional[int]:
         """
         Execute the command.
 

taskiq/abc/cmd.py

@@ -40,6 +40,7 @@ class WorkerArgs:
     receiver_arg: List[Tuple[str, str]] = field(default_factory=list)
     max_prefetch: int = 0
     no_propagate_errors: bool = False
+    max_fails: int = -1
 
     @classmethod
     def from_cli(  # noqa: WPS213
@@ -179,6 +180,13 @@ def from_cli(  # noqa: WPS213
             dest="configure_logging",
             help="Use this parameter if your application configures custom logging.",
         )
+        parser.add_argument(
+            "--max-fails",
+            type=int,
+            dest="max_fails",
+            default=-1,
+            help="Maximum number of child process exits.",
+        )
 
         namespace = parser.parse_args(args)
         return WorkerArgs(**namespace.__dict__)

taskiq/cli/worker/args.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Sequence
+from typing import Optional, Sequence
 
 from taskiq.abc.cmd import TaskiqCMD
 from taskiq.cli.worker.args import WorkerArgs
@@ -13,14 +13,15 @@ class WorkerCMD(TaskiqCMD):
 
     short_help = "Helper to run workers"
 
-    def exec(self, args: Sequence[str]) -> None:
+    def exec(self, args: Sequence[str]) -> Optional[int]:
         """
         Start worker process.
 
         Worker process creates several small
         processes in which tasks are actually processed.
 
         :param args: CLI arguments.
+        :returns: status code.
         """
         wargs = WorkerArgs.from_cli(args)
-        run_worker(wargs)
+        return run_worker(wargs)

taskiq/cli/worker/cmd.py

@@ -43,14 +43,15 @@ def handle(
         :param action_queue: queue to send events to.
         """
         for worker_id in range(workers_num):
-            action_queue.put(ReloadOneAction(worker_num=worker_id))
+            action_queue.put(ReloadOneAction(worker_num=worker_id, is_reload_all=True))
 
 
 @dataclass
 class ReloadOneAction(ProcessActionBase):
     """This action reloads single worker with particular id."""
 
     worker_num: int
+    is_reload_all: bool
 
     def handle(
         self,
@@ -153,6 +154,7 @@ def __init__(
         args: WorkerArgs,
         worker_function: Callable[[WorkerArgs, EventType], None],
         observer: Optional[Observer] = None,  # type: ignore[valid-type]
+        max_restarts: Optional[int] = None,
     ) -> None:
         self.worker_function = worker_function
         self.action_queue: "Queue[ProcessActionBase]" = Queue(-1)
@@ -198,7 +200,7 @@ def prepare_workers(self) -> None:
         for worker, event in zip(self.workers, events):
             _wait_for_worker_startup(worker, event)
 
-    def start(self) -> None:  # noqa: C901, WPS213
+    def start(self) -> Optional[int]:  # noqa: C901, WPS213
         """
         Start managing child processes.
 
@@ -223,7 +225,10 @@ def start(self) -> None:  # noqa: C901, WPS213
         After all events are handled, it iterates over all child processes and
         checks that all processes are healthy. If process was terminated for
         some reason, it schedules a restart for dead process.
+
+        :returns: status code or None.
         """
+        restarts = 0
         self.prepare_workers()
         while True:
             sleep(1)
@@ -238,16 +243,30 @@ def start(self) -> None:  # noqa: C901, WPS213
                         action_queue=self.action_queue,
                     )
                 elif isinstance(action, ReloadOneAction):
+                    # We check if max_fails is set.
+                    # If it's true, we check how many times
+                    # worker was reloaded.
+                    if not action.is_reload_all and self.args.max_fails >= 1:
+                        restarts += 1
+                        if restarts >= self.args.max_fails:
+                            logger.warning("Max restarts reached. Exiting.")
+                            # Returning error status.
+                            return -1
                     # If we just reloaded this worker, skip handling.
                     if action.worker_num in reloaded_workers:
                         continue
                     action.handle(self.workers, self.args, self.worker_function)
                     reloaded_workers.add(action.worker_num)
                 elif isinstance(action, ShutdownAction):
                     logger.debug("Process manager closed.")
-                    return
+                    return None
 
             for worker_num, worker in enumerate(self.workers):
                 if not worker.is_alive():
                     logger.info(f"{worker.name} is dead. Scheduling reload.")
-                    self.action_queue.put(ReloadOneAction(worker_num=worker_num))
+                    self.action_queue.put(
+                        ReloadOneAction(
+                            worker_num=worker_num,
+                            is_reload_all=False,
+                        ),
+                    )

taskiq/cli/worker/process_manager.py

@@ -5,7 +5,7 @@
 from multiprocessing import set_start_method
 from multiprocessing.synchronize import Event
 from sys import platform
-from typing import Any, Type
+from typing import Any, Optional, Type
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.cli.utils import import_object, import_tasks
@@ -149,7 +149,7 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
         loop.run_until_complete(shutdown_broker(broker, args.shutdown_timeout))
 
 
-def run_worker(args: WorkerArgs) -> None:  # noqa: WPS213
+def run_worker(args: WorkerArgs) -> Optional[int]:  # noqa: WPS213
     """
     This function starts worker processes.
 
@@ -159,6 +159,7 @@ def run_worker(args: WorkerArgs) -> None:  # noqa: WPS213
     :param args: CLI arguments.
 
     :raises ValueError: if reload flag is used, but dependencies are not installed.
+    :returns: Optional status code.
     """
     if platform == "darwin":
         set_start_method("fork")
@@ -187,10 +188,11 @@ def run_worker(args: WorkerArgs) -> None:  # noqa: WPS213
 
     manager = ProcessManager(args=args, observer=observer, worker_function=start_listen)
 
-    manager.start()
+    status = manager.start()
 
     if observer is not None and observer.is_alive():
         if args.reload:
             logger.info("Stopping watching files.")
         observer.stop()
-    logger.info("Stopping logging thread.")
+
+    return status