fix: 批量删除任务完成前等待所有进度更新的异步回调完成

收集 run_coroutine_threadsafe 返回的 Future,在标记任务完成前
通过 asyncio.gather 等待所有进度更新回调执行完毕,避免任务
状态提前跳到 completed 而进度数据尚未写入 Redis 的竞态问题。
This commit is contained in:
fawney19
2026-03-09 02:52:38 +08:00
parent fd32597015
commit 654ce89541
2 changed files with 115 additions and 1 deletions

View File

@@ -0,0 +1,98 @@
from __future__ import annotations
import asyncio
from collections.abc import Coroutine, Generator
from concurrent.futures import Future
from contextlib import contextmanager
import pytest
from src.services.provider_keys import batch_delete_task as taskmod
@contextmanager
def _fake_db_context() -> Generator[object, None, None]:
yield object()
async def _fake_get_redis_client(*, require_redis: bool = False) -> object:
_ = require_redis
return object()
async def _noop_delete_side_effects(**_kwargs: object) -> None:
return None
@pytest.mark.asyncio
async def test_run_batch_delete_waits_for_progress_updates_before_completion(
monkeypatch: pytest.MonkeyPatch,
) -> None:
release_progress = asyncio.Event()
updates: list[dict[str, object]] = []
async def fake_update_task_field(
task_id: str,
r: object | None = None,
**fields: object,
) -> None:
_ = task_id, r
updates.append(dict(fields))
def fake_sync_delete(
provider_id: str,
key_ids: list[str],
progress_callback: object | None = None,
) -> int:
_ = provider_id, key_ids
assert callable(progress_callback)
progress_callback(1)
return 1
def fake_run_coroutine_threadsafe(
coro: Coroutine[object, object, object],
loop: asyncio.AbstractEventLoop,
) -> Future[object]:
future: Future[object] = Future()
async def runner() -> None:
await release_progress.wait()
try:
result = await coro
except Exception as exc:
future.set_exception(exc)
else:
future.set_result(result)
loop.call_soon_threadsafe(lambda: asyncio.create_task(runner()))
return future
monkeypatch.setattr(taskmod, "get_redis_client", _fake_get_redis_client)
monkeypatch.setattr(taskmod, "_update_task_field", fake_update_task_field)
monkeypatch.setattr(taskmod, "_sync_delete", fake_sync_delete)
monkeypatch.setattr(taskmod.asyncio, "run_coroutine_threadsafe", fake_run_coroutine_threadsafe)
monkeypatch.setattr("src.database.get_db_context", _fake_db_context)
monkeypatch.setattr(
"src.services.provider_keys.key_side_effects.run_delete_key_side_effects",
_noop_delete_side_effects,
)
task = asyncio.create_task(taskmod._run_batch_delete("task-1", "provider-1", ["key-1"]))
await asyncio.sleep(0.05)
assert not task.done()
assert updates == [{"status": taskmod.STATUS_RUNNING}]
release_progress.set()
await task
assert updates == [
{"status": taskmod.STATUS_RUNNING},
{"deleted": 1},
{
"status": taskmod.STATUS_COMPLETED,
"deleted": 1,
"message": "1 keys deleted",
},
]