mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
fix: 批量删除任务完成前等待所有进度更新的异步回调完成
收集 run_coroutine_threadsafe 返回的 Future,在标记任务完成前 通过 asyncio.gather 等待所有进度更新回调执行完毕,避免任务 状态提前跳到 completed 而进度数据尚未写入 Redis 的竞态问题。
This commit is contained in:
@@ -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:
|
||||||
|
|||||||
98
tests/services/test_batch_delete_task.py
Normal file
98
tests/services/test_batch_delete_task.py
Normal 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",
|
||||||
|
},
|
||||||
|
]
|
||||||
Reference in New Issue
Block a user