vent_agent_2.0/api/concurrency.py

"""长任务并发闸。

超限**直接拒绝（429），不排队**。这是刻意的取舍：LLM 长任务动辄跑几分钟，
排队意味着用户盯着一个不动的界面等一个不知道多久才开始的请求；而 429 带
``retry_after`` 能让前端明确告诉用户「稍后再试」，体验反而更可控。

两道上限同时生效：
- 全局（``TFAGENTS_MAX_CONCURRENT_TASKS``）保护下游 provider 与自身资源
- 每用户（``TFAGENTS_MAX_USER_TASKS``）防止单个用户占满全局额度

计数是**进程内**的，所以这道闸只在 ``workers=1`` 下有效；多 worker 时每个
进程各算各的，实际并发会翻倍。多进程部署需要把计数挪到 Redis/MySQL。
"""

from __future__ import annotations

import threading
import time
from contextlib import contextmanager
from typing import Any

import config
from obs import get_logger

logger = get_logger(__name__)


class TooManyTasksError(RuntimeError):
    """并发已满。携带建议的重试间隔，供 429 响应的 retry_after 使用。"""

    def __init__(self, message: str, scope: str, retry_after: int = 10) -> None:
        self.scope = scope
        self.retry_after = retry_after
        super().__init__(message)


class TaskWallTimeoutError(RuntimeError):
    """任务超过墙钟上限。"""


class SessionBusyError(RuntimeError):
    """同一会话正在执行或修改。"""


class SessionMissingError(RuntimeError):
    """取得执行锁时会话已被删除或不属于当前用户。"""


_sessions: set[str] = set()


class SessionLease:
    def __init__(self, session_id: str) -> None:
        self.session_id = session_id
        self.released = False

    def release(self) -> None:
        with _state_lock:
            if not self.released:
                _sessions.discard(self.session_id)
                self.released = True


def acquire_session(session_id: str) -> SessionLease:
    with _state_lock:
        if session_id in _sessions:
            raise SessionBusyError("该会话正在执行，请等待结束后重试")
        _sessions.add(session_id)
    return SessionLease(session_id)


@contextmanager
def session_guard(session_id: str):
    lease = acquire_session(session_id)
    try:
        yield
    finally:
        lease.release()


class TaskSlot:
    """一个已占用的并发位。

    必须成对使用：``acquire_task_slot()`` 拿到后，在 finally 里 ``release()``。
    重复 release 安全（幂等）——SSE 生成器有多条退出路径，很容易重复调用。
    """

    __slots__ = ("user_name", "acquired_at", "_released", "_lock")

    def __init__(self, user_name: str) -> None:
        self.user_name = user_name
        self.acquired_at = time.time()
        self._released = False
        self._lock = threading.Lock()

    @property
    def held_seconds(self) -> float:
        return time.time() - self.acquired_at

    @property
    def released(self) -> bool:
        return self._released

    def release(self) -> bool:
        """释放占位。返回 True 表示本次调用真正释放了（首次调用）。"""
        with self._lock:
            if self._released:
                return False
            self._released = True
        with _state_lock:
            _global_count[0] = max(_global_count[0] - 1, 0)
            user_count = _per_user_count.get(self.user_name, 0)
            if user_count <= 1:
                _per_user_count.pop(self.user_name, None)
            else:
                _per_user_count[self.user_name] = user_count - 1
        logger.info(
            "释放并发位 user=%s 持有时长=%.1fs 剩余全局=%d 剩余用户=%d",
            self.user_name, self.held_seconds, _global_count[0],
            _per_user_count.get(self.user_name, 0),
        )
        return True


# 进程内计数。用 list/dict 而不是裸 int，避免到处写 global
_state_lock = threading.Lock()
_global_count = [0]
_per_user_count: dict[str, int] = {}


def acquire_task_slot(user_name: str) -> TaskSlot:
    """尝试占一个并发位；超限抛 :class:`TooManyTasksError`。"""
    user = user_name or "default"
    with _state_lock:
        if _global_count[0] >= config.MAX_CONCURRENT_TASKS:
            raise TooManyTasksError(
                f"服务当前并发任务已达上限（{config.MAX_CONCURRENT_TASKS}），请稍后重试",
                scope="global",
            )
        current_user = _per_user_count.get(user, 0)
        if current_user >= config.MAX_USER_TASKS:
            raise TooManyTasksError(
                f"您的并发任务已达上限（{config.MAX_USER_TASKS}），请等待当前任务结束后重试",
                scope="user",
            )
        _global_count[0] += 1
        _per_user_count[user] = current_user + 1
        logger.info(
            "占用并发位 user=%s 全局=%d/%d 用户=%d/%d",
            user, _global_count[0], config.MAX_CONCURRENT_TASKS,
            _per_user_count[user], config.MAX_USER_TASKS,
        )
    return TaskSlot(user)


def stats() -> dict[str, Any]:
    """当前并发占用快照，供 /health 展示。"""
    with _state_lock:
        return {
            "global_active": _global_count[0],
            "global_limit": config.MAX_CONCURRENT_TASKS,
            "user_limit": config.MAX_USER_TASKS,
            "users": dict(_per_user_count),
            "wall_timeout_seconds": config.TASK_WALL_TIMEOUT,
        }


def reset_for_tests() -> None:
    """仅供测试：清零计数。"""
    with _state_lock:
        _global_count[0] = 0
        _per_user_count.clear()
        _sessions.clear()