fix: scan past time keys once per minute in ListRedisScheduleSource

taskiq_redis/list_schedule_source.py

@@ -46,7 +46,8 @@ def __init__(
         if serializer is None:
             serializer = PickleSerializer()
         self._serializer = serializer
-        self._is_first_run = True
+        self._last_catch_up_minute: datetime.datetime | None = None
+        self._past_schedule_ids: list[bytes] = []
         self._previous_schedule_source: ScheduleSource | None = None
         self._delete_schedules_after_migration: bool = True
         self._skip_past_schedules = skip_past_schedules
@@ -113,8 +114,8 @@ async def _get_previous_time_schedules(self) -> list[bytes]:
         We do this by getting all the time keys and checking if the time
         is less than the current time.
 
-        This function is called only during the first run to minimize
-        the number of requests to the Redis server.
+        It runs on the first get_schedules() call of each minute
+        to minimize the number of requests to the Redis server.
         """
         logger.info("Getting previous time schedules")
         minute_before = datetime.datetime.now(
@@ -127,7 +128,10 @@ async def _get_previous_time_schedules(self) -> list[bytes]:
             time_keys: list[str] = []
             # We need to get all the time keys and check if the time is less than
             # the current time.
-            async for key in redis.scan_iter(f"{self._prefix}:time:*"):
+            async for key in redis.scan_iter(
+                f"{self._prefix}:time:*",
+                count=1000,
+            ):
                 key_time = self._parse_time_key(key.decode())
                 if key_time and key_time <= minute_before:
                     time_keys.append(key.decode())
@@ -191,16 +195,20 @@ async def get_schedules(self) -> list["ScheduledTask"]:
         What it does is get all the cron schedules and time schedules
         for the current time and return them.
 
-        If it's the first run, it also gets all the time schedules
-        that are in the past and haven't been sent yet.
+        It also returns the time schedules that are in the past and
+        haven't been sent yet. Their ids are looked up on the first call
+        of every minute and reused by the other calls of that minute.
         """
         schedules = []
         current_time = datetime.datetime.now(datetime.timezone.utc)
-        timed: list[bytes] = []
-        # Only during first run, we need to get previous time schedules
-        if not self._skip_past_schedules:
-            timed = await self._get_previous_time_schedules()
-            self._is_first_run = False
+        current_minute = current_time.replace(second=0, microsecond=0)
+        if (
+            not self._skip_past_schedules
+            and current_minute != self._last_catch_up_minute
+        ):
+            self._past_schedule_ids = await self._get_previous_time_schedules()
+            self._last_catch_up_minute = current_minute
+        timed = list(self._past_schedule_ids)
         async with Redis(connection_pool=self._connection_pool) as redis:
             buffer: list[bytes] = []
             crons = await redis.lrange(self._get_cron_key(), 0, -1)

tests/test_list_schedule_source.py

@@ -1,5 +1,6 @@
 import datetime
 import uuid
+from unittest.mock import patch
 
 import pytest
 from freezegun import freeze_time
@@ -64,13 +65,49 @@ async def test_schedule_from_past(redis_url: str) -> None:
     # schedules that are in the past.
     schedules = await source.get_schedules()
     assert schedules == [schedule]
+    # Past schedules are returned on every call until they are sent.
+    assert await source.get_schedules() == [schedule]
     for schedule in schedules:
         await source.post_send(schedule)
-    # After getting the schedules for the second time,
-    # all the schedules in the past are ignored.
+    # Sent schedules are removed from the source.
     schedules = await source.get_schedules()
     assert schedules == []
 
+    # Schedule that lands in an already passed minute after the first run
+    # is picked up on the first call of the next minute.
+    late_schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        time=datetime.datetime.now(datetime.timezone.utc)
+        - datetime.timedelta(minutes=2),
+    )
+    await source.add_schedule(late_schedule)
+    with freeze_time("2025-01-01 00:01:00"):
+        assert await source.get_schedules() == [late_schedule]
+
+
+@pytest.mark.anyio
+async def test_past_schedules_lookup_once_per_minute(redis_url: str) -> None:
+    """Test that past time keys are scanned once per minute, not on every call."""
+    prefix = uuid.uuid4().hex
+    source = ListRedisScheduleSource(redis_url, prefix=prefix)
+    with (
+        patch.object(
+            source,
+            "_get_previous_time_schedules",
+            wraps=source._get_previous_time_schedules,
+        ) as lookup,
+        freeze_time("2025-01-01 00:00:00") as frozen,
+    ):
+        for _ in range(60):
+            await source.get_schedules()
+            frozen.tick()
+        assert lookup.await_count == 1
+        await source.get_schedules()
+        assert lookup.await_count == 2
+
 
 @pytest.mark.anyio
 @freeze_time("2025-01-01 00:00:00")