video-gen/video-gen-api/app/tasks/private_portrait_asset_tasks.py

from __future__ import annotations

import logging
import uuid
from datetime import datetime, timedelta, timezone
from typing import Any

from sqlalchemy import select

from app.config import settings
from app.enums.celery_queue import CeleryQueue, CeleryTaskName
from app.enums.celery_runtime import CeleryRuntimeDomain
from app.enums.private_portrait import (
    PRIVATE_PORTRAIT_ASSET_POLL_BATCH_SIZE,
    PRIVATE_PORTRAIT_REMOTE_DELETE_RECOVERY_BATCH_SIZE,
    PrivatePortraitAssetStatus,
    PrivatePortraitEventSource,
    PrivatePortraitEventType,
    PrivatePortraitRemoteDeleteStatus,
)
from app.models import async_session
from app.models.private_portrait import PrivatePortraitAsset, PrivatePortraitAssetGroup
from app.services.celery_runtime.recovery_service import guard_periodic_recovery
from app.services.celery_runtime.runtime_service import CeleryRuntimeLease, RuntimeIdentity
from app.services.operation_log_service import log_operation_error, log_operation_event
from app.services.private_portrait.asset_service import (
    DOMAIN,
    delete_asset_group_remote,
    delete_asset_remote,
    sync_asset_status,
)
from app.services.redis_registry_service import RedisExecutionLockLease
from app.tasks.async_runner import run_async
from app.tasks.celery_app import celery_app

logger = logging.getLogger(__name__)
QUEUE = CeleryQueue.GEN_PRIVATE_PORTRAIT.value


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _retry_countdown(retries: int) -> int:
    return min(300, 30 * (2 ** max(0, retries)))


async def _rollback_and_reraise(
    db,
    *,
    event_type: str,
    exc: BaseException,
    detail: dict[str, Any] | None = None,
    **kwargs: Any,
):
    await db.rollback()
    log_operation_error(
        domain=DOMAIN,
        event_type=event_type,
        source=PrivatePortraitEventSource.CELERY.value,
        exc=exc,
        detail=detail,
        **kwargs,
    )
    raise exc


async def _acquire_runtime(
    *,
    domain: str,
    owner_type: str,
    owner_id: str,
    task_name: str,
    hash_key: str,
    zset_key: str,
    lock_prefix: str,
) -> CeleryRuntimeLease | None:
    token = uuid.uuid4().hex
    return await CeleryRuntimeLease.acquire(
        identity=RuntimeIdentity(
            domain=domain,
            owner_type=owner_type,
            owner_id=owner_id,
            attempt_no=1,
            task_name=task_name,
            queue=QUEUE,
        ),
        lock_key=f"{lock_prefix}:{owner_type}:{owner_id}:attempt:1",
        hash_key=hash_key,
        zset_key=zset_key,
        token=token,
        ttl_seconds=max(60, int(settings.PRIVATE_PORTRAIT_RUNTIME_LOCK_TTL_SECONDS or 180)),
        heartbeat_interval_seconds=max(10, int(settings.PRIVATE_PORTRAIT_RUNTIME_HEARTBEAT_SECONDS or 30)),
        pipeline_stage="processing",
    )


async def _run_poll_asset(asset_id: str) -> None:
    lease = await _acquire_runtime(
        domain=CeleryRuntimeDomain.PRIVATE_PORTRAIT_POLL.value,
        owner_type="asset",
        owner_id=asset_id,
        task_name=CeleryTaskName.PRIVATE_PORTRAIT_POLL_ASSET.value,
        hash_key=settings.PRIVATE_PORTRAIT_POLL_ACTIVE_REDIS_HASH_KEY,
        zset_key=settings.PRIVATE_PORTRAIT_POLL_ACTIVE_REDIS_ZSET_KEY,
        lock_prefix=settings.PRIVATE_PORTRAIT_POLL_LOCK_KEY_PREFIX,
    )
    if lease is None:
        return
    try:
        async with async_session() as db:
            try:
                await sync_asset_status(
                    db,
                    user_id=None,
                    asset_id=asset_id,
                    execution_guard=lease.ensure_owned,
                )
                await lease.ensure_owned()
                await db.commit()
            except Exception as exc:
                logger.exception("poll private portrait asset failed: %s", asset_id)
                await _rollback_and_reraise(
                    db,
                    event_type=PrivatePortraitEventType.ASSET_POLL_FAILED.value,
                    exc=exc,
                    asset_id=asset_id,
                    detail={"celery_task": CeleryTaskName.PRIVATE_PORTRAIT_POLL_ASSET.value},
                )
    finally:
        await lease.close()


async def _dispatch_due_assets() -> int:
    barrier = await guard_periodic_recovery()
    if barrier is not None:
        return 0
    lock = await RedisExecutionLockLease.acquire(
        lock_key=settings.PRIVATE_PORTRAIT_DISPATCH_LOCK_KEY,
        ttl_seconds=55,
        renew_interval_seconds=20,
        log_context="private_portrait_poll_dispatch",
    )
    if lock is None:
        return 0
    async with lock:
        async with async_session() as db:
            now = _now()
            rows = await db.execute(
                select(PrivatePortraitAsset)
                .where(
                    PrivatePortraitAsset.deleted_at.is_(None),
                    PrivatePortraitAsset.status == PrivatePortraitAssetStatus.PROCESSING.value,
                    PrivatePortraitAsset.next_poll_at.is_not(None),
                    PrivatePortraitAsset.next_poll_at <= now,
                )
                .order_by(PrivatePortraitAsset.next_poll_at.asc(), PrivatePortraitAsset.id.asc())
                .limit(PRIVATE_PORTRAIT_ASSET_POLL_BATCH_SIZE)
                .with_for_update(skip_locked=True)
            )
            assets = list(rows.scalars().all())
            dispatches: list[tuple[str, int]] = []
            queue_hold_until = now + timedelta(seconds=120)
            for asset in assets:
                dispatches.append((str(asset.id), int(asset.poll_count or 0) + 1))
                asset.next_poll_at = queue_hold_until
            await db.commit()

        for asset_id, poll_no in dispatches:
            poll_private_portrait_asset_status.apply_async(
                args=[asset_id],
                queue=QUEUE,
                countdown=0,
                task_id=f"private-portrait-poll:{asset_id}:attempt:{poll_no}",
            )
        log_operation_event(
            domain=DOMAIN,
            event_type=PrivatePortraitEventType.SYNC_DUE_ASSETS_DONE.value,
            event_status="success",
            source=PrivatePortraitEventSource.CELERY.value,
            detail={"matched_count": len(dispatches), "dispatched_count": len(dispatches)},
        )
        return len(dispatches)


async def _run_delete_asset(asset_id: str) -> None:
    lease = await _acquire_runtime(
        domain=CeleryRuntimeDomain.PRIVATE_PORTRAIT_DELETE.value,
        owner_type="asset",
        owner_id=asset_id,
        task_name=CeleryTaskName.PRIVATE_PORTRAIT_DELETE_ASSET.value,
        hash_key=settings.PRIVATE_PORTRAIT_DELETE_ACTIVE_REDIS_HASH_KEY,
        zset_key=settings.PRIVATE_PORTRAIT_DELETE_ACTIVE_REDIS_ZSET_KEY,
        lock_prefix=settings.PRIVATE_PORTRAIT_DELETE_LOCK_KEY_PREFIX,
    )
    if lease is None:
        return
    try:
        async with async_session() as db:
            try:
                await delete_asset_remote(
                    db,
                    asset_id=asset_id,
                    execution_guard=lease.ensure_owned,
                )
                await lease.ensure_owned()
                await db.commit()
            except Exception as exc:
                await _rollback_and_reraise(
                    db,
                    event_type=PrivatePortraitEventType.ASSET_DELETE_REMOTE_FAILED.value,
                    exc=exc,
                    asset_id=asset_id,
                    detail={"celery_task": CeleryTaskName.PRIVATE_PORTRAIT_DELETE_ASSET.value},
                )
    finally:
        await lease.close()


async def _run_delete_group(group_id: str) -> None:
    lease = await _acquire_runtime(
        domain=CeleryRuntimeDomain.PRIVATE_PORTRAIT_DELETE.value,
        owner_type="group",
        owner_id=group_id,
        task_name=CeleryTaskName.PRIVATE_PORTRAIT_DELETE_GROUP.value,
        hash_key=settings.PRIVATE_PORTRAIT_DELETE_ACTIVE_REDIS_HASH_KEY,
        zset_key=settings.PRIVATE_PORTRAIT_DELETE_ACTIVE_REDIS_ZSET_KEY,
        lock_prefix=settings.PRIVATE_PORTRAIT_DELETE_LOCK_KEY_PREFIX,
    )
    if lease is None:
        return
    try:
        async with async_session() as db:
            try:
                await delete_asset_group_remote(
                    db,
                    group_id=group_id,
                    execution_guard=lease.ensure_owned,
                )
                await lease.ensure_owned()
                await db.commit()
            except Exception as exc:
                await _rollback_and_reraise(
                    db,
                    event_type=PrivatePortraitEventType.PROJECT_DELETE_REMOTE_FAILED.value,
                    exc=exc,
                    group_id=group_id,
                    detail={"celery_task": CeleryTaskName.PRIVATE_PORTRAIT_DELETE_GROUP.value},
                )
    finally:
        await lease.close()


async def _dispatch_project_deletes(project_id: str) -> int:
    lease = await _acquire_runtime(
        domain=CeleryRuntimeDomain.PRIVATE_PORTRAIT_DELETE.value,
        owner_type="project",
        owner_id=project_id,
        task_name=CeleryTaskName.PRIVATE_PORTRAIT_DELETE_PROJECT.value,
        hash_key=settings.PRIVATE_PORTRAIT_DELETE_ACTIVE_REDIS_HASH_KEY,
        zset_key=settings.PRIVATE_PORTRAIT_DELETE_ACTIVE_REDIS_ZSET_KEY,
        lock_prefix=settings.PRIVATE_PORTRAIT_DELETE_LOCK_KEY_PREFIX,
    )
    if lease is None:
        return 0
    try:
        async with async_session() as db:
            asset_rows = await db.execute(select(PrivatePortraitAsset.id).where(PrivatePortraitAsset.project_id == project_id))
            group_rows = await db.execute(select(PrivatePortraitAssetGroup.id).where(PrivatePortraitAssetGroup.project_id == project_id))
            asset_ids = [str(value) for value in asset_rows.scalars().all()]
            group_ids = [str(value) for value in group_rows.scalars().all()]
            await db.rollback()
        await lease.ensure_owned()
        for asset_id in asset_ids:
            delete_private_portrait_asset_remote.apply_async(
                args=[asset_id], queue=QUEUE, task_id=f"private-portrait-delete-asset:{asset_id}"
            )
        for group_id in group_ids:
            delete_private_portrait_group_remote.apply_async(
                args=[group_id], queue=QUEUE, task_id=f"private-portrait-delete-group:{group_id}"
            )
        return len(asset_ids) + len(group_ids)
    finally:
        await lease.close()


async def _dispatch_remote_delete_recovery() -> dict[str, int]:
    barrier = await guard_periodic_recovery()
    if barrier is not None:
        return {"asset_count": 0, "group_count": 0, "total_count": 0}
    lock = await RedisExecutionLockLease.acquire(
        lock_key=settings.PRIVATE_PORTRAIT_DELETE_RECOVERY_LOCK_KEY,
        ttl_seconds=240,
        renew_interval_seconds=30,
        log_context="private_portrait_delete_recovery",
    )
    if lock is None:
        return {"asset_count": 0, "group_count": 0, "total_count": 0}
    async with lock:
        statuses = [PrivatePortraitRemoteDeleteStatus.PENDING.value, PrivatePortraitRemoteDeleteStatus.FAILED.value]
        async with async_session() as db:
            asset_rows = await db.execute(
                select(PrivatePortraitAsset.id)
                .where(PrivatePortraitAsset.remote_delete_status.in_(statuses))
                .order_by(PrivatePortraitAsset.updated_at.asc(), PrivatePortraitAsset.id.asc())
                .limit(PRIVATE_PORTRAIT_REMOTE_DELETE_RECOVERY_BATCH_SIZE)
            )
            asset_ids = [str(value) for value in asset_rows.scalars().all()]
            remaining = max(0, PRIVATE_PORTRAIT_REMOTE_DELETE_RECOVERY_BATCH_SIZE - len(asset_ids))
            group_ids: list[str] = []
            if remaining:
                group_rows = await db.execute(
                    select(PrivatePortraitAssetGroup.id)
                    .where(PrivatePortraitAssetGroup.remote_delete_status.in_(statuses))
                    .order_by(PrivatePortraitAssetGroup.updated_at.asc(), PrivatePortraitAssetGroup.id.asc())
                    .limit(remaining)
                )
                group_ids = [str(value) for value in group_rows.scalars().all()]
            await db.rollback()

        for asset_id in asset_ids:
            delete_private_portrait_asset_remote.apply_async(
                args=[asset_id], queue=QUEUE, task_id=f"private-portrait-delete-asset:{asset_id}"
            )
        for group_id in group_ids:
            delete_private_portrait_group_remote.apply_async(
                args=[group_id], queue=QUEUE, task_id=f"private-portrait-delete-group:{group_id}"
            )
        return {"asset_count": len(asset_ids), "group_count": len(group_ids), "total_count": len(asset_ids) + len(group_ids)}


@celery_app.task(name=CeleryTaskName.PRIVATE_PORTRAIT_POLL_ASSET.value, queue=QUEUE, bind=True, max_retries=5, default_retry_delay=30)
def poll_private_portrait_asset_status(self, asset_id: str) -> None:
    try:
        return run_async(_run_poll_asset(asset_id))
    except Exception as exc:
        raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))


@celery_app.task(name=CeleryTaskName.PRIVATE_PORTRAIT_SYNC_DUE_ASSETS.value, queue=QUEUE, bind=True, max_retries=3, default_retry_delay=60)
def sync_private_portrait_due_assets(self) -> int:
    try:
        return run_async(_dispatch_due_assets())
    except Exception as exc:
        raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))


@celery_app.task(name=CeleryTaskName.PRIVATE_PORTRAIT_DELETE_ASSET.value, queue=QUEUE, bind=True, max_retries=3, default_retry_delay=60)
def delete_private_portrait_asset_remote(self, asset_id: str) -> None:
    try:
        return run_async(_run_delete_asset(asset_id))
    except Exception as exc:
        raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))


@celery_app.task(name=CeleryTaskName.PRIVATE_PORTRAIT_DELETE_GROUP.value, queue=QUEUE, bind=True, max_retries=3, default_retry_delay=60)
def delete_private_portrait_group_remote(self, group_id: str) -> None:
    try:
        return run_async(_run_delete_group(group_id))
    except Exception as exc:
        raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))


@celery_app.task(name=CeleryTaskName.PRIVATE_PORTRAIT_DELETE_PROJECT.value, queue=QUEUE, bind=True, max_retries=3, default_retry_delay=60)
def delete_private_portrait_project_remote(self, project_id: str) -> int:
    try:
        return run_async(_dispatch_project_deletes(project_id))
    except Exception as exc:
        raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))


@celery_app.task(name=CeleryTaskName.PRIVATE_PORTRAIT_RECOVER_REMOTE_DELETES.value, queue=QUEUE, bind=True, max_retries=3, default_retry_delay=60)
def recover_private_portrait_remote_deletes(self) -> dict[str, int]:
    try:
        return run_async(_dispatch_remote_delete_recovery())
    except Exception as exc:
        raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))