Bug/task-scheduled-by-time-runs-multiple-times (#164)

Author: MishaVyb

.gitignore

@@ -165,3 +165,6 @@ cython_debug/
 node_modules
 .temp
 .cache
+
+# macOS
+.DS_Store

docs/examples/extending/schedule_source.py

@@ -19,10 +19,32 @@ async def get_schedules(self) -> List["ScheduledTask"]:
                 args=[],
                 kwargs={},
                 cron="* * * * *",
+                #
+                # We need point on self source for calling pre_send / post_send when
+                # task is ready to be enqueued.
+                source=self,
             ),
         ]
 
     # This method is optional. You may not implement this.
     # It's just a helper to people to be able to interact with your source.
     async def add_schedule(self, schedule: "ScheduledTask") -> None:
         return await super().add_schedule(schedule)
+
+    # This method is optional. You may not implement this.
+    # It's just a helper to people to be able to interact with your source.
+    async def pre_send(self, task: "ScheduledTask") -> None:
+        """
+        Actions to execute before task will be sent to broker.
+
+        :param task: task that will be sent
+        """
+
+    # This method is optional. You may not implement this.
+    # It's just a helper to people to be able to interact with your source.
+    async def post_send(self, task: "ScheduledTask") -> None:
+        """
+        Actions to execute after task was sent to broker.
+
+        :param task: task that just have sent
+        """

poetry.lock

@@ -59,6 +59,8 @@ pytest-xdist = { version = "^2.5.0", extras = ["psutil"] }
 types-mock = "^4.0.15"
 wemake-python-styleguide = "^0.18.0"
 tox = "^4.6.4"
+freezegun = "^1.2.2"
+pytest-mock = "^3.11.1"
 
 [tool.poetry.extras]
 zmq = ["pyzmq"]
@@ -89,6 +91,9 @@ warn_unused_ignores = false
 profile = "black"
 multi_line_output = 3
 
+[tool.pytest.ini_options]
+log_level='INFO'
+
 [tool.coverage.run]
 omit = [
     "taskiq/__main__.py",

pyproject.toml

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, Any, Coroutine, List, Union
 
 if TYPE_CHECKING:  # pragma: no cover
     from taskiq.scheduler.scheduler import ScheduledTask
@@ -33,3 +33,23 @@ async def add_schedule(self, schedule: "ScheduledTask") -> None:  # noqa: B027
 
         :param schedule: schedule to add.
         """
+
+    def pre_send(  # noqa: B027
+        self,
+        task: "ScheduledTask",
+    ) -> Union[None, Coroutine[Any, Any, None]]:
+        """
+        Actions to execute before task will be sent to broker.
+
+        :param task: task that will be sent
+        """
+
+    def post_send(  # noqa: B027
+        self,
+        task: "ScheduledTask",
+    ) -> Union[None, Coroutine[Any, Any, None]]:
+        """
+        Actions to execute after task was sent to broker.
+
+        :param task: task that just have sent
+        """

taskiq/abc/schedule_source.py

@@ -1,15 +1,16 @@
 from argparse import ZERO_OR_MORE, ArgumentDefaultsHelpFormatter, ArgumentParser
 from dataclasses import dataclass
-from typing import List, Optional, Sequence
+from typing import List, Optional, Sequence, Union
 
 from taskiq.cli.common_args import LogLevel
+from taskiq.scheduler.scheduler import TaskiqScheduler
 
 
 @dataclass
 class SchedulerArgs:
     """Arguments for scheduler."""
 
-    scheduler: str
+    scheduler: Union[str, TaskiqScheduler]
     modules: List[str]
     log_level: str = LogLevel.INFO.name
     fs_discover: bool = False

taskiq/cli/scheduler/args.py

@@ -5,9 +5,9 @@
 
 from pycron import is_now
 
+from taskiq.abc.broker import AsyncBroker
 from taskiq.cli.scheduler.args import SchedulerArgs
 from taskiq.cli.utils import import_object, import_tasks
-from taskiq.kicker import AsyncKicker
 from taskiq.scheduler.scheduler import ScheduledTask, TaskiqScheduler
 
 logger = getLogger(__name__)
@@ -40,8 +40,10 @@ async def schedules_updater(
                 )
                 logger.debug(exc, exc_info=True)
                 continue
+
             for schedule in scheduler.merge_func(new_schedules, schedules):
                 new_schedules.append(schedule)
+
         current_schedules.clear()
         current_schedules.extend(new_schedules)
         await asyncio.sleep(scheduler.refresh_delay)
@@ -55,7 +57,7 @@ def should_run(task: ScheduledTask) -> bool:
     :return: True if task must be sent.
     """
     if task.cron is not None:
-        return is_now(task.cron)
+        return is_now(task.cron, datetime.utcnow())
     if task.time is not None:
         return task.time <= datetime.utcnow()
     return False
@@ -70,6 +72,7 @@ async def run_scheduler(args: SchedulerArgs) -> None:  # noqa: C901, WPS210, WPS
 
     :param args: parsed CLI args.
     """
+    AsyncBroker.is_scheduler_process = True
     if isinstance(args.scheduler, str):
         scheduler = import_object(args.scheduler)
     else:
@@ -98,7 +101,6 @@ async def run_scheduler(args: SchedulerArgs) -> None:  # noqa: C901, WPS210, WPS
     await scheduler.startup()
     logger.info("Startup completed.")
     while True:  # noqa: WPS457
-        not_fired_tasks = []
         for task in tasks:
             try:
                 ready = should_run(task)
@@ -111,14 +113,8 @@ async def run_scheduler(args: SchedulerArgs) -> None:  # noqa: C901, WPS210, WPS
                 continue
             if ready:
                 logger.info("Sending task %s.", task.task_name)
-                loop.create_task(
-                    AsyncKicker(task.task_name, scheduler.broker, task.labels).kiq(
-                        *task.args,
-                        **task.kwargs,
-                    ),
-                )
-            else:
-                not_fired_tasks.append(task)
+                loop.create_task(scheduler.on_ready(task))
+
         delay = (
             datetime.now().replace(second=1, microsecond=0)
             + timedelta(minutes=1)

taskiq/cli/scheduler/run.py

@@ -1,9 +1,12 @@
+from logging import getLogger
 from typing import List
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.schedule_source import ScheduleSource
 from taskiq.scheduler.scheduler import ScheduledTask
 
+logger = getLogger(__name__)
+
 
 class LabelScheduleSource(ScheduleSource):
     """Schedule source based on labels."""
@@ -40,6 +43,29 @@ async def get_schedules(self) -> List["ScheduledTask"]:
                         kwargs=schedule.get("kwargs", {}),
                         cron=schedule.get("cron"),
                         time=schedule.get("time"),
+                        source=self,
                     ),
                 )
         return schedules
+
+    def post_send(self, scheduled_task: ScheduledTask) -> None:
+        """
+        Remove `time` schedule from task's scheduler list.
+
+        Task just have sent and won't be sent by that trigger anymore. Other triggers in
+        scheduler list left unchanged.
+
+        :param scheduled_task: task that just have sent
+        """
+        if scheduled_task.cron or not scheduled_task.time:
+            return  # it's scheduled task with cron label, do not remove this trigger.
+
+        for task_name, task in self.broker.available_tasks.items():
+            if task.broker != self.broker or scheduled_task.task_name != task_name:
+                continue
+
+            schedule_list = task.labels.get("schedule", []).copy()
+            for idx, schedule in enumerate(schedule_list):
+                if schedule.get("time") == scheduled_task.time:
+                    task.labels.get("schedule", []).pop(idx)
+                    return

taskiq/schedule_sources/label_based.py

@@ -3,7 +3,9 @@
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
 
 from taskiq.abc.broker import AsyncBroker
+from taskiq.kicker import AsyncKicker
 from taskiq.scheduler.merge_functions import preserve_all
+from taskiq.utils import maybe_awaitable
 
 if TYPE_CHECKING:  # pragma: no cover
     from taskiq.abc.schedule_source import ScheduleSource
@@ -17,6 +19,7 @@ class ScheduledTask:
     labels: Dict[str, Any]
     args: List[Any]
     kwargs: Dict[str, Any]
+    source: "ScheduleSource"  # Backward point to source which this task belongs to
     cron: Optional[str] = field(default=None)
     time: Optional[datetime] = field(default=None)
 
@@ -56,3 +59,17 @@ async def startup(self) -> None:  # pragma: no cover
         connections or anything you'd like.
         """
         await self.broker.startup()
+
+    async def on_ready(self, task: ScheduledTask) -> None:
+        """
+        This method is called when task is ready to be enqueued.
+
+        It's triggered on proper time depending on `task.cron` or `task.time` attribute.
+        :param task: task to send
+        """
+        await maybe_awaitable(task.source.pre_send(task))
+        await AsyncKicker(task.task_name, self.broker, task.labels).kiq(
+            *task.args,
+            **task.kwargs,
+        )
+        await maybe_awaitable(task.source.post_send(task))

taskiq/scheduler/scheduler.py

@@ -1,6 +1,8 @@
+import asyncio
 from typing import Generator
 
 import pytest
+from pytest_mock import MockerFixture
 
 from taskiq.abc.broker import AsyncBroker
 
@@ -26,3 +28,14 @@ def reset_broker() -> Generator[None, None, None]:
     """
     yield
     AsyncBroker.available_tasks = {}
+    AsyncBroker.is_worker_process = False
+    AsyncBroker.is_scheduler_process = False
+
+
+@pytest.fixture
+def mock_sleep(mocker: MockerFixture) -> None:
+    async def _fast_sleep(delay: float) -> None:
+        await asyncio_sleep(delay / 10000)
+
+    asyncio_sleep = asyncio.sleep
+    mocker.patch("asyncio.sleep", _fast_sleep)