Added second-based scheduler intervals. (#450)

Author: s3rius

taskiq/api/scheduler.py

@@ -1,10 +1,14 @@
+from datetime import timedelta
+from typing import Optional
+
 from taskiq.cli.scheduler.run import run_scheduler_loop
 from taskiq.scheduler.scheduler import TaskiqScheduler
 
 
 async def run_scheduler_task(
     scheduler: TaskiqScheduler,
     run_startup: bool = False,
+    interval: Optional[timedelta] = None,
 ) -> None:
     """
     Run scheduler task.
@@ -20,4 +24,4 @@ async def run_scheduler_task(
     if run_startup:
         await scheduler.startup()
     while True:
-        await run_scheduler_loop(scheduler)
+        await run_scheduler_loop(scheduler, interval)

taskiq/cli/scheduler/args.py

@@ -17,6 +17,7 @@ class SchedulerArgs:
     fs_discover: bool = False
     tasks_pattern: Sequence[str] = ("**/tasks.py",)
     skip_first_run: bool = False
+    update_interval: Optional[int] = None
 
     @classmethod
     def from_cli(cls, args: Optional[Sequence[str]] = None) -> "SchedulerArgs":
@@ -80,6 +81,15 @@ def from_cli(cls, args: Optional[Sequence[str]] = None) -> "SchedulerArgs":
                 "This option skips running tasks immediately after scheduler start."
             ),
         )
+        parser.add_argument(
+            "--update-interval",
+            type=int,
+            default=None,
+            help=(
+                "Interval in seconds to check for new tasks. "
+                "If not specified, scheduler will run once a minute."
+            ),
+        )
 
         namespace = parser.parse_args(args)
         # If there are any patterns specified, remove default.

taskiq/cli/scheduler/run.py

@@ -3,7 +3,7 @@
 import sys
 from datetime import datetime, timedelta
 from logging import basicConfig, getLevelName, getLogger
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional, Set, Tuple
 
 import pytz
 from pycron import is_now
@@ -55,7 +55,7 @@ async def get_schedules(source: ScheduleSource) -> List[ScheduledTask]:
 
 async def get_all_schedules(
     scheduler: TaskiqScheduler,
-) -> Dict[ScheduleSource, List[ScheduledTask]]:
+) -> List[Tuple[ScheduleSource, List[ScheduledTask]]]:
     """
     Task to update all schedules.
 
@@ -71,7 +71,7 @@ async def get_all_schedules(
     schedules = await asyncio.gather(
         *[get_schedules(source) for source in scheduler.sources],
     )
-    return dict(zip(scheduler.sources, schedules))
+    return list(zip(scheduler.sources, schedules))
 
 
 def get_task_delay(task: ScheduledTask) -> Optional[int]:
@@ -98,12 +98,10 @@ def get_task_delay(task: ScheduledTask) -> Optional[int]:
         task_time = to_tz_aware(task.time)
         if task_time <= now:
             return 0
-        one_min_ahead = (now + timedelta(minutes=1)).replace(second=1, microsecond=0)
-        if task_time <= one_min_ahead:
-            delay = task_time - now
-            if delay.microseconds:
-                return int(delay.total_seconds()) + 1
-            return int(delay.total_seconds())
+        delay = task_time - now
+        if delay.microseconds:
+            return int(delay.total_seconds()) + 1
+        return int(delay.total_seconds())
     return None
 
 
@@ -145,21 +143,41 @@ async def delayed_send(
     await scheduler.on_ready(source, task)
 
 
-async def run_scheduler_loop(scheduler: TaskiqScheduler) -> None:
+async def run_scheduler_loop(  # noqa: C901
+    scheduler: TaskiqScheduler,
+    interval: Optional[timedelta] = None,
+) -> None:
     """
     Runs scheduler loop.
 
     This function imports taskiq scheduler
     and runs tasks when needed.
 
     :param scheduler: current scheduler.
+    :param interval: interval to check for schedule updates.
     """
     loop = asyncio.get_event_loop()
-    running_schedules = set()
+    running_schedules: Dict[str, asyncio.Task[Any]] = {}
+    ran_cron_jobs: Set[str] = set()
+    current_minute = datetime.now(tz=pytz.UTC).minute
     while True:
-        # We use this method to correctly sleep for one minute.
+        now = datetime.now(tz=pytz.UTC)
+        # If minute changed, we need to clear
+        # ran_cron_jobs set and update current minute.
+        if now.minute != current_minute:
+            current_minute = now.minute
+            ran_cron_jobs.clear()
+        # If interval is not None, we need to
+        # calculate next run time using it.
+        if interval is not None:
+            next_run = now + interval
+        # otherwise we need assume that
+        # we will run it at the start of the next minute.
+        # as crontab does.
+        else:
+            next_run = (now + timedelta(minutes=1)).replace(second=1, microsecond=0)
         scheduled_tasks = await get_all_schedules(scheduler)
-        for source, task_list in scheduled_tasks.items():
+        for source, task_list in scheduled_tasks:
             logger.debug("Got %d schedules from source %s.", len(task_list), source)
             for task in task_list:
                 try:
@@ -172,16 +190,37 @@ async def run_scheduler_loop(scheduler: TaskiqScheduler) -> None:
                         task.schedule_id,
                     )
                     continue
-                if task_delay is not None:
-                    send_task = loop.create_task(
-                        delayed_send(scheduler, source, task, task_delay),
-                    )
-                    running_schedules.add(send_task)
-                    send_task.add_done_callback(running_schedules.discard)
-        next_minute = datetime.now().replace(second=0, microsecond=0) + timedelta(
-            minutes=1,
-        )
-        delay = next_minute - datetime.now()
+                # If task delay is None, we don't need to run it.
+                if task_delay is None:
+                    continue
+                # If task is delayed for more than next_run,
+                # we don't need to run it, because we will
+                # run it in the next iteration.
+                if now + timedelta(seconds=task_delay) >= next_run:
+                    continue
+                # If task is already running, we don't need to run it again.
+                if task.schedule_id in running_schedules and task_delay < 1:
+                    continue
+                # If task is cron job, we need to check if
+                # we already ran it this minute.
+                if task.cron is not None:
+                    if task.schedule_id in ran_cron_jobs:
+                        continue
+                    ran_cron_jobs.add(task.schedule_id)
+                send_task = loop.create_task(
+                    delayed_send(scheduler, source, task, task_delay),
+                    # We need to set the name of the task
+                    # to be able to discard its reference
+                    # after it is done.
+                    name=f"schedule_{task.schedule_id}",
+                )
+                running_schedules[task.schedule_id] = send_task
+                send_task.add_done_callback(
+                    lambda task_future: running_schedules.pop(
+                        task_future.get_name().removeprefix("schedule_"),
+                    ),
+                )
+        delay = next_run - datetime.now(tz=pytz.UTC)
         logger.debug(
             "Sleeping for %.2f seconds before getting schedules.",
             delay.total_seconds(),
@@ -226,6 +265,10 @@ async def run_scheduler(args: SchedulerArgs) -> None:
     for source in scheduler.sources:
         await source.startup()
 
+    interval = None
+    if args.update_interval:
+        interval = timedelta(seconds=args.update_interval)
+
     logger.info("Starting scheduler.")
     await scheduler.startup()
     logger.info("Startup completed.")
@@ -239,7 +282,7 @@ async def run_scheduler(args: SchedulerArgs) -> None:
         await asyncio.sleep(delay.total_seconds())
         logger.info("First run skipped. The scheduler is now running.")
     try:
-        await run_scheduler_loop(scheduler)
+        await run_scheduler_loop(scheduler, interval)
     except asyncio.CancelledError:
         logger.warning("Shutting down scheduler.")
         await scheduler.shutdown()

taskiq/schedule_sources/label_based.py

@@ -1,5 +1,6 @@
+import uuid
 from logging import getLogger
-from typing import List
+from typing import Dict, List
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.schedule_source import ScheduleSource
@@ -13,20 +14,26 @@ class LabelScheduleSource(ScheduleSource):
 
     def __init__(self, broker: AsyncBroker) -> None:
         self.broker = broker
+        self.schedules: Dict[str, ScheduledTask] = {}
 
-    async def get_schedules(self) -> List["ScheduledTask"]:
+    async def startup(self) -> None:
         """
-        Collect schedules for all tasks.
-
-        this function checks labels for all
-        tasks available to the broker.
+        Startup the schedule source.
 
+        This function iterates over all tasks
+        available to the broker and collects
+        schedules from their labels.
         If task has a schedule label,
-        it will be parsed and returned.
+        it will be parsed and added to the
+        scheduler list.
 
-        :return: list of schedules.
+        Every time schedule is added, the random
+        schedule id is generated. Please be aware that
+        they are different for every startup.
+
+        :return: None
         """
-        schedules = []
+        self.schedules.clear()
         for task_name, task in self.broker.get_all_tasks().items():
             if task.broker != self.broker:
                 # if task broker doesn't match self, something is probably wrong
@@ -40,20 +47,36 @@ async def get_schedules(self) -> List["ScheduledTask"]:
                     continue
                 labels = schedule.get("labels", {})
                 labels.update(task.labels)
-                schedules.append(
-                    ScheduledTask(
-                        task_name=task_name,
-                        labels=labels,
-                        args=schedule.get("args", []),
-                        kwargs=schedule.get("kwargs", {}),
-                        cron=schedule.get("cron"),
-                        time=schedule.get("time"),
-                        cron_offset=schedule.get("cron_offset"),
-                    ),
+                schedule_id = uuid.uuid4().hex
+
+                self.schedules[schedule_id] = ScheduledTask(
+                    task_name=task_name,
+                    labels=labels,
+                    schedule_id=schedule_id,
+                    args=schedule.get("args", []),
+                    kwargs=schedule.get("kwargs", {}),
+                    cron=schedule.get("cron"),
+                    time=schedule.get("time"),
+                    cron_offset=schedule.get("cron_offset"),
                 )
-        return schedules
 
-    def post_send(self, scheduled_task: ScheduledTask) -> None:
+        return await super().startup()
+
+    async def get_schedules(self) -> List["ScheduledTask"]:
+        """
+        Collect schedules for all tasks.
+
+        this function checks labels for all
+        tasks available to the broker.
+
+        If task has a schedule label,
+        it will be parsed and returned.
+
+        :return: list of schedules.
+        """
+        return list(self.schedules.values())
+
+    def post_send(self, task: "ScheduledTask") -> None:
         """
         Remove `time` schedule from task's scheduler list.
 
@@ -62,22 +85,7 @@ def post_send(self, scheduled_task: ScheduledTask) -> None:
 
         :param scheduled_task: task that just have sent
         """
-        if scheduled_task.cron or not scheduled_task.time:
+        if task.cron or not task.time:
             return  # it's scheduled task with cron label, do not remove this trigger.
 
-        for task_name, task in self.broker.get_all_tasks().items():
-            if task.broker != self.broker:
-                # if task broker doesn't match self, something is probably wrong
-                logger.warning(
-                    f"Broker for {task_name} `{task.broker}` doesn't "
-                    f"match scheduler's broker `{self.broker}`",
-                )
-                continue
-            if 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
+        self.schedules.pop(task.schedule_id, None)

tests/cli/scheduler/test_task_delays.py

@@ -9,7 +9,7 @@
 
 
 def test_should_run_success() -> None:
-    hour = datetime.datetime.utcnow().hour
+    hour = datetime.datetime.now(datetime.timezone.utc).hour
     delay = get_task_delay(
         ScheduledTask(
             task_name="",
@@ -97,18 +97,26 @@ def test_time_utc_with_local_zone() -> None:
     assert delay is not None and delay >= 0
 
 
+@freeze_time("2023-01-14 12:00:00")
 def test_time_localtime_without_zone() -> None:
     time = datetime.datetime.now(tz=pytz.FixedOffset(240)).replace(tzinfo=None)
+    time_to_run = time - datetime.timedelta(seconds=1)
+
     delay = get_task_delay(
         ScheduledTask(
             task_name="",
             labels={},
             args=[],
             kwargs={},
-            time=time - datetime.timedelta(seconds=1),
+            time=time_to_run,
         ),
     )
-    assert delay is None
+
+    expected_delay = time_to_run.replace(tzinfo=pytz.UTC) - datetime.datetime.now(
+        pytz.UTC,
+    )
+
+    assert delay == int(expected_delay.total_seconds())
 
 
 @freeze_time("2023-01-14 12:00:00")

tests/cli/scheduler/test_updater.py

@@ -56,10 +56,10 @@ async def test_get_schedules_success() -> None:
     schedules = await get_all_schedules(
         TaskiqScheduler(InMemoryBroker(), sources),
     )
-    assert schedules == {
-        sources[0]: schedules1,
-        sources[1]: schedules2,
-    }
+    assert schedules == [
+        (sources[0], schedules1),
+        (sources[1], schedules2),
+    ]
 
 
 @pytest.mark.anyio
@@ -81,7 +81,7 @@ async def test_get_schedules_error() -> None:
     schedules = await get_all_schedules(
         TaskiqScheduler(InMemoryBroker(), [source1, source2]),
     )
-    assert schedules == {
-        source1: source1.schedules,
-        source2: [],
-    }
+    assert schedules == [
+        (source1, source1.schedules),
+        (source2, []),
+    ]