Fixed delay for specific time values. (#200)

Author: s3rius

taskiq/cli/scheduler/run.py

@@ -1,7 +1,7 @@
 import asyncio
 from datetime import datetime, timedelta
 from logging import basicConfig, getLevelName, getLogger
-from typing import List
+from typing import List, Optional
 
 import pytz
 from pycron import is_now
@@ -64,15 +64,15 @@ async def schedules_updater(
         await asyncio.sleep(scheduler.refresh_delay)
 
 
-def should_run(task: ScheduledTask) -> bool:
+def get_task_delay(task: ScheduledTask) -> Optional[int]:  # noqa: C901
     """
-    Checks if it's time to run a task.
+    Get delay of the task in seconds.
 
     :param task: task to check.
     :return: True if task must be sent.
     """
+    now = datetime.now(tz=pytz.UTC)
     if task.cron is not None:
-        now = datetime.now(tz=pytz.UTC)
         # If user specified cron offset we apply it.
         # If it's timedelta, we simply add the delta to current time.
         if task.cron_offset and isinstance(task.cron_offset, timedelta):
@@ -81,10 +81,43 @@ def should_run(task: ScheduledTask) -> bool:
         # offset and then apply.
         elif task.cron_offset and isinstance(task.cron_offset, str):
             now = now.astimezone(pytz.timezone(task.cron_offset))
-        return is_now(task.cron, now)
-    if task.time is not None:
-        return to_tz_aware(task.time) <= datetime.now(tz=pytz.UTC)
-    return False
+        if is_now(task.cron, now):
+            return 0
+        return None
+    elif task.time is not None:
+        task_time = to_tz_aware(task.time).replace(microsecond=0)
+        if task_time <= now:
+            return 0
+        one_min_ahead = (now + timedelta(minutes=1)).replace(second=1, microsecond=0)
+        if task_time <= one_min_ahead:
+            return int((task_time - now).total_seconds())
+    return None
+
+
+async def delayed_send(
+    scheduler: TaskiqScheduler,
+    task: ScheduledTask,
+    delay: int,
+) -> None:
+    """
+    Send a task with a delay.
+
+    This function waits for some time and then
+    sends a task.
+
+    The main idea is that scheduler gathers
+    tasks every minute and some of them have
+    specfic time. To respect the time, we calculate
+    the delay and send the task after some delay.
+
+    :param scheduler: current scheduler.
+    :param task: task to send.
+    :param delay: how long to wait.
+    """
+    if delay > 0:
+        await asyncio.sleep(delay)
+    logger.info("Sending task %s.", task.task_name)
+    await scheduler.on_ready(task)
 
 
 async def _run_loop(scheduler: TaskiqScheduler) -> None:
@@ -105,17 +138,16 @@ async def _run_loop(scheduler: TaskiqScheduler) -> None:
     while True:  # noqa: WPS457
         for task in tasks:
             try:
-                ready = should_run(task)
+                task_delay = get_task_delay(task)
             except ValueError:
                 logger.warning(
                     "Cannot parse cron: %s for task: %s",
                     task.cron,
                     task.task_name,
                 )
                 continue
-            if ready:
-                logger.info("Sending task %s.", task.task_name)
-                loop.create_task(scheduler.on_ready(task))
+            if task_delay is not None:
+                loop.create_task(delayed_send(scheduler, task, task_delay))
 
         delay = (
             datetime.now().replace(second=1, microsecond=0)

tests/cli/scheduler/test_is_now.py renamed to tests/cli/scheduler/test_task_delays.py

@@ -1,9 +1,10 @@
 import datetime
 
 import pytz
+from freezegun import freeze_time
 from tzlocal import get_localzone
 
-from taskiq.cli.scheduler.run import should_run
+from taskiq.cli.scheduler.run import get_task_delay
 from taskiq.schedule_sources.label_based import LabelScheduleSource
 from taskiq.scheduler.scheduler import ScheduledTask
 
@@ -12,7 +13,7 @@
 
 def test_should_run_success() -> None:
     hour = datetime.datetime.utcnow().hour
-    assert should_run(
+    delay = get_task_delay(
         ScheduledTask(
             task_name="",
             labels={},
@@ -22,12 +23,13 @@ def test_should_run_success() -> None:
             cron=f"* {hour} * * *",
         ),
     )
+    assert delay is not None and delay >= 0
 
 
 def test_should_run_cron_str_offset() -> None:
     hour = datetime.datetime.now().hour
     zone = get_localzone()
-    assert should_run(
+    delay = get_task_delay(
         ScheduledTask(
             task_name="",
             labels={},
@@ -38,12 +40,13 @@ def test_should_run_cron_str_offset() -> None:
             cron_offset=str(zone),
         ),
     )
+    assert delay is not None and delay >= 0
 
 
 def test_should_run_cron_td_offset() -> None:
     offset = 2
     hour = datetime.datetime.utcnow().hour + offset
-    assert should_run(
+    delay = get_task_delay(
         ScheduledTask(
             task_name="",
             labels={},
@@ -54,11 +57,12 @@ def test_should_run_cron_td_offset() -> None:
             cron_offset=datetime.timedelta(hours=offset),
         ),
     )
+    assert delay is not None and delay >= 0
 
 
 def test_time_utc_without_zone() -> None:
     time = datetime.datetime.utcnow()
-    assert should_run(
+    delay = get_task_delay(
         ScheduledTask(
             task_name="",
             labels={},
@@ -68,11 +72,12 @@ def test_time_utc_without_zone() -> None:
             time=time - datetime.timedelta(seconds=1),
         ),
     )
+    assert delay is not None and delay >= 0
 
 
 def test_time_utc_with_zone() -> None:
     time = datetime.datetime.now(tz=pytz.UTC)
-    assert should_run(
+    delay = get_task_delay(
         ScheduledTask(
             task_name="",
             labels={},
@@ -82,12 +87,13 @@ def test_time_utc_with_zone() -> None:
             time=time - datetime.timedelta(seconds=1),
         ),
     )
+    assert delay is not None and delay >= 0
 
 
 def test_time_utc_with_local_zone() -> None:
     localtz = get_localzone()
     time = datetime.datetime.now(tz=localtz)
-    assert should_run(
+    delay = get_task_delay(
         ScheduledTask(
             task_name="",
             labels={},
@@ -97,11 +103,12 @@ def test_time_utc_with_local_zone() -> None:
             time=time - datetime.timedelta(seconds=1),
         ),
     )
+    assert delay is not None and delay >= 0
 
 
 def test_time_localtime_without_zone() -> None:
     time = datetime.datetime.now(tz=pytz.FixedOffset(240)).replace(tzinfo=None)
-    assert not should_run(
+    delay = get_task_delay(
         ScheduledTask(
             task_name="",
             labels={},
@@ -111,3 +118,20 @@ def test_time_localtime_without_zone() -> None:
             time=time - datetime.timedelta(seconds=1),
         ),
     )
+    assert delay is None
+
+
+@freeze_time("2023-01-14 12:00:00")
+def test_time_delay() -> None:
+    time = datetime.datetime.now(tz=pytz.UTC) + datetime.timedelta(seconds=15)
+    delay = get_task_delay(
+        ScheduledTask(
+            task_name="",
+            labels={},
+            args=[],
+            kwargs={},
+            source=DUMMY_SOURCE,
+            time=time,
+        ),
+    )
+    assert delay is not None and 15 == delay