diff --git a/taskiq_redis/list_schedule_source.py b/taskiq_redis/list_schedule_source.py index c2004ed..0f295d8 100644 --- a/taskiq_redis/list_schedule_source.py +++ b/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) diff --git a/tests/test_list_schedule_source.py b/tests/test_list_schedule_source.py index c21486b..29d043f 100644 --- a/tests/test_list_schedule_source.py +++ b/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")