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))