667 lines
22 KiB
Python
667 lines
22 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
from datetime import datetime, timedelta, timezone
|
||
from typing import Any
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.config import settings
|
||
from app.models.chat_generation_task import ChatGenerationTask
|
||
from app.services.celery_download_recovery_service import (
|
||
ensure_aware_utc,
|
||
get_download_active_payloads,
|
||
get_due_download_record_ids,
|
||
postpone_download_active_check,
|
||
remove_download_active,
|
||
)
|
||
from app.services.generation_log_service import log_provider_call, log_task_event
|
||
from app.services.generation_module_hook_service import notify_chat_generation_task_finished
|
||
from app.services.generation_provider_service import poll_provider_task
|
||
from app.services.generation_refund_service import mark_chat_generation_task_failed_and_refund_once
|
||
from app.services.redis_registry_service import (
|
||
redis_get_due_registry_ids,
|
||
redis_get_registry_payloads,
|
||
redis_postpone_registry_item,
|
||
redis_remove_registry_item,
|
||
)
|
||
|
||
logger = logging.getLogger("video_gen")
|
||
|
||
ALLOWED_GENERATION_MODES = {"chatapi_async", "hot_opening_replicate", "shot_replicate"}
|
||
POLL_QUEUE = "gen_provider_poll"
|
||
|
||
|
||
def _now() -> datetime:
|
||
return datetime.now(timezone.utc)
|
||
|
||
|
||
def _is_expired(value: datetime | None, now: datetime | None = None) -> bool:
|
||
checked = ensure_aware_utc(value)
|
||
if checked is None:
|
||
return True
|
||
return checked <= (now or _now())
|
||
|
||
|
||
def _queue_timeout_at(task: ChatGenerationTask, now: datetime | None = None) -> datetime:
|
||
current_time = now or _now()
|
||
enqueued_at = ensure_aware_utc(task.download_enqueued_at)
|
||
if enqueued_at is None:
|
||
return current_time
|
||
return enqueued_at + timedelta(seconds=int(settings.DOWNLOAD_TASK_QUEUE_TIMEOUT_SECONDS or 300))
|
||
|
||
|
||
def _is_queue_timeout(task: ChatGenerationTask, now: datetime | None = None) -> bool:
|
||
current_time = now or _now()
|
||
return _queue_timeout_at(task, current_time) <= current_time
|
||
|
||
|
||
def _is_final_task_state(task: ChatGenerationTask) -> bool:
|
||
return task.status in ("completed", "failed") or task.pipeline_stage in (
|
||
"done",
|
||
"failed",
|
||
"timeout",
|
||
"download_failed",
|
||
)
|
||
|
||
|
||
def _is_success(status: str | None) -> bool:
|
||
return str(status or "").lower() in ("succeeded", "success", "completed", "done")
|
||
|
||
|
||
def _is_failed(status: str | None) -> bool:
|
||
return str(status or "").lower() in ("failed", "error", "canceled", "cancelled")
|
||
|
||
|
||
def _engine_snapshot(task: ChatGenerationTask) -> dict[str, Any]:
|
||
try:
|
||
value = json.loads(task.engine_snapshot_json or "{}")
|
||
return value if isinstance(value, dict) else {}
|
||
except Exception:
|
||
return {}
|
||
|
||
|
||
def _poll_queue_timeout_at(now: datetime | None = None) -> datetime:
|
||
current_time = now or _now()
|
||
return current_time + timedelta(seconds=int(settings.POLL_TASK_QUEUE_TIMEOUT_SECONDS or 120))
|
||
|
||
|
||
async def _remove_poll_active(task_id: str) -> None:
|
||
await redis_remove_registry_item(
|
||
hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
|
||
zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
|
||
item_id=task_id,
|
||
log_context="poll_active",
|
||
)
|
||
|
||
|
||
async def _postpone_poll_active(
|
||
*,
|
||
task_id: str,
|
||
payload: dict[str, Any] | None = None,
|
||
check_at: datetime | int | float | None = None,
|
||
) -> None:
|
||
await redis_postpone_registry_item(
|
||
hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
|
||
zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
|
||
item_id=task_id,
|
||
payload=payload,
|
||
check_at=check_at or _poll_queue_timeout_at(),
|
||
log_context="poll_active",
|
||
)
|
||
|
||
|
||
async def recover_one_download_task(
|
||
db: AsyncSession,
|
||
task: ChatGenerationTask,
|
||
*,
|
||
payload: dict[str, Any] | None = None,
|
||
source: str = "startup_db",
|
||
) -> str:
|
||
from app.tasks.generation_download_tasks import (
|
||
DOWNLOAD_STAGE_DOWNLOADING,
|
||
DOWNLOAD_STAGE_QUEUED,
|
||
DOWNLOAD_STAGE_RETRY_WAITING,
|
||
enqueue_download_task,
|
||
)
|
||
|
||
current_time = _now()
|
||
|
||
if not task:
|
||
return "skip_missing_task"
|
||
if task.generation_mode not in ALLOWED_GENERATION_MODES:
|
||
await remove_download_active(task.id)
|
||
return "clean_invalid_mode"
|
||
if _is_final_task_state(task):
|
||
await remove_download_active(task.id)
|
||
return "clean_final_state"
|
||
if task.status != "generating":
|
||
await remove_download_active(task.id)
|
||
return "clean_not_generating"
|
||
if not task.remote_result_url:
|
||
return "skip_no_remote_result_url"
|
||
|
||
stage = task.pipeline_stage
|
||
redis_payload = payload or {}
|
||
|
||
if stage == "result_ready":
|
||
await log_task_event(
|
||
task,
|
||
event_type="DOWNLOAD_RECOVERY_ENQUEUE",
|
||
message=f"{source} 发现 result_ready 未完成下载，启动时恢复投递下载任务",
|
||
detail={"payload": redis_payload},
|
||
)
|
||
await enqueue_download_task(
|
||
db,
|
||
task,
|
||
recover=True,
|
||
reason=f"{source}_result_ready",
|
||
)
|
||
return "recover_result_ready"
|
||
|
||
if stage == DOWNLOAD_STAGE_QUEUED:
|
||
if _is_queue_timeout(task, current_time):
|
||
await log_task_event(
|
||
task,
|
||
event_type="DOWNLOAD_RECOVERY_ENQUEUE",
|
||
message=f"{source} 发现 download_queued 长时间未消费，启动时恢复投递下载任务",
|
||
detail={"payload": redis_payload},
|
||
)
|
||
await enqueue_download_task(
|
||
db,
|
||
task,
|
||
recover=True,
|
||
reason=f"{source}_download_queued_timeout",
|
||
)
|
||
return "recover_queued_timeout"
|
||
|
||
await postpone_download_active_check(
|
||
record_id=task.id,
|
||
payload=payload,
|
||
check_at=_queue_timeout_at(task, current_time),
|
||
)
|
||
return "skip_queued_not_timeout"
|
||
|
||
if stage == DOWNLOAD_STAGE_DOWNLOADING:
|
||
if _is_expired(task.download_lease_until, current_time):
|
||
await log_task_event(
|
||
task,
|
||
event_type="DOWNLOAD_RECOVERY_ENQUEUE",
|
||
message=f"{source} 发现 downloading lease 过期，启动时恢复投递下载任务",
|
||
detail={"payload": redis_payload},
|
||
)
|
||
await enqueue_download_task(
|
||
db,
|
||
task,
|
||
recover=True,
|
||
reason=f"{source}_downloading_lease_expired",
|
||
)
|
||
return "recover_downloading_expired"
|
||
|
||
await postpone_download_active_check(
|
||
record_id=task.id,
|
||
payload=payload,
|
||
check_at=task.download_lease_until,
|
||
)
|
||
return "skip_downloading_alive"
|
||
|
||
if stage == DOWNLOAD_STAGE_RETRY_WAITING:
|
||
if _is_expired(task.download_next_retry_at, current_time):
|
||
await log_task_event(
|
||
task,
|
||
event_type="DOWNLOAD_RECOVERY_ENQUEUE",
|
||
message=f"{source} 发现 retry_waiting 到期，启动时恢复投递下载任务",
|
||
detail={"payload": redis_payload},
|
||
)
|
||
await enqueue_download_task(
|
||
db,
|
||
task,
|
||
recover=True,
|
||
reason=f"{source}_retry_waiting_due",
|
||
)
|
||
return "recover_retry_due"
|
||
|
||
await postpone_download_active_check(
|
||
record_id=task.id,
|
||
payload=payload,
|
||
check_at=task.download_next_retry_at,
|
||
)
|
||
return "skip_retry_waiting_not_due"
|
||
|
||
return f"skip_stage_{stage}"
|
||
|
||
|
||
async def recover_download_tasks_once(db: AsyncSession) -> dict[str, Any]:
|
||
"""启动时下载容灾扫描。
|
||
|
||
先按 Redis active_index 找到到期下载任务；Redis 不可用或索引丢失时，
|
||
再通过 DB fallback 扫描 result_ready/download_* 状态，避免任务永久卡住。
|
||
"""
|
||
checked_ids: set[str] = set()
|
||
results: dict[str, int] = {}
|
||
|
||
due_ids = await get_due_download_record_ids(
|
||
limit=settings.DOWNLOAD_RECOVERY_BATCH_SIZE,
|
||
)
|
||
payloads = await get_download_active_payloads(due_ids)
|
||
|
||
for task_id in due_ids:
|
||
result = await db.execute(
|
||
select(ChatGenerationTask)
|
||
.where(
|
||
ChatGenerationTask.id == task_id,
|
||
ChatGenerationTask.deleted_at.is_(None),
|
||
)
|
||
.with_for_update()
|
||
.limit(1)
|
||
)
|
||
task = result.scalar_one_or_none()
|
||
if task is None:
|
||
await remove_download_active(task_id)
|
||
action = "clean_missing_task"
|
||
else:
|
||
checked_ids.add(task.id)
|
||
action = await recover_one_download_task(
|
||
db,
|
||
task,
|
||
payload=payloads.get(task_id),
|
||
source="startup_redis",
|
||
)
|
||
results[action] = results.get(action, 0) + 1
|
||
|
||
# DB fallback：不依赖 Redis active 注册表。
|
||
fallback_result = await db.execute(
|
||
select(ChatGenerationTask)
|
||
.where(
|
||
ChatGenerationTask.deleted_at.is_(None),
|
||
ChatGenerationTask.generation_mode.in_(["chatapi_async", "hot_opening_replicate", "shot_replicate"]),
|
||
ChatGenerationTask.status == "generating",
|
||
ChatGenerationTask.remote_result_url.is_not(None),
|
||
ChatGenerationTask.pipeline_stage.in_(
|
||
["result_ready", "download_queued", "downloading", "retry_waiting"]
|
||
),
|
||
)
|
||
.order_by(ChatGenerationTask.updated_at.asc())
|
||
.limit(int(settings.DOWNLOAD_RECOVERY_BATCH_SIZE or 100))
|
||
.with_for_update(skip_locked=True)
|
||
)
|
||
fallback_tasks = fallback_result.scalars().all()
|
||
|
||
for task in fallback_tasks:
|
||
if task.id in checked_ids:
|
||
continue
|
||
action = await recover_one_download_task(
|
||
db,
|
||
task,
|
||
payload=None,
|
||
source="startup_db",
|
||
)
|
||
results[action] = results.get(action, 0) + 1
|
||
checked_ids.add(task.id)
|
||
|
||
return {"checked": len(checked_ids), "results": results}
|
||
|
||
|
||
async def _mark_timeout(
|
||
db: AsyncSession,
|
||
task: ChatGenerationTask,
|
||
*,
|
||
error_message: str = "任务超时",
|
||
) -> str:
|
||
await mark_chat_generation_task_failed_and_refund_once(
|
||
db,
|
||
task=task,
|
||
error_message=error_message,
|
||
pipeline_stage="timeout",
|
||
)
|
||
await notify_chat_generation_task_finished(db, task)
|
||
await db.commit()
|
||
await _remove_poll_active(task.id)
|
||
await log_task_event(
|
||
task,
|
||
event_type="TASK_TIMEOUT",
|
||
to_status="failed",
|
||
to_stage="timeout",
|
||
)
|
||
return "mark_timeout"
|
||
|
||
|
||
async def _mark_failed(
|
||
db: AsyncSession,
|
||
task: ChatGenerationTask,
|
||
*,
|
||
error_message: str,
|
||
event_type: str = "POLL_FAILED",
|
||
detail: Any = None,
|
||
) -> str:
|
||
await mark_chat_generation_task_failed_and_refund_once(
|
||
db,
|
||
task=task,
|
||
error_message=error_message,
|
||
pipeline_stage="failed",
|
||
)
|
||
await notify_chat_generation_task_finished(db, task)
|
||
await db.commit()
|
||
await _remove_poll_active(task.id)
|
||
await log_task_event(task, event_type=event_type, message=task.error_message, detail=detail)
|
||
return "mark_failed"
|
||
|
||
|
||
async def _try_final_poll_before_timeout(db: AsyncSession, task: ChatGenerationTask) -> str:
|
||
"""超时前最后查一次供应商，避免 Celery 中断导致本地假超时。
|
||
|
||
如果供应商已经成功，继续进入下载；如果仍 running 或查询失败，再按超时处理。
|
||
"""
|
||
from app.tasks.generation_download_tasks import enqueue_download_task
|
||
|
||
if not (task.provider_task_id or task.seedance_task_id):
|
||
return await _mark_timeout(db, task)
|
||
|
||
try:
|
||
poll_result = await poll_provider_task(db, task)
|
||
status = poll_result.get("status")
|
||
response_data = poll_result.get("response_data")
|
||
except Exception as exc:
|
||
await log_task_event(
|
||
task,
|
||
event_type="FINAL_POLL_BEFORE_TIMEOUT_ERROR",
|
||
message=str(exc),
|
||
)
|
||
return await _mark_timeout(db, task)
|
||
|
||
try:
|
||
provider_response = json.loads(response_data or "{}")
|
||
except Exception:
|
||
provider_response = {"raw": response_data}
|
||
|
||
snapshot = _engine_snapshot(task)
|
||
await log_provider_call(
|
||
task,
|
||
provider=snapshot.get("provider") or "ark",
|
||
api_type=f"{task.gen_type}_final_poll_before_timeout",
|
||
model=snapshot.get("model_name"),
|
||
engine_id=task.engine_id,
|
||
status="success",
|
||
provider_task_id=task.seedance_task_id or task.provider_task_id,
|
||
response_data=provider_response,
|
||
)
|
||
|
||
if _is_success(status):
|
||
if task.gen_type == "image":
|
||
task.remote_result_url = poll_result.get("image_url")
|
||
task.image_tokens_used = poll_result.get("image_tokens", 0) or 0
|
||
else:
|
||
task.remote_result_url = poll_result.get("video_url")
|
||
task.video_tokens_used = poll_result.get("video_tokens", 0) or 0
|
||
|
||
task.provider_response_json = response_data
|
||
if not task.remote_result_url:
|
||
return await _mark_failed(
|
||
db,
|
||
task,
|
||
error_message="供应商任务成功但未返回结果URL",
|
||
detail=poll_result,
|
||
)
|
||
|
||
task.pipeline_stage = "result_ready"
|
||
task.retry_count = 0
|
||
await db.commit()
|
||
await _remove_poll_active(task.id)
|
||
await log_task_event(
|
||
task,
|
||
event_type="POLL_SUCCESS_AFTER_TIMEOUT_RECOVERY",
|
||
to_stage="result_ready",
|
||
detail=poll_result,
|
||
)
|
||
await enqueue_download_task(db, task, recover=True, reason="final_poll_before_timeout_success")
|
||
return "recover_timeout_success_to_download"
|
||
|
||
if _is_failed(status):
|
||
task.provider_response_json = response_data
|
||
return await _mark_failed(
|
||
db,
|
||
task,
|
||
error_message=poll_result.get("error") or f"供应商任务失败: {status}",
|
||
detail=poll_result,
|
||
)
|
||
|
||
await log_task_event(
|
||
task,
|
||
event_type="FINAL_POLL_BEFORE_TIMEOUT_PENDING",
|
||
message=f"status={status}",
|
||
detail=poll_result,
|
||
)
|
||
return await _mark_timeout(db, task)
|
||
|
||
|
||
async def recover_one_generation_task(
|
||
db: AsyncSession,
|
||
task: ChatGenerationTask,
|
||
*,
|
||
payload: dict[str, Any] | None = None,
|
||
source: str = "startup_db",
|
||
) -> str:
|
||
from app.tasks.generation_create_tasks import chatapi_create_generation_task
|
||
from app.tasks.generation_download_tasks import enqueue_download_task
|
||
from app.tasks.generation_poll_tasks import poll_generation_task, register_poll_active
|
||
|
||
current_time = _now()
|
||
redis_payload = payload or {}
|
||
|
||
if not task:
|
||
return "skip_missing_task"
|
||
if task.generation_mode not in ALLOWED_GENERATION_MODES:
|
||
await _remove_poll_active(task.id)
|
||
return "clean_invalid_mode"
|
||
if _is_final_task_state(task):
|
||
await _remove_poll_active(task.id)
|
||
return "clean_final_state"
|
||
if task.status != "generating":
|
||
await _remove_poll_active(task.id)
|
||
return "clean_not_generating"
|
||
|
||
if task.deadline_at and _is_expired(task.deadline_at, current_time):
|
||
if task.pipeline_stage in ("waiting_remote", "polling"):
|
||
return await _try_final_poll_before_timeout(db, task)
|
||
return await _mark_timeout(db, task)
|
||
|
||
if task.pipeline_stage in ("queued", "preparing", "creating_provider_task"):
|
||
if task.provider_task_id or task.seedance_task_id:
|
||
task.pipeline_stage = "waiting_remote"
|
||
await db.commit()
|
||
await log_task_event(
|
||
task,
|
||
event_type="GENERATION_RECOVERY_ENQUEUE",
|
||
message=f"{source} 发现创建阶段已存在供应商任务ID，恢复投递轮询队列",
|
||
detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
|
||
)
|
||
poll_generation_task.apply_async(args=[task.id], queue=POLL_QUEUE, countdown=0)
|
||
await register_poll_active(
|
||
task,
|
||
check_at=_poll_queue_timeout_at(),
|
||
reason=f"{source}_create_stage_has_provider_id",
|
||
)
|
||
return "recover_poll_from_create_stage"
|
||
|
||
await log_task_event(
|
||
task,
|
||
event_type="GENERATION_RECOVERY_ENQUEUE",
|
||
message=f"{source} 发现创建阶段任务未完成，恢复投递创建队列",
|
||
detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
|
||
)
|
||
chatapi_create_generation_task.apply_async(
|
||
args=[task.id],
|
||
queue="gen_chatapi_create",
|
||
countdown=0,
|
||
)
|
||
return "recover_create"
|
||
|
||
if task.pipeline_stage in ("waiting_remote", "polling"):
|
||
if task.remote_result_url:
|
||
await _remove_poll_active(task.id)
|
||
await enqueue_download_task(
|
||
db,
|
||
task,
|
||
recover=True,
|
||
reason=f"{source}_waiting_remote_has_result",
|
||
)
|
||
return "recover_waiting_has_result"
|
||
|
||
if task.provider_task_id or task.seedance_task_id:
|
||
await log_task_event(
|
||
task,
|
||
event_type="GENERATION_RECOVERY_ENQUEUE",
|
||
message=f"{source} 发现远程等待/轮询阶段任务未完成，恢复投递轮询队列",
|
||
detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
|
||
)
|
||
task.pipeline_stage = "waiting_remote"
|
||
await db.commit()
|
||
poll_generation_task.apply_async(
|
||
args=[task.id],
|
||
queue=POLL_QUEUE,
|
||
countdown=0,
|
||
)
|
||
await register_poll_active(
|
||
task,
|
||
check_at=_poll_queue_timeout_at(),
|
||
reason=f"{source}_recover_poll",
|
||
)
|
||
return "recover_poll"
|
||
|
||
await log_task_event(
|
||
task,
|
||
event_type="GENERATION_RECOVERY_ENQUEUE",
|
||
message=f"{source} 发现任务缺少供应商任务ID，恢复投递创建队列",
|
||
detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
|
||
)
|
||
task.pipeline_stage = "queued"
|
||
await db.commit()
|
||
await _remove_poll_active(task.id)
|
||
chatapi_create_generation_task.apply_async(
|
||
args=[task.id],
|
||
queue="gen_chatapi_create",
|
||
countdown=0,
|
||
)
|
||
return "recover_create_missing_provider_id"
|
||
|
||
if task.pipeline_stage == "result_ready":
|
||
await _remove_poll_active(task.id)
|
||
if task.remote_result_url:
|
||
await enqueue_download_task(
|
||
db,
|
||
task,
|
||
recover=True,
|
||
reason=f"{source}_generation_result_ready",
|
||
)
|
||
return "recover_result_ready"
|
||
return "skip_result_ready_no_url"
|
||
|
||
return f"skip_stage_{task.pipeline_stage}"
|
||
|
||
|
||
async def recover_generation_tasks_once(db: AsyncSession) -> dict[str, Any]:
|
||
"""启动时生成链路容灾扫描。
|
||
|
||
不新增 Celery beat，不新增 worker 命令；worker 启动时由 Redis 锁保证只投递一次。
|
||
恢复顺序：
|
||
1. Redis poll active_index 到期任务；
|
||
2. DB fallback 扫描 queued/creating/waiting_remote/polling/result_ready；
|
||
3. 下载阶段仍由 recover_download_tasks_once 兜底。
|
||
"""
|
||
checked_ids: set[str] = set()
|
||
results: dict[str, int] = {}
|
||
|
||
due_poll_ids = await redis_get_due_registry_ids(
|
||
zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
|
||
limit=int(settings.POLL_RECOVERY_BATCH_SIZE or settings.GENERATION_RECOVERY_BATCH_SIZE or 100),
|
||
log_context="poll_active",
|
||
)
|
||
poll_payloads = await redis_get_registry_payloads(
|
||
hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
|
||
item_ids=due_poll_ids,
|
||
log_context="poll_active",
|
||
)
|
||
|
||
for task_id in due_poll_ids:
|
||
result = await db.execute(
|
||
select(ChatGenerationTask)
|
||
.where(
|
||
ChatGenerationTask.id == task_id,
|
||
ChatGenerationTask.deleted_at.is_(None),
|
||
)
|
||
.with_for_update()
|
||
.limit(1)
|
||
)
|
||
task = result.scalar_one_or_none()
|
||
if task is None:
|
||
await _remove_poll_active(task_id)
|
||
action = "clean_missing_poll_task"
|
||
else:
|
||
checked_ids.add(task.id)
|
||
action = await recover_one_generation_task(
|
||
db,
|
||
task,
|
||
payload=poll_payloads.get(task_id),
|
||
source="startup_poll_redis",
|
||
)
|
||
results[action] = results.get(action, 0) + 1
|
||
|
||
batch_size = int(settings.GENERATION_RECOVERY_BATCH_SIZE or settings.DOWNLOAD_RECOVERY_BATCH_SIZE or 100)
|
||
max_rounds = max(1, int(settings.GENERATION_RECOVERY_MAX_ROUNDS or 1))
|
||
total_db_checked = 0
|
||
|
||
for _round in range(max_rounds):
|
||
query_result = await db.execute(
|
||
select(ChatGenerationTask)
|
||
.where(
|
||
ChatGenerationTask.deleted_at.is_(None),
|
||
ChatGenerationTask.generation_mode.in_(["chatapi_async", "hot_opening_replicate", "shot_replicate"]),
|
||
ChatGenerationTask.status == "generating",
|
||
ChatGenerationTask.pipeline_stage.in_(
|
||
[
|
||
"queued",
|
||
"preparing",
|
||
"creating_provider_task",
|
||
"waiting_remote",
|
||
"polling",
|
||
"result_ready",
|
||
]
|
||
),
|
||
)
|
||
.order_by(ChatGenerationTask.updated_at.asc())
|
||
.limit(batch_size)
|
||
.with_for_update(skip_locked=True)
|
||
)
|
||
tasks = query_result.scalars().all()
|
||
if not tasks:
|
||
break
|
||
|
||
progressed_this_round = 0
|
||
for task in tasks:
|
||
if task.id in checked_ids:
|
||
continue
|
||
action = await recover_one_generation_task(
|
||
db,
|
||
task,
|
||
payload=None,
|
||
source="startup_db",
|
||
)
|
||
results[action] = results.get(action, 0) + 1
|
||
checked_ids.add(task.id)
|
||
total_db_checked += 1
|
||
progressed_this_round += 1
|
||
|
||
if len(tasks) < batch_size or progressed_this_round <= 0:
|
||
break
|
||
|
||
# 下载阶段单独跑 DB fallback。
|
||
download_result = await recover_download_tasks_once(db)
|
||
return {
|
||
"checked": len(checked_ids),
|
||
"db_checked": total_db_checked,
|
||
"results": results,
|
||
"download_recovery": download_result,
|
||
}
|