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

@@ -10,6 +10,7 @@ import asyncio
import json import json
import uuid import uuid
from collections.abc import Callable from collections.abc import Callable
from concurrent.futures import Future
import redis.asyncio as aioredis import redis.asyncio as aioredis
from sqlalchemy import delete as sa_delete from sqlalchemy import delete as sa_delete
@@ -232,18 +233,33 @@ async def _run_batch_delete(
# 从工作线程安全地触发 Redis 进度更新 # 从工作线程安全地触发 Redis 进度更新
loop = asyncio.get_running_loop() loop = asyncio.get_running_loop()
progress_futures: list[Future[object]] = []
def on_progress(current: int) -> None: def on_progress(current: int) -> None:
try: try:
asyncio.run_coroutine_threadsafe( future = asyncio.run_coroutine_threadsafe(
_update_task_field(task_id, r=r, deleted=current), loop _update_task_field(task_id, r=r, deleted=current), loop
) )
progress_futures.append(future)
except RuntimeError: except RuntimeError:
pass pass
try: try:
affected = await asyncio.to_thread(_sync_delete, provider_id, key_ids, on_progress) affected = await asyncio.to_thread(_sync_delete, provider_id, key_ids, on_progress)
if progress_futures:
results = await asyncio.gather(
*(asyncio.wrap_future(f) for f in progress_futures),
return_exceptions=True,
)
for exc in results:
if isinstance(exc, Exception):
logger.debug(
"[BATCH_DELETE_TASK] progress update failed task={}: {}",
task_id,
exc,
)
# 副作用(缓存失效等)是异步操作,在事件循环中执行 # 副作用(缓存失效等)是异步操作,在事件循环中执行
if affected > 0: if affected > 0:
try: try:

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",
},
]