feat: add time index functionality to ListRedisScheduleSource

- Introduced `populate_time_index` parameter to backfill the time index from existing keys.
- Updated `startup` method to populate the time index if `populate_time_index` is set to True.
- Modified schedule addition and deletion to manage the time index sorted set.
- Added tests to verify time index population and cleanup behavior.

Author: tmkarthi

taskiq_redis/list_schedule_source.py

@@ -23,6 +23,7 @@ def __init__(
         serializer: TaskiqSerializer | None = None,
         buffer_size: int = 50,
         skip_past_schedules: bool = False,
+        populate_time_index: bool = False,
         **connection_kwargs: Any,
     ) -> None:
         """
@@ -34,6 +35,11 @@ def __init__(
         :param serializer: Serializer to use for the schedules
         :param buffer_size: Buffer size for getting schedules
         :param skip_past_schedules: Skip schedules that are in the past.
+        :param populate_time_index: If True, on startup run a one-time SCAN
+            to populate the time index sorted set from existing time keys.
+            This is needed for migrating from an older version that did not
+            maintain the time index. Set this to True once to backfill the
+            index, then set it back to False for subsequent runs.
         :param connection_kwargs: Additional connection kwargs
         """
         super().__init__()
@@ -51,6 +57,7 @@ def __init__(
         self._previous_schedule_source: ScheduleSource | None = None
         self._delete_schedules_after_migration: bool = True
         self._skip_past_schedules = skip_past_schedules
+        self._populate_time_index = populate_time_index
 
     async def startup(self) -> None:
         """
@@ -59,6 +66,9 @@ async def startup(self) -> None:
         By default this function does nothing.
         But if the previous schedule source is set,
         it will try to migrate schedules from it.
+
+        If populate_time_index is True, it will scan for existing
+        time keys and populate the time index sorted set.
         """
         if self._previous_schedule_source is not None:
             logger.info("Migrating schedules from previous source")
@@ -74,13 +84,36 @@ async def startup(self) -> None:
             await self._previous_schedule_source.shutdown()
             logger.info("Migration complete")
 
+        if self._populate_time_index:
+            logger.info("Populating time index from existing keys via scan")
+            async with Redis(connection_pool=self._connection_pool) as redis:
+                batch: dict[str, float] = {}
+                async for key in redis.scan_iter(f"{self._prefix}:time:*"):
+                    key_str = key.decode()
+                    key_time = self._parse_time_key(key_str)
+                    if key_time:
+                        batch[key_str] = key_time.timestamp()
+                        if len(batch) >= self._buffer_size:
+                            await redis.zadd(
+                                self._get_time_index_key(),
+                                batch,
+                            )
+                            batch = {}
+                if batch:
+                    await redis.zadd(self._get_time_index_key(), batch)
+            logger.info("Time index population complete")
+
     def _get_time_key(self, time: datetime.datetime) -> str:
         """Get the key for a time-based schedule."""
         if time.tzinfo is None:
             time = time.replace(tzinfo=datetime.timezone.utc)
         iso_time = time.astimezone(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M")
         return f"{self._prefix}:time:{iso_time}"
 
+    def _get_time_index_key(self) -> str:
+        """Get the key for the time index sorted set."""
+        return f"{self._prefix}:time_index"
+
     def _get_cron_key(self) -> str:
         """Get the key for a cron-based schedule."""
         return f"{self._prefix}:cron"
@@ -111,8 +144,8 @@ async def _get_previous_time_schedules(self) -> list[bytes]:
         we need to get all the schedules that are in the past and haven't
         been sent yet.
 
-        We do this by getting all the time keys and checking if the time
-        is less than the current time.
+        Uses the time index sorted set to look up past time keys
+        instead of scanning all Redis keys.
 
         This function is called only during the first run to minimize
         the number of requests to the Redis server.
@@ -125,13 +158,12 @@ async def _get_previous_time_schedules(self) -> list[bytes]:
         )
         schedules = []
         async with Redis(connection_pool=self._connection_pool) as redis:
-            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:*"):
-                key_time = self._parse_time_key(key.decode())
-                if key_time and key_time <= minute_before:
-                    time_keys.append(key.decode())
+            max_score = minute_before.timestamp()
+            time_keys: list[bytes] = await redis.zrangebyscore(
+                self._get_time_index_key(),
+                "-inf",
+                max_score,
+            )
             for key in time_keys:
                 schedules.extend(await redis.lrange(key, 0, -1))  # type: ignore[misc]
 
@@ -153,6 +185,14 @@ async def delete_schedule(self, schedule_id: str) -> None:
                 elif schedule.time is not None:
                     time_key = self._get_time_key(schedule.time)
                     await redis.lrem(time_key, 0, schedule_id)  # type: ignore[misc]
+                    # If the time key list is now empty, clean up both
+                    # the list key and its entry in the time index.
+                    if await redis.llen(time_key) == 0:
+                        await redis.delete(time_key)
+                        await redis.zrem(
+                            self._get_time_index_key(),
+                            time_key,
+                        )
                 elif schedule.interval:
                     await redis.lrem(self._get_interval_key(), 0, schedule_id)  # type: ignore[misc]
 
@@ -170,9 +210,21 @@ async def add_schedule(self, schedule: "ScheduledTask") -> None:
             if schedule.cron is not None:
                 await redis.rpush(self._get_cron_key(), schedule.schedule_id)  # type: ignore[misc]
             elif schedule.time is not None:
-                await redis.rpush(  # type: ignore[misc]
-                    self._get_time_key(schedule.time),
-                    schedule.schedule_id,
+                time_key = self._get_time_key(schedule.time)
+                await redis.rpush(time_key, schedule.schedule_id)  # type: ignore[misc]
+                # Add to the time index sorted set so we can look up
+                # past time keys without scanning all Redis keys.
+                time_val = schedule.time
+                if time_val.tzinfo is None:
+                    time_val = time_val.replace(tzinfo=datetime.timezone.utc)
+                score = (
+                    time_val.astimezone(datetime.timezone.utc)
+                    .replace(second=0, microsecond=0)
+                    .timestamp()
+                )
+                await redis.zadd(  # type: ignore[misc]
+                    self._get_time_index_key(),
+                    {time_key: score},
                 )
             elif schedule.interval:
                 await redis.rpush(  # type: ignore[misc]
@@ -200,7 +252,7 @@ async def get_schedules(self) -> list["ScheduledTask"]:
         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:
+        if not self._skip_past_schedules and self._is_first_run:
             timed = await self._get_previous_time_schedules()
             self._is_first_run = False
         async with Redis(connection_pool=self._connection_pool) as redis:

tests/test_list_schedule_source.py

@@ -3,6 +3,7 @@
 
 import pytest
 from freezegun import freeze_time
+from redis.asyncio import BlockingConnectionPool, Redis
 from taskiq import ScheduledTask
 
 from taskiq_redis.list_schedule_source import ListRedisScheduleSource
@@ -179,3 +180,224 @@ async def test_migration(redis_url: str) -> None:
     for old_schedule in old_schedules:
         with freeze_time(old_schedule.time):
             assert await source.get_schedules() == [old_schedule]
+
+
+@pytest.mark.anyio
+@freeze_time("2025-01-01 00:00:00")
+async def test_time_index_populated_on_add(redis_url: str) -> None:
+    """Test that adding a time schedule populates the time index sorted set."""
+    prefix = uuid.uuid4().hex
+    source = ListRedisScheduleSource(redis_url, prefix=prefix)
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        time=datetime.datetime.now(datetime.timezone.utc)
+        + datetime.timedelta(minutes=5),
+    )
+    await source.add_schedule(schedule)
+
+    # Verify the time index sorted set has an entry.
+    async with Redis(connection_pool=source._connection_pool) as redis:
+        members = await redis.zrange(source._get_time_index_key(), 0, -1)
+        assert len(members) == 1
+        assert members[0].decode() == source._get_time_key(schedule.time)
+
+
+@pytest.mark.anyio
+@freeze_time("2025-01-01 00:00:00")
+async def test_time_index_cleaned_on_delete(redis_url: str) -> None:
+    """Test that deleting last schedule from a time key cleans the index."""
+    prefix = uuid.uuid4().hex
+    source = ListRedisScheduleSource(redis_url, prefix=prefix)
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        time=datetime.datetime.now(datetime.timezone.utc)
+        + datetime.timedelta(minutes=5),
+    )
+    await source.add_schedule(schedule)
+
+    # Index has 1 entry.
+    async with Redis(connection_pool=source._connection_pool) as redis:
+        assert await redis.zcard(source._get_time_index_key()) == 1
+
+    await source.delete_schedule(schedule.schedule_id)
+
+    # After deletion, the index should be empty.
+    async with Redis(connection_pool=source._connection_pool) as redis:
+        assert await redis.zcard(source._get_time_index_key()) == 0
+        # The time key list itself should also be deleted.
+        assert not await redis.exists(source._get_time_key(schedule.time))
+
+
+@pytest.mark.anyio
+@freeze_time("2025-01-01 00:00:00")
+async def test_time_index_not_cleaned_when_other_schedules_remain(
+    redis_url: str,
+) -> None:
+    """Test that deleting one schedule doesn't remove the index entry
+    when other schedules still exist at the same time."""
+    prefix = uuid.uuid4().hex
+    source = ListRedisScheduleSource(redis_url, prefix=prefix)
+    schedule_time = datetime.datetime.now(
+        datetime.timezone.utc,
+    ) + datetime.timedelta(minutes=5)
+    schedule1 = ScheduledTask(
+        task_name="test_task_1",
+        labels={},
+        args=[],
+        kwargs={},
+        time=schedule_time,
+    )
+    schedule2 = ScheduledTask(
+        task_name="test_task_2",
+        labels={},
+        args=[],
+        kwargs={},
+        time=schedule_time,
+    )
+    await source.add_schedule(schedule1)
+    await source.add_schedule(schedule2)
+
+    await source.delete_schedule(schedule1.schedule_id)
+
+    # Index should still have the entry because schedule2 remains.
+    async with Redis(connection_pool=source._connection_pool) as redis:
+        assert await redis.zcard(source._get_time_index_key()) == 1
+
+    await source.delete_schedule(schedule2.schedule_id)
+
+    # Now the index should be empty.
+    async with Redis(connection_pool=source._connection_pool) as redis:
+        assert await redis.zcard(source._get_time_index_key()) == 0
+
+
+@pytest.mark.anyio
+@freeze_time("2025-01-01 00:00:00")
+async def test_past_schedules_found_via_time_index(redis_url: str) -> None:
+    """Test that past schedules are discovered via the time index
+    instead of a full SCAN."""
+    prefix = uuid.uuid4().hex
+    source = ListRedisScheduleSource(redis_url, prefix=prefix)
+    past_time = datetime.datetime.now(
+        datetime.timezone.utc,
+    ) - datetime.timedelta(minutes=5)
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        time=past_time,
+    )
+    await source.add_schedule(schedule)
+
+    # First call to get_schedules should find the past schedule via time index.
+    schedules = await source.get_schedules()
+    assert schedules == [schedule]
+
+
+@pytest.mark.anyio
+@freeze_time("2025-01-01 00:00:00")
+async def test_populate_time_index_from_existing_keys(redis_url: str) -> None:
+    """Test that populate_time_index=True backfills the sorted set
+    from existing time keys created without the index."""
+    prefix = uuid.uuid4().hex
+
+    # Simulate old-style data: create time key lists directly in Redis
+    # without populating the time index sorted set.
+    pool = BlockingConnectionPool.from_url(url=redis_url)
+    past_times = [
+        datetime.datetime(2024, 12, 31, 23, 55, tzinfo=datetime.timezone.utc),
+        datetime.datetime(2024, 12, 31, 23, 56, tzinfo=datetime.timezone.utc),
+        datetime.datetime(2024, 12, 31, 23, 57, tzinfo=datetime.timezone.utc),
+    ]
+
+    source_for_keys = ListRedisScheduleSource(redis_url, prefix=prefix)
+    async with Redis(connection_pool=pool) as redis:
+        for t in past_times:
+            time_key = source_for_keys._get_time_key(t)
+            # Push a dummy schedule ID directly (bypassing add_schedule
+            # to simulate old behavior without time index).
+            await redis.rpush(time_key, f"sched_{t.minute}")  # type: ignore[misc]
+
+        # Verify no time index exists yet.
+        assert await redis.zcard(source_for_keys._get_time_index_key()) == 0
+    await pool.disconnect()
+
+    # Now create a source with populate_time_index=True.
+    source = ListRedisScheduleSource(
+        redis_url,
+        prefix=prefix,
+        populate_time_index=True,
+    )
+    await source.startup()
+
+    # The time index should now be populated.
+    async with Redis(connection_pool=source._connection_pool) as redis:
+        count = await redis.zcard(source._get_time_index_key())
+        assert count == len(past_times)
+
+
+@pytest.mark.anyio
+@freeze_time("2025-01-01 00:00:00")
+async def test_post_send_cleans_time_index(redis_url: str) -> None:
+    """Test that post_send (which calls delete_schedule for time tasks)
+    properly cleans up the time index."""
+    prefix = uuid.uuid4().hex
+    source = ListRedisScheduleSource(redis_url, prefix=prefix)
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        time=datetime.datetime.now(datetime.timezone.utc)
+        - datetime.timedelta(minutes=3),
+    )
+    await source.add_schedule(schedule)
+
+    # First run picks up past schedules.
+    schedules = await source.get_schedules()
+    assert schedules == [schedule]
+
+    # Simulate sending the task.
+    for s in schedules:
+        await source.post_send(s)
+
+    # Time index should be empty now.
+    async with Redis(connection_pool=source._connection_pool) as redis:
+        assert await redis.zcard(source._get_time_index_key()) == 0
+
+    # Second run should return nothing.
+    schedules = await source.get_schedules()
+    assert schedules == []
+
+
+@pytest.mark.anyio
+@freeze_time("2025-01-01 00:00:00")
+async def test_cron_and_interval_not_in_time_index(redis_url: str) -> None:
+    """Test that cron and interval schedules do not affect the time index."""
+    prefix = uuid.uuid4().hex
+    source = ListRedisScheduleSource(redis_url, prefix=prefix)
+    cron_schedule = ScheduledTask(
+        task_name="cron_task",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+    )
+    interval_schedule = ScheduledTask(
+        task_name="interval_task",
+        labels={},
+        args=[],
+        kwargs={},
+        interval=datetime.timedelta(seconds=30),
+    )
+    await source.add_schedule(cron_schedule)
+    await source.add_schedule(interval_schedule)
+
+    async with Redis(connection_pool=source._connection_pool) as redis:
+        assert await redis.zcard(source._get_time_index_key()) == 0