超大用户多租户 LLM 并发解决方案

from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional
from enum import Enum
import threading


class CredentialStatus(Enum):
    ACTIVE = "active"           # 正常运行，可分配
    BUSY = "busy"               # 达到并发上限，暂停分配
    CIRCUIT_BROKEN = "circuit_broken"  # 熔断中，不分配
    RECOVERING = "recovering"   # 恢复中，逐步放量
    RETIRED = "retired"         # 已下线，不再使用


@dataclass
class Credential:
    """
    单个 LLM API 凭证的抽象封装。
    包含认证信息、状态信息、运行时统计三大类属性。
    """
    # === 认证信息（创建时确定，运行时不变）===
    api_key: str                          # API Key 明文（仅在创建时传入）
    model: str                            # 对应模型名称，如 "MiniMax-M2.7-highspeed"
    base_url: str                         # API 端点，如 "https://api.minimax.chat"
    max_concurrency: int = 20             # 该 credential 的最大并发上限
    weight: float = 1.0                   # 权重因子，用于加权负载均衡

    # === 运行时状态（线程安全访问）===
    current_load: int = 0                 # 当前正在使用的数量（原子操作）
    status: CredentialStatus = CredentialStatus.ACTIVE
    last_used: Optional[datetime] = None  # 上一次使用时间
    acquired_count: int = 0               # 累计获取次数
    released_count: int = 0               # 累计释放次数

    # === 健康相关统计 ===
    health_score: float = 1.0            # 健康分 0.0-1.0，由 HealthChecker 更新
    consecutive_errors: int = 0           # 连续错误次数（成功一次清零）
    consecutive_timeouts: int = 0          # 连续超时次数（成功一次清零）
    recent_success_count: int = 0         # 窗口内成功次数
    recent_failure_count: int = 0         # 窗口内失败次数
    recent_request_count: int = 0         # 窗口内总请求数
    recent_latency_sum: float = 0.0       # 窗口内延迟总和（ms）
    recent_latency_count: int =           # 窗口内请求计数

    # === 熔断恢复相关 ===
    circuit_break_time: Optional[datetime] = None  # 熔断触发时间
    recovery_attempts: int = 0            # 恢复尝试次数
    last_health_check: Optional[datetime] = None    # 上次健康检查时间

    # === 内部锁 ===
    _lock: threading.RLock = field(default_factory=threading.RLock)

    @property
    def recent_success_rate(self) -> float:
        """计算窗口内成功率"""
        total = self.recent_success_count + self.recent_failure_count
        if total == 0:
            return 1.0
        return self.recent_success_count / total

    @property
    def recent_avg_latency(self) -> float:
        """计算窗口内平均延迟"""
        if self.recent_latency_count == 0:
            return 0.0
        return self.recent_latency_sum / self.recent_latency_count

    @property
    def is_available(self) -> bool:
        """判断当前是否可分配"""
        return (
            self.status in (CredentialStatus.ACTIVE, CredentialStatus.RECOVERING)
            and self.current_load < self.max_concurrency
        )

    def try_acquire(self) -> bool:
        """
        尝试获取该 credential。
        成功返回 True 并递增 current_load；失败返回 False。
        """
        with self._lock:
            if not self.is_available:
                return False
            self.current_load += 1
            self.last_used = datetime.now()
            self.acquired_count += 1
            return True

    def release(self):
        """释放该 credential，递减 current_load"""
        with self._lock:
            if self.current_load > 0:
                self.current_load -= 1
            self.released_count += 1

    def record_success(self, latency_ms: float):
        """记录一次成功的请求"""
        with self._lock:
            self.recent_success_count += 1
            self.recent_request_count += 1
            self.recent_latency_sum += latency_ms
            self.recent_latency_count += 1
            self.consecutive_errors = 0
            self.consecutive_timeouts = 0

    def record_failure(self, is_timeout: bool = False):
        """记录一次失败的请求"""
        with self._lock:
            self.recent_failure_count += 1
            self.recent_request_count += 1
            self.consecutive_errors += 1
            if is_timeout:
                self.consecutive_timeouts += 1
            else:
                self.consecutive_timeouts = 0  # 非超时错误不计入连续超时

    def reset_window_stats(self):
        """重置窗口统计（由定时任务调用）"""
        with self._lock:
            self.recent_success_count = 0
            self.recent_failure_count = 0
            self.recent_request_count = 0
            self.recent_latency_sum = 0.0
            self.recent_latency_count = 0


class CredentialPool:
    """
    同一模型的多个 Credential 的统一调度器。
    提供 acquire / release / health_check / rotate 四个核心接口。
    """

    def __init__(
        self,
        model: str,
        base_url: str,
        strategy: str = "least_load",  # round_robin | least_load | weighted_random | sticky
        sticky_tenant_window_seconds: int = 300,
    ):
        self.model = model
        self.base_url = base_url
        self.strategy = strategy
        self.sticky_tenant_window = sticky_tenant_window_seconds

        self._credentials: list[Credential] = []
        self._lock = threading.RLock()
        self._round_robin_index = 0  # 轮询策略的当前位置

        # 粘性会话：tenant_id -> (credential, expire_time)
        self._sticky_map: dict[str, tuple[Credential, datetime]] = {}

    # === 核心管理接口 ===

    def register(self, api_key: str, max_concurrency: int = 20, weight: float = 1.0):
        """向池中注册一个新的 Credential"""
        cred = Credential(
            api_key=api_key,
            model=self.model,
            base_url=self.base_url,
            max_concurrency=max_concurrency,
            weight=weight,
        )
        with self._lock:
            self._credentials.append(cred)
        return cred

    def acquire(self, tenant_id: str = None, priority: int = 0) -> Optional[Credential]:
        """
        从池中获取一个可用的 Credential。

        参数:
            tenant_id: 租户 ID，用于粘性会话策略
            priority: 优先级（暂未使用，预留）

        返回:
            可用的 Credential 实例，或 None（池中没有可用 Credential）
        """
        with self._lock:
            candidates = [c for c in self._credentials if c.is_available]
            if not candidates:
                return None

            chosen = self._select_one(candidates, tenant_id)
            if chosen and chosen.try_acquire():
                return chosen
            return None

    def release(self, credential: Credential):
        """将 Credential 归还到池中"""
        credential.release()

    def health_check(self, credential: Credential) -> bool:
        """
        对指定 Credential 执行主动健康检查。
        返回 True 表示健康，False 表示应熔断。
        """
        # 构造一个最小化的 test 请求
        import time
        start = time.time()
        try:
            import httpx
            resp = httpx.get(
                credential.base_url.rstrip("/") + "/health",
                headers={"Authorization": f"Bearer {credential.api_key}"},
                timeout=5.0,
            )
            latency = (time.time() - start) * 1000
            if resp.status_code == 200:
                credential.record_success(latency)
                return True
            credential.record_failure()
            return False
        except Exception:
            credential.record_failure(is_timeout=True)
            return False

    def rotate(self) -> list[Credential]:
        """
        负载均衡策略触发时调用，返回当前所有 active credential。
        """
        with self._lock:
            # 清理过期粘性映射
            now = datetime.now()
            expired = [
                tid for tid, (_, exp) in self._sticky_map.items() if now > exp
            ]
            for tid in expired:
                del self._sticky_map[tid]

            return [
                c for c in self._credentials
                if c.status in (CredentialStatus.ACTIVE, CredentialStatus.RECOVERING)
            ]

    # === 负载均衡策略实现 ===

    def _select_one(
        self, candidates: list[Credential], tenant_id: str = None
    ) -> Optional[Credential]:
        if not candidates:
            return None

        if self.strategy == "round_robin":
            return self._round_robin(candidates)
        elif self.strategy == "least_load":
            return self._least_load(candidates)
        elif self.strategy == "weighted_random":
            return self._weighted_random(candidates)
        elif self.strategy == "sticky":
            return self._sticky(candidates, tenant_id)
        else:
            return self._least_load(candidates)

    def _round_robin(self, candidates: list[Credential]) -> Credential:
        """轮询策略：按注册顺序均匀分发请求"""
        n = len(candidates)
        idx = self._round_robin_index % n
        self._round_robin_index = idx + 1
        return candidates[idx]

    def _least_load(self, candidates: list[Credential]) -> Credential:
        """最少负载策略：优先选择当前 load 最低的"""
        return min(candidates, key=lambda c: (c.current_load, -c.health_score))

    def _weighted_random(self, candidates: list[Credential]) -> Credential:
        """加权随机策略：根据健康分和权重综合计算概率"""
        import random
        weights = [c.health_score * c.weight for c in candidates]
        total = sum(weights)
        if total <= 0:
            return random.choice(candidates)
        normalized = [w / total for w in weights]
        return random.choices(candidates, weights=normalized, k=1)[0]

    def _sticky(
        self, candidates: list[Credential], tenant_id: str = None
    ) -> Optional[Credential]:
        """粘性会话策略：同一租户尽量复用同一个 Credential"""
        if not tenant_id:
            return self._least_load(candidates)

        now = datetime.now()
        if tenant_id in self._sticky_map:
            cached_cred, expire_time = self._sticky_map[tenant_id]
            # 如果缓存的 credential 仍然可用，直接复用
            if cached_cred.is_available:
                return cached_cred
            else:
                # 已不可用，删除缓存，fallback 到 least_load
                del self._sticky_map[tenant_id]

        # 选择一个可用的 credential 并建立粘性关联
        chosen = self._least_load(candidates)
        self._sticky_map[tenant_id] = (
            chosen,
            now + timedelta(seconds=self.sticky_tenant_window),
        )
        return chosen

    def get_stats(self) -> dict:
        """获取池子的整体统计信息"""
        with self._lock:
            return {
                "model": self.model,
                "total_credentials": len(self._credentials),
                "active_count": sum(
                    1 for c in self._credentials
                    if c.status == CredentialStatus.ACTIVE
                ),
                "busy_count": sum(
                    1 for c in self._credentials
                    if c.status == CredentialStatus.BUSY
                ),
                "circuit_broken_count": sum(
                    1 for c in self._credentials
                    if c.status == CredentialStatus.CIRCUIT_BROKEN
                ),
                "total_load": sum(c.current_load for c in self._credentials),
                "total_acquired": sum(c.acquired_count for c in self._credentials),
            }

def _least_load(self, candidates):
    return min(candidates, key=lambda c: (c.current_load, -c.health_score))

def _weighted_random(self, candidates):
    weights = [c.health_score * c.weight for c in candidates]
    return random.choices(candidates, weights=weights, k=1)[0]

def _sticky(self, candidates, tenant_id):
    if tenant_id in self._sticky_map:
        cached_cred, expire_time = self._sticky_map[tenant_id]
        if cached_cred.is_available:
            return cached_cred
    chosen = self._least_load(candidates)
    self._sticky_map[tenant_id] = (chosen, now + timedelta(seconds=self.sticky_window))
    return chosen

class HealthChecker:
    def __init__(self, pool: CredentialPool, config: dict):
        self.pool = pool
        self.interval = config.get("interval_seconds", 30)
        self.timeout = config.get("test_request_timeout", 5.0)
        self._running = False

    async def run(self):
        """定时任务入口"""
        while self._running:
            for cred in self.pool.rotate():
                is_healthy = self.pool.health_check(cred)
                self._update_status(cred, is_healthy)
            await asyncio.sleep(self.interval)

    def _update_status(self, cred: Credential, is_healthy: bool):
        if is_healthy:
            if cred.status == CredentialStatus.RECOVERING:
                # 恢复中，逐步放量
                cred.recovery_attempts += 1
                # 恢复进度达到阈值后转为 ACTIVE
                if cred.recovery_attempts >= 3:
                    cred.status = CredentialStatus.ACTIVE
                    cred.recovery_attempts = 0
            elif cred.status == CredentialStatus.CIRCUIT_BROKEN:
                # 从熔断中恢复，先进入 RECOVERING 状态
                cred.status = CredentialStatus.RECOVERING
                cred.recovery_attempts = 0
        else:
            if cred.status in (CredentialStatus.ACTIVE, CredentialStatus.RECOVERING):
                # 触发熔断
                cred.status = CredentialStatus.CIRCUIT_BROKEN
                cred.circuit_break_time = datetime.now()

class HealthChecker:
    def should_circuit_break(self, cred: Credential) -> bool:
        """
        连续3次超时 或 成功率<80% → 熔断
        """
        # 条件1：连续超时达到阈值
        if cred.consecutive_timeouts >= 3:
            return True
        # 条件2：滑动窗口内成功率低于阈值
        if cred.recent_request_count >= 10:  # 至少积累10次请求才判断
            if cred.recent_success_rate < 0.8:
                return True
        return False

def compute_health_score(cred: Credential) -> float:
    """
    综合考虑：成功率、延迟、连续错误次数
    """
    # 基础分 1.0
    score = 1.0

    # 成功率扣分（权重 50%）
    success_rate = cred.recent_success_rate
    score -= (1.0 - success_rate) * 0.5

    # 延迟扣分（权重 30%）
    # 以 500ms 为基准，每超过 100ms 扣 0.05
    avg_latency = cred.recent_avg_latency
    if avg_latency > 500:
        penalty = ((avg_latency - 500) / 100) * 0.05
        score -= min(penalty * 0.3, 0.3)

    # 连续错误扣分（权重 20%）
    # 每连续错误 1 次扣 0.05
    error_penalty = min(cred.consecutive_errors * 0.05, 0.2)
    score -= error_penalty

    return max(0.0, min(1.0, score))

class AutoScaler:
    def __init__(self, pool: CredentialPool, config: dict):
        self.pool = pool
        self.max_active = config.get("max_active", 8)
        self.min_active = config.get("min_active", 1)
        self.scale_up_threshold = config.get("scale_up_threshold", 0.8)
        self.scale_down_threshold = config.get("scale_down_threshold", 0.3)
        self.scale_down_cooldown_minutes = config.get("scale_down_cooldown", 10)

        self._last_scale_down_time = None

    def maybe_scale(self):
        stats = self.pool.get_stats()
        total_max = sum(c.max_concurrency for c in self.pool._credentials
                        if c.status != CredentialStatus.RETIRED)
        total_load = stats["total_load"]
        load_ratio = total_load / total_max if total_max > 0 else 0

        active_count = stats["active_count"] + stats["busy_count"]

        if load_ratio > self.scale_up_threshold and active_count < self.max_active:
            self._scale_up()
        elif load_ratio < self.scale_down_threshold and active_count > self.min_active:
            if self._can_scale_down():
                self._scale_down()

    def _scale_up(self):
        """激活一个备用 Credential"""
        for cred in self.pool._credentials:
            if cred.status == CredentialStatus.RETIRED:
                cred.status = CredentialStatus.ACTIVE
                logger.info(f"Scaled up credential: {cred.api_key[:8]}...")
                break

    def _scale_down(self):
        """将一个低负载 Credential 置为休眠态"""
        for cred in self.pool._credentials:
            if cred.status == CredentialStatus.ACTIVE and cred.current_load == 0:
                cred.status = CredentialStatus.RETIRED
                self._last_scale_down_time = datetime.now()
                logger.info(f"Scaled down credential: {cred.api_key[:8]}...")
                break

def warmup_credential(cred: Credential, target_qps: float = 5.0, duration_seconds: int = 60):
    """
    新 Credential 激活后，先以低 QPS 预热，逐渐提升到目标 QPS
    """
    import time
    interval = 1.0 / target_qps
    elapsed = 0
    while elapsed < duration_seconds:
        # 发送少量预热请求
        try:
            httpx.get(
                cred.base_url + "/chat/completions",
                headers={"Authorization": f"Bearer {cred.api_key}"},
                json={"model": cred.model, "messages": [{"role": "user", "content": "hi"}]},
                timeout=5.0,
            )
        except Exception:
            pass
        time.sleep(interval)
        elapsed += interval
        # 线性提升 QPS
        interval = max(0.2, interval * 0.95)

# 示例：租户组配置
tenant_groups:
  enterprise:
    pool_name: "minimax_enterprise"  # 专属池
    priority: 100
  standard:
    pool_name: "minimax_shared"       # 共享池
    priority: 50
  trial:
    pool_name: "minimax_trial"        # 试用池
    priority: 10

class TwoLevelRateLimiter:
    def __init__(self, pool: CredentialPool, tenant_limiter: TenantLimiter):
        self.pool = pool
        self.tenant_limiter = tenant_limiter

    def acquire_slot(self, tenant_id: str) -> Optional[tuple[Credential, str]]:
        """
        尝试获取一个请求槽位。
        返回 (credential, request_id) 或 None（被限流）。
        """
        # 第一层：检查租户配额
        if not self.tenant_limiter.try_acquire(tenant_id):
            return None  # 租户级限流触发

        # 第二层：从池中获取可用 Credential
        cred = self.pool.acquire(tenant_id=tenant_id)
        if not cred:
            # Credential 池耗尽，释放租户配额
            self.tenant_limiter.release(tenant_id)
            return None  # 降级到下一种模型或返回限流

        return cred

    def release_slot(self, tenant_id: str, cred: Credential):
        """释放 slot"""
        self.pool.release(cred)
        self.tenant_limiter.release(tenant_id)

# 限流系统订阅 Credential 状态变更事件
class CredentialEventSubscriber:
    def __init__(self, limiter):
        self.limiter = limiter
        self.pool = None  # 注入 pool 引用

    def on_credential_status_change(self, cred: Credential, old_status, new_status):
        if new_status == CredentialStatus.CIRCUIT_BROKEN:
            self.limiter.set_key_weight(cred.api_key, 0)  # 完全移除
        elif new_status == CredentialStatus.RECOVERING:
            self.limiter.set_key_weight(cred.api_key, 0.25)  # 25% 权重
        elif new_status == CredentialStatus.ACTIVE:
            self.limiter.set_key_weight(cred.api_key, cred.weight)  # 恢复

def route(model_request: ModelRequest) -> ModelResponse:
    """
    多模型路由决策算法

    决策流程：
    1. 解析请求，提取任务类型、租户信息、成本约束
    2. 按优先级尝试每个模型
    3. 每个模型检查：可用性、负载、配额、熔断状态
    4. 成功则返回响应；失败则尝试下一个模型或降级
    """
    # Step 1: 解析请求上下文
    tenant_id = model_request.tenant_id
    task_type = model_request.task_type  # chat | completion | embedding | summarize
    cost_budget = model_request.cost_budget  # 可选的最高成本限制

    # Step 2: 确定候选模型列表（按优先级排序）
    candidates = get_candidate_models(task_type, tenant_id)

    # Step 3: 依次尝试每个候选模型
    last_error = None
    for model_group, model_name, credential_pool in candidates:
        # 3.1: 检查模型是否启用
        if not is_model_enabled(model_group, model_name):
            continue

        # 3.2: 检查租户配额是否允许使用该模型
        if not check_tenant_quota(tenant_id, model_group, model_name):
            continue

        # 3.3: 从模型对应的 Credential Pool 获取可用 Credential
        credential = credential_pool.acquire(tenant_id=tenant_id)
        if credential is None:
            continue  # 该模型池无可用资源，尝试下一个

        # 3.4: 尝试调用模型
        try:
            response = call_model_with_timeout(credential, model_request)
            credential_pool.release(credential)
            return response  # 成功，直接返回
        except ModelTimeoutError:
            credential_pool.release(credential)
            last_error = "timeout"
            continue  # 超时，尝试下一个模型
        except ModelRateLimitError:
            credential_pool.release(credential)
            # 429 错误，记录并尝试下一个模型
            record_rate_limit(model_group, model_name)
            continue
        except ModelServerError:
            credential_pool.release(credential)
            last_error = "server_error"
            # 5xx 错误，触发该模型的熔断检查
            maybe_circuit_break(model_group, model_name, credential)
            continue
        except Exception as e:
            credential_pool.release(credential)
            last_error = str(e)
            continue

    # Step 4: 所有模型都失败，进入降级或排队
    return enqueue_or_degrade(model_request, last_error)


def get_candidate_models(task_type: str, tenant_id: str) -> list:
    """
    根据任务类型和租户信息获取候选模型列表
    返回格式：[(model_group, model_name, credential_pool), ...]
    """
    # 基于任务类型确定路由目标
    if task_type in ("embedding", "similarity"):
        # 向量嵌入任务，使用辅助模型池
        return [("auxiliary", "text-embedding-3-small", embedding_pool)]
    elif task_type in ("summarize", "translate", "keyword"):
        # 摘要翻译任务，使用辅助模型池
        return [("auxiliary", "gpt-3.5-turbo", auxiliary_pool),
                ("fallback", "deepseek-chat", fallback_pool_1)]
    else:
        # 对话生成任务，使用主模型 + Fallback 链
        return [("primary", "minimax-m2.7-highspeed", primary_pool),
                ("fallback", "deepseek-chat", fallback_pool_1),
                ("fallback", "gpt-4o-mini", fallback_pool_2)]

from abc import ABC, abstractmethod
from typing import Any, AsyncIterator
from dataclasses import dataclass

@dataclass
class ModelResponse:
    """统一的模型响应格式"""
    model: str                    # 模型名称
    content: str                  # 生成的文本内容
    prompt_tokens: int            # 输入 token 数
    completion_tokens: int       # 输出 token 数
    total_tokens: int             # 总 token 数
    finish_reason: str            # 结束原因（stop / length / error）
    latency_ms: float             # 请求耗时
    raw_response: Any             # 原始响应（用于调试）


class ModelAdapter(ABC):
    """模型适配器抽象基类"""

    @property
    @abstractmethod
    def provider(self) -> str:
        """返回模型提供商名称"""
        pass

    @property
    @abstractmethod
    def model_name(self) -> str:
        """返回模型名称"""
        pass

    @abstractmethod
    async def chat_complete(
        self,
        messages: list[dict],
        temperature: float = 0.7,
        max_tokens: int = 2048,
        **kwargs
    ) -> ModelResponse:
        """
        发送对话补全请求

        参数:
            messages: 对话消息列表，格式为 [{"role": "user", "content": "..."}]
            temperature: 温度参数，控制随机性
            max_tokens: 最大输出 token 数
            **kwargs: 其他提供商特定参数

        返回:
            ModelResponse: 统一的响应格式
        """
        pass

    @abstractmethod
    async def embedding(
        self,
        text: str,
        **kwargs
    ) -> list[float]:
        """
        获取文本的向量嵌入

        参数:
            text: 输入文本
            **kwargs: 其他参数

        返回:
            list[float]: 嵌入向量
        """
        pass

class TenantQuota:
    """租户资源配额模型"""
    tenant_id: str                           # 租户唯一标识
    plan_type: str                           # 套餐类型：basic/premium/enterprise

    # === 并发控制 ===
    max_concurrent_requests: int            # 最大并发请求数，超过此数将排队等待

    # === 请求量控制 ===
    max_daily_requests: int                  # 日请求上限，-1 表示无限制
    max_weekly_requests: int                 # 周请求上限，-1 表示无限制
    max_monthly_requests: int                # 月请求上限，-1 表示无限制

    # === Token 消耗控制 ===
    max_tokens_per_day: int                 # 日 Token 上限，-1 表示无限制
    max_tokens_per_month: int                # 月 Token 上限，-1 表示无限制

    # === Credential 池访问控制 ===
    credential_pool_access: list[str]        # 该租户可用的 Credential 池列表
    exclusive_pools: list[str]              # 该租户的独占 Credential 池

    # === 配额使用统计 ===
    used_daily_requests: int                 # 当日已使用请求数
    used_monthly_requests: int              # 当月已使用请求数
    used_daily_tokens: int                  # 当日已使用 Token 数
    current_concurrent: int                  # 当前并发请求数

    # === 告警阈值配置 ===
    daily_alert_threshold: float             # 日配额告警阈值（如 0.8 表示 80%）
    monthly_alert_threshold: float          # 月配额告警阈值

    # === 配额元数据 ===
    quota_updated_at: datetime              # 配额最后更新时间
    quota_updated_by: str                   # 配额最后更新操作者

def check_hard_quota(tenant_id: str, request: LLMRequest) -> QuotaResult:
    """
    检查硬配额，返回配额检查结果。
    超配额时立即拒绝，不进入队列。
    """
    quota = get_tenant_quota(tenant_id)

    # 1. 检查并发配额
    if quota.current_concurrent >= quota.max_concurrent_requests:
        return QuotaResult(
            allowed=False,
            reason="CONCURRENT_LIMIT_EXCEEDED",
            retry_after=calculate_queue_time(quota),
            message=f"并发数超限，当前{quota.current_concurrent}，上限{quota.max_concurrent_requests}"
        )

    # 2. 检查日请求配额
    if quota.max_daily_requests > 0 and quota.used_daily_requests >= quota.max_daily_requests:
        return QuotaResult(
            allowed=False,
            reason="DAILY_REQUEST_LIMIT_EXCEEDED",
            retry_after=seconds_until_midnight(),
            message=f"日请求量超限，请明日重试"
        )

    # 3. 检查月请求配额
    if quota.max_monthly_requests > 0 and quota.used_monthly_requests >= quota.max_monthly_requests:
        return QuotaResult(
            allowed=False,
            reason="MONTHLY_REQUEST_LIMIT_EXCEEDED",
            retry_after=seconds_until_month_start(),
            message=f"月请求量超限，请次月重试"
        )

    # 4. 检查日 Token 配额
    if quota.max_tokens_per_day > 0 and quota.used_daily_tokens >= quota.max_tokens_per_day:
        return QuotaResult(
            allowed=False,
            reason="DAILY_TOKEN_LIMIT_EXCEEDED",
            retry_after=seconds_until_midnight(),
            message=f"日 Token 消耗超限，请明日重试"
        )

    # 5. 检查单次请求 Token 上限
    estimated_tokens = estimate_tokens(request)
    if quota.max_tokens_per_request > 0 and estimated_tokens > quota.max_tokens_per_request:
        return QuotaResult(
            allowed=False,
            reason="SINGLE_REQUEST_TOKEN_LIMIT_EXCEEDED",
            retry_after=0,
            message=f"单次请求 Token 数超限，最大支持{quota.max_tokens_per_request}"
        )

    return QuotaResult(allowed=True)

async def check_quota_alert(tenant_id: str, quota: TenantQuota):
    """检查配额告警，异步发送通知"""

    # 检查日请求配额告警
    if quota.max_daily_requests > 0:
        daily_ratio = quota.used_daily_requests / quota.max_daily_requests
        if daily_ratio >= 0.95:
            await send_critical_alert(tenant_id, "daily_requests", daily_ratio)
        elif daily_ratio >= 0.80:
            await send_warning_alert(tenant_id, "daily_requests", daily_ratio)

    # 检查月请求配额告警
    if quota.max_monthly_requests > 0:
        monthly_ratio = quota.used_monthly_requests / quota.max_monthly_requests
        if monthly_ratio >= 0.95:
            await send_critical_alert(tenant_id, "monthly_requests", monthly_ratio)
        elif monthly_ratio >= 0.80:
            await send_warning_alert(tenant_id, "monthly_requests", monthly_ratio)

    # 检查并发配额告警
    if quota.max_concurrent_requests > 0:
        concurrent_ratio = quota.current_concurrent / quota.max_concurrent_requests
        if concurrent_ratio >= 0.95:
            await send_critical_alert(tenant_id, "concurrent", concurrent_ratio)
        elif concurrent_ratio >= 0.80:
            await send_warning_alert(tenant_id, "concurrent", concurrent_ratio)

# 配额缓存读写逻辑
async def get_tenant_quota_cached(tenant_id: str) -> TenantQuota:
    cache_key = f"tenant:{tenant_id}:quota:cache"

    # 尝试从缓存获取
    cached = await redis.get(cache_key)
    if cached:
        return deserialize_quota(cached)

    # 缓存未命中，从数据库加载
    quota = await db.tenant_quotas.find_one(tenant_id=tenant_id)
    if quota:
        # 写入缓存，TTL 60 秒
        await redis.setex(cache_key, 60, serialize_quota(quota))

    return quota

# 启动时批量加载配额
async def preload_all_quotas():
    all_quotas = await db.tenant_quotas.find_many()
    quota_cache = {q.tenant_id: q for q in all_quotas}
    logger.info(f"Loaded {len(quota_cache)} tenant quotas into memory")
    return quota_cache

# 内存缓存 + 异步更新
class InMemoryQuotaCache:
    def __init__(self):
        self._cache: dict[str, TenantQuota] = {}
        self._lock = asyncio.Lock()
        self._last_sync = datetime.min

    async def get(self, tenant_id: str) -> Optional[TenantQuota]:
        async with self._lock:
            return self._cache.get(tenant_id)

    async def reload_if_needed(self):
        """检查是否需要全量刷新"""
        if datetime.now() - self._last_sync > timedelta(minutes=5):
            await self._full_sync()

    async def _full_sync(self):
        """全量同步数据库最新配额"""
        async with self._lock:
            all_quotas = await db.tenant_quotas.find_many()
            self._cache = {q.tenant_id: q for q in all_quotas}
            self._last_sync = datetime.now()

# 配额变更发布到消息队列
async def update_tenant_quota(tenant_id: str, updates: dict):
    # 1. 更新数据库
    await db.tenant_quotas.update_one(tenant_id, updates)

    # 2. 发布配额变更事件
    await mq.publish("quota_updates", {
        "tenant_id": tenant_id,
        "updates": updates,
        "timestamp": datetime.now().isoformat()
    })

# 消费者异步更新缓存
async def consume_quota_updates():
    async for msg in mq.subscribe("quota_updates"):
        tenant_id = msg["tenant_id"]
        updates = msg["updates"]

        # 更新 Redis 缓存
        cache_key = f"tenant:{tenant_id}:quota:cache"
        existing = await redis.get(cache_key)
        if existing:
            quota = deserialize_quota(existing)
            quota.update(updates)
            await redis.setex(cache_key, 60, serialize_quota(quota))

        # 更新内存缓存（如果存在）
        if tenant_id in in_memory_cache._cache:
            in_memory_cache._cache[tenant_id].update(updates)

class TrieNode:
    def __init__(self):
        self.children: dict[str, TrieNode] = {}
        self.tenant_ids: list[str] = []  # 存储该前缀下的所有租户 ID

class TenantIdTrie:
    """租户 ID 前缀 Trie 树，用于高效的前缀匹配查询"""

    def __init__(self):
        self.root = TrieNode()

    def insert(self, tenant_id: str):
        node = self.root
        for char in tenant_id:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
            # 在每个节点存储经过它的租户 ID（用于前缀匹配）
            if tenant_id not in node.tenant_ids:
                node.tenant_ids.append(tenant_id)

    def search_prefix(self, prefix: str) -> list[str]:
        """查找所有以 prefix 开头的租户 ID"""
        node = self.root
        for char in prefix:
            if char not in node.children:
                return []
            node = node.children[char]
        return node.tenant_ids

    def batch_load(self, tenant_ids: list[str]):
        """批量加载租户 ID 到 Trie 树"""
        for tid in tenant_ids:
            self.insert(tid)

# 使用示例
trie = TenantIdTrie()
trie.batch_load([t.tenant_id for t in all_tenants])

# 查找所有 ent_ 前缀的租户
enterprise_tenants = trie.search_prefix("ent_")

class TimeSliceScheduler:
    def __init__(self, slice_duration_ms: int = 100):
        self.slice_duration = slice_duration_ms / 1000.0  # 转换为秒
        self.tenant_queue = deque()  # 租户轮转队列
        self.tenant_last_slice = {}  # 租户上次获得时间片的时间

    def enqueue(self, tenant_id: str):
        if tenant_id not in self.tenant_last_slice:
            self.tenant_queue.append(tenant_id)
            self.tenant_last_slice[tenant_id] = 0

    def select_next(self) -> Optional[str]:
        if not self.tenant_queue:
            return None

        tenant_id = self.tenant_queue.popleft()
        now = time.time()

        # 检查该租户是否在冷却期
        elapsed = now - self.tenant_last_slice[tenant_id]
        if elapsed < self.slice_duration:
            # まだ時間切れではない場合、キューの末尾に回す
            self.tenant_queue.append(tenant_id)
            return None

        # 選択された租户を更新
        self.tenant_last_slice[tenant_id] = now
        self.tenant_queue.append(tenant_id)  # 下一轮继续轮转
        return tenant_id

class CreditBucket:
    def __init__(
        self,
        max_credits: int,
        refill_rate: float,  # 每秒补充的信用量
        refill_interval: float = 1.0  # 补充间隔（秒）
    ):
        self.max_credits = max_credits
        self.credits = max_credits
        self.refill_rate = refill_rate
        self.refill_interval = refill_interval
        self.last_refill = time.time()

    def try_consume(self, cost: int = 1) -> bool:
        """尝试消耗信用，返回是否成功"""
        self._refill()

        if self.credits >= cost:
            self.credits -= cost
            return True
        return False

    def _refill(self):
        """补充信用"""
        now = time.time()
        elapsed = now - self.last_refill
        if elapsed >= self.refill_interval:
            # 按时间比例补充信用
            refill_amount = elapsed * self.refill_rate
            self.credits = min(self.max_credits, self.credits + refill_amount)
            self.last_refill = now

    @property
    def available_credits(self) -> float:
        self._refill()
        return self.credits

class TenantCreditManager:
    """租户信用管理器"""

    def __init__(self, default_max_credits: int = 100, default_refill_rate: float = 10.0):
        self.default_max = default_max_credits
        self.default_refill = default_refill_rate
        self.buckets: dict[str, CreditBucket] = {}
        self._lock = threading.Lock()

    def get_bucket(self, tenant_id: str) -> CreditBucket:
        with self._lock:
            if tenant_id not in self.buckets:
                self.buckets[tenant_id] = CreditBucket(self.default_max, self.default_refill)
            return self.buckets[tenant_id]

    def try_request(self, tenant_id: str, cost: int = 1) -> bool:
        bucket = self.get_bucket(tenant_id)
        return bucket.try_consume(cost)

class PriorityScheduler:
    """多优先级调度器"""

    PRIORITY_LEVELS = {
        "enterprise": 0,   # 最高优先级
        "premium": 1,
        "basic": 2,
        "trial": 3,         # 试用用户，最低优先级
    }

    def __init__(self):
        self.queues: dict[int, deque] = {
            level: deque() for level in set(self.PRIORITY_LEVELS.values())
        }
        self.current_priority = 0  # 当前轮询到的优先级

    def enqueue(self, tenant_id: str, request: Request):
        priority = self._get_tenant_priority(tenant_id)
        self.queues[priority].append((tenant_id, request))

    def select_next(self) -> Optional[tuple]:
        # 从高优先级到低优先级遍历
        for level in sorted(self.queues.keys()):
            if self.queues[level]:
                return self.queues[level].popleft()
        return None

    def _get_tenant_priority(self, tenant_id: str) -> int:
        # 从租户配置获取优先级
        plan = get_tenant_plan(tenant_id)
        return self.PRIORITY_LEVELS.get(plan, 2)

class FairSchedulerWithAging:
    """
    带老化机制的公平调度器。
    等待时间越长的租户，其有效优先级越高。
    """

    def __init__(self, aging_factor: float = 0.1):
        self.aging_factor = aging_factor  # 老化系数
        self.tenant_queue: dict[str, tuple] = {}  # tenant_id -> (priority, enqueue_time, request)
        self._lock = threading.Lock()

    def enqueue(self, tenant_id: str, request: Request, base_priority: int):
        with self._lock:
            self.tenant_queue[tenant_id] = (base_priority, time.time(), request)

    def select_next(self) -> Optional[tuple]:
        with self._lock:
            if not self.tenant_queue:
                return None

            now = time.time()

            # 计算每个租户的有效优先级（考虑等待时间）
            def effective_priority(item):
                tenant_id, (base_priority, enqueue_time, _) = item
                wait_time = now - enqueue_time
                # 等待时间越长，有效优先级越高（数值越小越优先）
                age_bonus = wait_time * self.aging_factor
                return base_priority - age_bonus

            # 选择有效优先级最高的租户
            selected_id, (_, _, request) = min(
                self.tenant_queue.items(),
                key=effective_priority
            )

            del self.tenant_queue[selected_id]
            return (selected_id, request)