feat: implement lazy cleanup for time index in ListRedisScheduleSource

- Added `_maybe_cleanup_time_index` method to manage time index cleanup at most once per minute.
- Introduced `_cleanup_time_index` method to remove stale entries older than one hour with empty time key lists.
- Updated `delete_schedule` to call `_maybe_cleanup_time_index` for efficient cleanup.
- Enhanced tests to verify the behavior of the new cleanup methods, ensuring proper handling of stale and recent entries.

Author: tmkarthi

taskiq_redis/list_schedule_source.py

@@ -1,4 +1,5 @@
 import datetime
+import time as _time
 from logging import getLogger
 from typing import Any
 
@@ -58,6 +59,7 @@ def __init__(
         self._delete_schedules_after_migration: bool = True
         self._skip_past_schedules = skip_past_schedules
         self._populate_time_index = populate_time_index
+        self._last_cleanup_time: float = 0
 
     async def startup(self) -> None:
         """
@@ -136,6 +138,42 @@ def _parse_time_key(self, key: str) -> datetime.datetime | None:
             logger.debug("Failed to parse time key %s", key)
             return None
 
+    async def _maybe_cleanup_time_index(self, redis: Redis) -> None:  # type: ignore[type-arg]
+        """
+        Run time index cleanup at most once per minute.
+
+        Called from delete_schedule after removing a time-based schedule,
+        since that's the path where time key lists become empty.
+        """
+        now = _time.monotonic()
+        if now - self._last_cleanup_time < 60:
+            return
+        self._last_cleanup_time = now
+        await self._cleanup_time_index(redis)
+
+    async def _cleanup_time_index(self, redis: Redis) -> None:  # type: ignore[type-arg]
+        """
+        Remove stale entries from the time index sorted set.
+
+        Only removes entries that are older than 1 hour AND whose
+        corresponding time key list is empty (or no longer exists).
+        This avoids a race condition where an eager cleanup in
+        delete_schedule could remove an index entry right as
+        add_schedule is creating a new schedule at the same minute.
+        """
+        one_hour_ago = (
+            datetime.datetime.now(datetime.timezone.utc)
+            - datetime.timedelta(hours=1)
+        ).timestamp()
+        stale_keys: list[bytes] = await redis.zrangebyscore(
+            self._get_time_index_key(),
+            "-inf",
+            one_hour_ago,
+        )
+        for key in stale_keys:
+            if await redis.llen(key) == 0:
+                await redis.zrem(self._get_time_index_key(), key)
+
     async def _get_previous_time_schedules(self) -> list[bytes]:
         """
         Function that gets all timed schedules that are in the past.
@@ -185,14 +223,7 @@ 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,
-                        )
+                    await self._maybe_cleanup_time_index(redis)
                 elif schedule.interval:
                     await redis.lrem(self._get_interval_key(), 0, schedule_id)  # type: ignore[misc]
 

tests/test_list_schedule_source.py

@@ -207,8 +207,10 @@ async def test_time_index_populated_on_add(redis_url: str) -> None:
 
 @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."""
+async def test_time_index_not_eagerly_cleaned_on_delete(redis_url: str) -> None:
+    """Test that delete_schedule does NOT eagerly remove the index entry.
+    This avoids a race condition where a concurrent add_schedule at the
+    same minute could lose its index entry."""
     prefix = uuid.uuid4().hex
     source = ListRedisScheduleSource(redis_url, prefix=prefix)
     schedule = ScheduledTask(
@@ -227,55 +229,110 @@ async def test_time_index_cleaned_on_delete(redis_url: str) -> None:
 
     await source.delete_schedule(schedule.schedule_id)
 
-    # After deletion, the index should be empty.
+    # Index entry is still present (lazy cleanup handles it later).
     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))
+        assert await redis.zcard(source._get_time_index_key()) == 1
 
 
 @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."""
+async def test_cleanup_removes_old_empty_entries(redis_url: str) -> None:
+    """Test that _cleanup_time_index removes index entries that are
+    older than 1 hour and whose time key lists are empty."""
     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)
+    with freeze_time("2025-01-01 00:00:00"):
+        source = ListRedisScheduleSource(redis_url, prefix=prefix)
+        old_time = datetime.datetime(
+            2024, 12, 31, 22, 0, tzinfo=datetime.timezone.utc,
+        )
+        schedule = ScheduledTask(
+            task_name="test_task",
+            labels={},
+            args=[],
+            kwargs={},
+            time=old_time,
+        )
+        await source.add_schedule(schedule)
+        # Prevent delete_schedule from triggering cleanup by pretending
+        # cleanup just ran (rate limiter blocks it).
+        import time
 
-    await source.delete_schedule(schedule1.schedule_id)
+        source._last_cleanup_time = time.monotonic()
+        await source.delete_schedule(schedule.schedule_id)
 
-    # Index should still have the entry because schedule2 remains.
+    # Index still has the stale entry (cleanup was rate-limited).
     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)
+    # Run cleanup directly — entry is > 1 hour old and empty.
+    with freeze_time("2025-01-01 00:00:00"):
+        async with Redis(connection_pool=source._connection_pool) as redis:
+            await source._cleanup_time_index(redis)
 
-    # Now the index should be empty.
+    # Now it should be cleaned up.
     async with Redis(connection_pool=source._connection_pool) as redis:
         assert await redis.zcard(source._get_time_index_key()) == 0
 
 
+@pytest.mark.anyio
+async def test_cleanup_keeps_non_empty_entries(redis_url: str) -> None:
+    """Test that _cleanup_time_index does NOT remove index entries whose
+    time key lists still have schedules, even if older than 1 hour."""
+    prefix = uuid.uuid4().hex
+    with freeze_time("2025-01-01 00:00:00"):
+        source = ListRedisScheduleSource(redis_url, prefix=prefix)
+        old_time = datetime.datetime(
+            2024, 12, 31, 22, 0, tzinfo=datetime.timezone.utc,
+        )
+        schedule = ScheduledTask(
+            task_name="test_task",
+            labels={},
+            args=[],
+            kwargs={},
+            time=old_time,
+        )
+        await source.add_schedule(schedule)
+
+    # Run cleanup — entry is > 1 hour old but list is NOT empty.
+    with freeze_time("2025-01-01 00:00:00"):
+        async with Redis(connection_pool=source._connection_pool) as redis:
+            await source._cleanup_time_index(redis)
+
+    # Entry should still be present.
+    async with Redis(connection_pool=source._connection_pool) as redis:
+        assert await redis.zcard(source._get_time_index_key()) == 1
+
+
+@pytest.mark.anyio
+async def test_cleanup_keeps_recent_empty_entries(redis_url: str) -> None:
+    """Test that _cleanup_time_index does NOT remove index entries that
+    are less than 1 hour old, even if their time key lists are empty."""
+    prefix = uuid.uuid4().hex
+    with freeze_time("2025-01-01 00:00:00"):
+        source = ListRedisScheduleSource(redis_url, prefix=prefix)
+        # 30 minutes ago — within the 1-hour safety window.
+        recent_time = datetime.datetime(
+            2024, 12, 31, 23, 30, tzinfo=datetime.timezone.utc,
+        )
+        schedule = ScheduledTask(
+            task_name="test_task",
+            labels={},
+            args=[],
+            kwargs={},
+            time=recent_time,
+        )
+        await source.add_schedule(schedule)
+        await source.delete_schedule(schedule.schedule_id)
+
+    # Run cleanup — entry is empty but only 30 min old.
+    with freeze_time("2025-01-01 00:00:00"):
+        async with Redis(connection_pool=source._connection_pool) as redis:
+            await source._cleanup_time_index(redis)
+
+    # Entry should still be present (not old enough).
+    async with Redis(connection_pool=source._connection_pool) as redis:
+        assert await redis.zcard(source._get_time_index_key()) == 1
+
+
 @pytest.mark.anyio
 @freeze_time("2025-01-01 00:00:00")
 async def test_past_schedules_found_via_time_index(redis_url: str) -> None:
@@ -343,37 +400,83 @@ async def test_populate_time_index_from_existing_keys(redis_url: str) -> None:
 
 
 @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."""
+async def test_post_send_triggers_cleanup(redis_url: str) -> None:
+    """Test the full lifecycle: add schedule, get it, post_send it,
+    then verify cleanup (triggered from delete_schedule) removes
+    the stale index entry when it's > 1 hour old."""
     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]
+    with freeze_time("2025-01-01 02:00:00"):
+        source = ListRedisScheduleSource(redis_url, prefix=prefix)
+        schedule = ScheduledTask(
+            task_name="test_task",
+            labels={},
+            args=[],
+            kwargs={},
+            time=datetime.datetime(
+                2025, 1, 1, 0, 30, tzinfo=datetime.timezone.utc,
+            ),
+        )
+        await source.add_schedule(schedule)
 
-    # Simulate sending the task.
-    for s in schedules:
-        await source.post_send(s)
+        # First run picks up past schedules.
+        schedules = await source.get_schedules()
+        assert schedules == [schedule]
 
-    # 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
+        # post_send -> delete_schedule -> _maybe_cleanup_time_index.
+        # The entry is > 1 hour old and the list becomes empty,
+        # so cleanup should remove it.
+        for s in schedules:
+            await source.post_send(s)
+
+        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 == []
+    with freeze_time("2025-01-01 02:01:00"):
+        schedules = await source.get_schedules()
+        assert schedules == []
+
+
+@pytest.mark.anyio
+async def test_cleanup_rate_limited(redis_url: str) -> None:
+    """Test that _maybe_cleanup_time_index only runs once per minute."""
+    prefix = uuid.uuid4().hex
+
+    with freeze_time("2025-01-01 02:00:00"):
+        source = ListRedisScheduleSource(redis_url, prefix=prefix)
+        old_time = datetime.datetime(
+            2025, 1, 1, 0, 30, tzinfo=datetime.timezone.utc,
+        )
+        sched1 = ScheduledTask(
+            task_name="task1",
+            labels={},
+            args=[],
+            kwargs={},
+            time=old_time,
+        )
+        sched2 = ScheduledTask(
+            task_name="task2",
+            labels={},
+            args=[],
+            kwargs={},
+            time=old_time,
+        )
+        await source.add_schedule(sched1)
+        await source.add_schedule(sched2)
+
+        # First delete triggers cleanup (first call, _last_cleanup_time=0).
+        # But the time key list still has sched2, so the entry is kept.
+        await source.delete_schedule(sched1.schedule_id)
+        async with Redis(connection_pool=source._connection_pool) as redis:
+            assert await redis.zcard(source._get_time_index_key()) == 1
+
+        # Second delete happens within the same minute, so cleanup
+        # is rate-limited and does NOT run — index entry remains
+        # even though the list is now empty.
+        await source.delete_schedule(sched2.schedule_id)
+        async with Redis(connection_pool=source._connection_pool) as redis:
+            assert await redis.zcard(source._get_time_index_key()) == 1
 
 
 @pytest.mark.anyio