fix(tasks): clean up singleflight cache on cancellation (#6174)

Signed-off-by: daixiheguu <daixihegu@outlook.com>

Author: daixiheguu

src/task_scheduler.py

@@ -84,19 +84,30 @@ async def _cached(key: Tuple, ttl: float, fetch: Callable[[], Awaitable[Any]]) -
             pending = fut
             owner = True
     if not owner:
-        return await pending
+        # A cancelled waiter must not cancel the shared Future for the owner
+        # and every other waiter.
+        return await asyncio.shield(pending)
     try:
         val = await fetch()
         async with _shared_cache_lock:
             _shared_cache[key] = (time.monotonic() + ttl, val)
-            _shared_cache_pending.pop(key, None)
         pending.set_result(val)
         return val
+    except asyncio.CancelledError:
+        # Cancellation is a BaseException on supported Python versions, so it
+        # bypasses the Exception handler below. Wake all current waiters while
+        # allowing a later caller to retry the fetch.
+        pending.cancel()
+        raise
     except Exception as e:
-        async with _shared_cache_lock:
-            _shared_cache_pending.pop(key, None)
         pending.set_exception(e)
         raise
+    finally:
+        # Keep this cleanup synchronous so a second cancellation cannot
+        # interrupt it and leave a permanently pending Future behind. All
+        # access runs on the scheduler's event-loop thread.
+        if _shared_cache_pending.get(key) is pending:
+            _shared_cache_pending.pop(key, None)
 
 
 def compute_next_run(schedule: str, scheduled_time: str,

tests/test_task_scheduler_cache.py

@@ -0,0 +1,86 @@
+import asyncio
+
+import pytest
+
+from src import task_scheduler
+
+
+@pytest.fixture(autouse=True)
+def clear_shared_cache():
+    task_scheduler._shared_cache.clear()
+    task_scheduler._shared_cache_pending.clear()
+    yield
+    task_scheduler._shared_cache.clear()
+    task_scheduler._shared_cache_pending.clear()
+
+
+async def test_cached_owner_cancellation_wakes_waiters_and_allows_retry():
+    key = ("cancelled-owner",)
+    fetch_started = asyncio.Event()
+
+    async def blocked_fetch():
+        fetch_started.set()
+        await asyncio.Event().wait()
+
+    owner = asyncio.create_task(task_scheduler._cached(key, 60, blocked_fetch))
+    await fetch_started.wait()
+
+    async def unexpected_fetch():
+        pytest.fail("a waiter must share the owner's fetch")
+
+    waiter = asyncio.create_task(task_scheduler._cached(key, 60, unexpected_fetch))
+    await asyncio.sleep(0)
+
+    owner.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await owner
+    with pytest.raises(asyncio.CancelledError):
+        await asyncio.wait_for(waiter, timeout=1)
+
+    assert key not in task_scheduler._shared_cache_pending
+
+    async def retry_fetch():
+        return "fresh"
+
+    result = await asyncio.wait_for(
+        task_scheduler._cached(key, 60, retry_fetch),
+        timeout=1,
+    )
+    assert result == "fresh"
+
+
+async def test_cached_waiter_cancellation_does_not_cancel_shared_fetch():
+    key = ("cancelled-waiter",)
+    fetch_started = asyncio.Event()
+    release_fetch = asyncio.Event()
+
+    async def blocked_fetch():
+        fetch_started.set()
+        await release_fetch.wait()
+        return "shared"
+
+    owner = asyncio.create_task(task_scheduler._cached(key, 60, blocked_fetch))
+    await fetch_started.wait()
+
+    async def unexpected_fetch():
+        pytest.fail("a waiter must share the owner's fetch")
+
+    waiter = asyncio.create_task(task_scheduler._cached(key, 60, unexpected_fetch))
+    await asyncio.sleep(0)
+    waiter.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await waiter
+
+    pending = task_scheduler._shared_cache_pending[key]
+    assert not pending.cancelled()
+    assert not owner.done()
+
+    release_fetch.set()
+    assert await asyncio.wait_for(owner, timeout=1) == "shared"
+    assert key not in task_scheduler._shared_cache_pending
+
+    async def cache_miss():
+        pytest.fail("the successful owner result should be cached")
+
+    assert await task_scheduler._cached(key, 60, cache_miss) == "shared"