Merge branch 'release/0.8.7'

Author: s3rius

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "taskiq"
-version = "0.8.6"
+version = "0.8.7"
 description = "Distributed task queue with full async support"
 authors = ["Pavel Kirilin <[email protected]>"]
 maintainers = ["Pavel Kirilin <[email protected]>"]

pyproject.toml

@@ -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

@@ -26,7 +26,7 @@ async def schedules_updater(
     :param scheduler: current scheduler.
     :param current_schedules: list of schedules.
     """
-    while True:
+    while True:  # noqa: WPS457
         logger.debug("Started schedule update.")
         new_schedules: "List[ScheduledTask]" = []
         for source in scheduler.sources:
@@ -40,8 +40,7 @@ async def schedules_updater(
                 logger.debug(exc, exc_info=True)
                 continue
 
-            for schedule in scheduler.merge_func(new_schedules, schedules):
-                new_schedules.append(schedule)
+            new_schedules = scheduler.merge_func(new_schedules, schedules)
 
         current_schedules.clear()
         current_schedules.extend(new_schedules)
@@ -62,38 +61,15 @@ def should_run(task: ScheduledTask) -> bool:
     return False
 
 
-async def run_scheduler(args: SchedulerArgs) -> None:  # noqa: C901, WPS210, WPS213
+async def _run_loop(scheduler: TaskiqScheduler) -> None:
     """
     Runs scheduler loop.
 
     This function imports taskiq scheduler
     and runs tasks when needed.
 
-    :param args: parsed CLI args.
+    :param scheduler: current scheduler.
     """
-    if isinstance(args.scheduler, str):
-        scheduler = import_object(args.scheduler)
-    else:
-        scheduler = args.scheduler
-    if not isinstance(scheduler, TaskiqScheduler):
-        print(  # noqa: WPS421
-            "Imported scheduler is not a subclass of TaskiqScheduler.",
-        )
-        exit(1)  # noqa: WPS421
-    scheduler.broker.is_scheduler_process = True
-    import_tasks(args.modules, args.tasks_pattern, args.fs_discover)
-    if args.configure_logging:
-        basicConfig(
-            level=getLevelName(args.log_level),
-            format=(
-                "[%(asctime)s][%(levelname)-7s]"
-                "[%(module)s:%(funcName)s:%(lineno)d]"
-                " %(message)s"
-            ),
-        )
-    getLogger("taskiq").setLevel(level=getLevelName(args.log_level))
-    for source in scheduler.sources:
-        await source.startup()
     loop = asyncio.get_event_loop()
     tasks: "List[ScheduledTask]" = []
     loop.create_task(schedules_updater(scheduler, tasks))
@@ -121,3 +97,44 @@ async def run_scheduler(args: SchedulerArgs) -> None:  # noqa: C901, WPS210, WPS
             - datetime.now()
         )
         await asyncio.sleep(delay.total_seconds())
+
+
+async def run_scheduler(args: SchedulerArgs) -> None:  # noqa: WPS213
+    """
+    Run scheduler.
+
+    This function takes all CLI arguments
+    and starts the scheduler process.
+
+    :param args: parsed CLI arguments.
+    """
+    if isinstance(args.scheduler, str):
+        scheduler = import_object(args.scheduler)
+    else:
+        scheduler = args.scheduler
+    if not isinstance(scheduler, TaskiqScheduler):
+        print(  # noqa: WPS421
+            "Imported scheduler is not a subclass of TaskiqScheduler.",
+        )
+        exit(1)  # noqa: WPS421
+    scheduler.broker.is_scheduler_process = True
+    import_tasks(args.modules, args.tasks_pattern, args.fs_discover)
+    if args.configure_logging:
+        basicConfig(
+            level=getLevelName(args.log_level),
+            format=(
+                "[%(asctime)s][%(levelname)-7s]"
+                "[%(module)s:%(funcName)s:%(lineno)d]"
+                " %(message)s"
+            ),
+        )
+    getLogger("taskiq").setLevel(level=getLevelName(args.log_level))
+    for source in scheduler.sources:
+        await source.startup()
+
+    try:
+        await _run_loop(scheduler)
+    except asyncio.CancelledError:
+        logger.warning("Shutting down scheduler.")
+        await scheduler.shutdown()
+        logger.info("Scheduler shut down. Good bye!")

taskiq/cli/scheduler/run.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

taskiq/cli/worker/run.py

@@ -73,3 +73,7 @@ async def on_ready(self, task: ScheduledTask) -> None:
             **task.kwargs,
         )
         await maybe_awaitable(task.source.post_send(task))
+
+    async def shutdown(self) -> None:
+        """Shutdown the scheduler process."""
+        await self.broker.shutdown()