mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
263 lines
9.6 KiB
Python
263 lines
9.6 KiB
Python
"""
|
||
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
|