Aether/_deprecated_py_src/services/task/polling/task_poller.py

"""
Task poller (Phase2)

Provides a generic polling skeleton for async tasks.
Currently wired with a video poller adapter.

优化：HTTP 请求期间不持有数据库连接，避免阻塞其他请求。
采用三阶段处理：准备数据 -> HTTP 请求 -> 更新数据库。
"""

from __future__ import annotations

import asyncio
from datetime import datetime, timezone
from typing import Any, Protocol, runtime_checkable
from uuid import uuid4

from sqlalchemy.orm import Session

from src.core.api_format.conversion.internal_video import InternalVideoPollResult
from src.core.logger import logger
from src.database import create_session
from src.services.system.scheduler import get_scheduler
from src.services.task.video.poller_adapter import VideoPollContext, VideoTaskPollerAdapter


@runtime_checkable
class TaskPollerAdapter(Protocol):
    task_type: str

    # scheduler
    job_id: str
    job_name: str
    interval_seconds: int

    # distributed lock (optional, best-effort)
    lock_key: str
    lock_ttl: int

    # execution
    batch_size: int
    concurrency: int
    consecutive_failure_alert_threshold: int

    def list_due_task_ids(self, db: Session, *, now: datetime, limit: int) -> list[str]: ...

    def get_task(self, db: Session, task_id: str) -> Any | None: ...

    # 分阶段处理方法（推荐使用）
    async def prepare_poll_context(
        self, db: Session, task: Any
    ) -> Any: ...  # Returns context or error result

    async def poll_task_http(self, ctx: Any) -> Any: ...  # Returns poll result

    async def update_task_after_poll(
        self,
        task_id: str,
        result: Any,
        ctx: Any,
        redis_client: Any | None,
        error_exception: Exception | None = None,
    ) -> None: ...

    # 旧版方法（保留兼容性）
    async def poll_single_task(
        self, db: Session, task: Any, *, redis_client: Any | None
    ) -> None: ...

    def sanitize_error_message(self, message: str) -> str: ...


class TaskPollerService:
    """Generic background poller for async tasks."""

    def __init__(self, adapter: TaskPollerAdapter) -> None:
        self.adapter = adapter
        self._lock = asyncio.Lock()
        self.redis: Any | None = None
        self._semaphore: asyncio.Semaphore | None = None
        self._consecutive_failures = 0

    async def start(self) -> None:
        if self._semaphore is None:
            self._semaphore = asyncio.Semaphore(self.adapter.concurrency)

        # lazy import to avoid redis hard dependency in local runs
        from src.clients.redis_client import get_redis_client

        if self.redis is None:
            self.redis = await get_redis_client(require_redis=False)

        scheduler = get_scheduler()
        scheduler.add_interval_job(
            self.poll_pending_tasks,
            seconds=self.adapter.interval_seconds,
            job_id=self.adapter.job_id,
            name=self.adapter.job_name,
        )

    async def stop(self) -> None:
        scheduler = get_scheduler()
        scheduler.remove_job(self.adapter.job_id)

    async def poll_pending_tasks(self) -> None:
        try:
            await self._do_poll()
        except asyncio.CancelledError:
            logger.debug("[{}] poll_pending_tasks cancelled (shutdown?)", self.adapter.task_type)
            return

    async def _do_poll(self) -> None:
        async with self._lock:
            token = await self._acquire_redis_lock()
            if token is None:
                return

            try:
                with create_session() as db:
                    now = datetime.now(timezone.utc)
                    task_ids = self.adapter.list_due_task_ids(
                        db, now=now, limit=self.adapter.batch_size
                    )

                if not task_ids:
                    self._consecutive_failures = 0
                    return

                poll_results: list[bool] = []

                if self._semaphore is None:
                    self._semaphore = asyncio.Semaphore(self.adapter.concurrency)
                semaphore = self._semaphore

                async def poll_with_semaphore(task_id: str) -> None:
                    async with semaphore:
                        try:
                            # ========== 阶段 1：准备数据（短暂持有连接）==========
                            with create_session() as task_db:
                                task_obj = self.adapter.get_task(task_db, task_id)
                                if not task_obj:
                                    logger.warning(
                                        "[{}] Task {} disappeared during poll",
                                        self.adapter.task_type,
                                        task_id,
                                    )
                                    poll_results.append(True)
                                    return

                                ctx_or_result = await self.adapter.prepare_poll_context(
                                    task_db, task_obj
                                )

                            # 检查是否是错误结果（而非上下文）
                            if isinstance(ctx_or_result, InternalVideoPollResult):
                                # 准备阶段就失败了，直接更新任务状态
                                await self.adapter.update_task_after_poll(
                                    task_id=task_id,
                                    result=ctx_or_result,
                                    ctx=None,  # type: ignore[arg-type]
                                    redis_client=self.redis,
                                )
                                poll_results.append(True)
                                return

                            ctx: VideoPollContext = ctx_or_result

                            # ========== 阶段 2：HTTP 请求（不持有数据库连接）==========
                            error_exception: Exception | None = None
                            try:
                                result = await self.adapter.poll_task_http(ctx)
                            except Exception as http_exc:
                                # HTTP 请求失败，记录异常以便后续处理
                                error_exception = http_exc
                                result = InternalVideoPollResult(
                                    status=None,  # type: ignore[arg-type]
                                    error_message=str(http_exc),
                                )

                            # ========== 阶段 3：更新数据库（获取新连接）==========
                            await self.adapter.update_task_after_poll(
                                task_id=task_id,
                                result=result,
                                ctx=ctx,
                                redis_client=self.redis,
                                error_exception=error_exception,
                            )
                            poll_results.append(True)
                        except Exception as exc:
                            logger.exception(
                                "[{}] Unexpected error polling task {}: {}",
                                self.adapter.task_type,
                                task_id,
                                self.adapter.sanitize_error_message(str(exc)),
                            )
                            poll_results.append(False)

                async with asyncio.TaskGroup() as tg:
                    for tid in task_ids:
                        tg.create_task(poll_with_semaphore(tid))

                batch_failures = sum(1 for r in poll_results if r is False)
                if batch_failures == len(task_ids):
                    self._consecutive_failures += 1
                    if (
                        self._consecutive_failures
                        >= self.adapter.consecutive_failure_alert_threshold
                    ):
                        logger.error(
                            "[ALERT] {} poller: {} consecutive batches failed.",
                            self.adapter.task_type,
                            self._consecutive_failures,
                        )
                else:
                    self._consecutive_failures = 0
            finally:
                await self._release_redis_lock(token)

    async def _acquire_redis_lock(self) -> str | None:
        if not self.redis:
            return "no_redis"
        token = str(uuid4())
        try:
            acquired = await self.redis.set(
                self.adapter.lock_key, token, nx=True, ex=self.adapter.lock_ttl
            )
        except Exception as exc:
            logger.warning(
                "[{}] Redis lock acquire failed (best-effort skip): {}",
                self.adapter.task_type,
                exc,
            )
            return "no_redis"
        return token if acquired else None

    async def _release_redis_lock(self, token: str) -> None:
        if not self.redis or token == "no_redis":
            return
        script = """
        if redis.call('GET', KEYS[1]) == ARGV[1] then
            return redis.call('DEL', KEYS[1])
        end
        return 0
        """
        try:
            await self.redis.eval(script, 1, self.adapter.lock_key, token)
        except Exception as exc:
            logger.warning(
                "[{}] Redis lock release failed (will expire via TTL): {}",
                self.adapter.task_type,
                exc,
            )


_task_poller: TaskPollerService | None = None


def get_task_poller() -> TaskPollerService:
    global _task_poller
    if _task_poller is None:
        _task_poller = TaskPollerService(VideoTaskPollerAdapter())
    return _task_poller