149 lines
5.5 KiB
Python
149 lines
5.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.config import settings
|
|
from app.enums.generation_status import GenerationRecordPipelineStage, GenerationStatus
|
|
from app.enums.generation_task import GenerationOwnerType
|
|
from app.models.generation_record import GenerationRecord
|
|
from app.services.generation.pipeline.owner_service import GenerationOwnerRef
|
|
from app.services.redis_registry_service import ensure_aware_utc
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class GenerationRecordRecoveryCursor:
|
|
owner_id: str
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class GenerationRecordRecoveryBatch:
|
|
create: list[GenerationOwnerRef]
|
|
poll: list[GenerationOwnerRef]
|
|
download: list[GenerationOwnerRef]
|
|
next_cursor: GenerationRecordRecoveryCursor | None
|
|
|
|
|
|
async def find_generation_record_recovery_batch(
|
|
db: AsyncSession,
|
|
*,
|
|
limit: int,
|
|
cursor: GenerationRecordRecoveryCursor | None = None,
|
|
) -> GenerationRecordRecoveryBatch:
|
|
"""按稳定游标读取恢复分流所需列，避免大字段加载和 offset 扫描。"""
|
|
stages = {
|
|
GenerationRecordPipelineStage.QUEUED.value,
|
|
GenerationRecordPipelineStage.PREPARING.value,
|
|
GenerationRecordPipelineStage.CREATING_PROVIDER_TASK.value,
|
|
GenerationRecordPipelineStage.WAITING_REMOTE.value,
|
|
GenerationRecordPipelineStage.POLLING.value,
|
|
GenerationRecordPipelineStage.RESULT_READY.value,
|
|
GenerationRecordPipelineStage.DOWNLOAD_QUEUED.value,
|
|
GenerationRecordPipelineStage.DOWNLOADING.value,
|
|
GenerationRecordPipelineStage.RETRY_WAITING.value,
|
|
}
|
|
page_size = max(1, int(limit))
|
|
now = datetime.now(timezone.utc)
|
|
queue_timeout = timedelta(
|
|
seconds=max(1, int(settings.DOWNLOAD_TASK_QUEUE_TIMEOUT_SECONDS or 300))
|
|
)
|
|
query = select(
|
|
GenerationRecord.id,
|
|
GenerationRecord.generation_attempt_no,
|
|
GenerationRecord.seedance_task_id,
|
|
GenerationRecord.remote_result_url,
|
|
GenerationRecord.pipeline_stage,
|
|
GenerationRecord.provider_create_lease_until,
|
|
GenerationRecord.next_poll_at,
|
|
GenerationRecord.poll_lease_until,
|
|
GenerationRecord.download_enqueued_at,
|
|
GenerationRecord.download_lease_until,
|
|
GenerationRecord.download_next_retry_at,
|
|
).where(
|
|
GenerationRecord.deleted_at.is_(None),
|
|
GenerationRecord.status == GenerationStatus.generating.value,
|
|
GenerationRecord.pipeline_stage.in_(stages),
|
|
)
|
|
if cursor is not None:
|
|
query = query.where(GenerationRecord.id > cursor.owner_id)
|
|
result = await db.execute(
|
|
query.order_by(GenerationRecord.id.asc()).limit(page_size)
|
|
)
|
|
rows = list(result.all())
|
|
|
|
create: list[GenerationOwnerRef] = []
|
|
poll: list[GenerationOwnerRef] = []
|
|
download: list[GenerationOwnerRef] = []
|
|
for (
|
|
owner_id,
|
|
attempt_no,
|
|
provider_task_id,
|
|
remote_result_url,
|
|
pipeline_stage,
|
|
provider_create_lease_until,
|
|
next_poll_at,
|
|
poll_lease_until,
|
|
download_enqueued_at,
|
|
download_lease_until,
|
|
download_next_retry_at,
|
|
) in rows:
|
|
ref = GenerationOwnerRef(
|
|
GenerationOwnerType.GENERATION_RECORD.value,
|
|
str(owner_id),
|
|
int(attempt_no or 1),
|
|
)
|
|
stage = str(pipeline_stage or "")
|
|
if str(remote_result_url or "").strip():
|
|
if stage == GenerationRecordPipelineStage.RESULT_READY.value:
|
|
download.append(ref)
|
|
elif stage == GenerationRecordPipelineStage.DOWNLOAD_QUEUED.value:
|
|
checked_enqueued_at = ensure_aware_utc(download_enqueued_at)
|
|
if (
|
|
checked_enqueued_at is None
|
|
or checked_enqueued_at + queue_timeout <= now
|
|
):
|
|
download.append(ref)
|
|
elif stage == GenerationRecordPipelineStage.DOWNLOADING.value:
|
|
if (
|
|
ensure_aware_utc(download_lease_until) is None
|
|
or ensure_aware_utc(download_lease_until) <= now
|
|
):
|
|
download.append(ref)
|
|
elif stage == GenerationRecordPipelineStage.RETRY_WAITING.value:
|
|
if (
|
|
ensure_aware_utc(download_next_retry_at) is None
|
|
or ensure_aware_utc(download_next_retry_at) <= now
|
|
):
|
|
download.append(ref)
|
|
elif str(provider_task_id or "").strip():
|
|
checked_next_poll_at = ensure_aware_utc(next_poll_at)
|
|
checked_poll_lease_until = ensure_aware_utc(poll_lease_until)
|
|
if (
|
|
(checked_next_poll_at is None or checked_next_poll_at <= now)
|
|
and (
|
|
checked_poll_lease_until is None
|
|
or checked_poll_lease_until <= now
|
|
)
|
|
):
|
|
poll.append(ref)
|
|
else:
|
|
if (
|
|
ensure_aware_utc(provider_create_lease_until) is None
|
|
or ensure_aware_utc(provider_create_lease_until) <= now
|
|
):
|
|
create.append(ref)
|
|
|
|
next_cursor = None
|
|
if len(rows) == page_size:
|
|
last = rows[-1]
|
|
next_cursor = GenerationRecordRecoveryCursor(owner_id=str(last.id))
|
|
return GenerationRecordRecoveryBatch(
|
|
create=create,
|
|
poll=poll,
|
|
download=download,
|
|
next_cursor=next_cursor,
|
|
)
|