load_balancer.py

"""
负载均衡器 - Round Robin + 自动故障转移
"""
from typing import Optional
from datetime import datetime, timedelta
from account_manager import account_manager
from models import Account


class LoadBalancer:
    """智能负载均衡器"""
    
    def __init__(self):
        self.cooldown_duration = 60  # 默认冷却时间（秒）
        self.max_retries = 3
    
    async def get_next_account(self) -> Optional[Account]:
        """获取下一个可用账号"""
        return await account_manager.get_next_token()
    
    def mark_account_error(self, account_id: str, error_code: int, error_msg: str = ""):
        """
        标记账号错误，根据错误类型决定是否冷却
        - 429: Too Many Requests - 立即冷却
        - 400: Bad Request - 可能是请求格式问题，不冷却
        - 401: Unauthorized - Token 失效，冷却
        - 403: Forbidden - 配额耗尽，长时间冷却
        """
        account_manager.update_account_stats(account_id, success=False, error=error_msg)
        
        if error_code == 429:
            # 速率限制，冷却 60 秒
            account_manager.set_account_cooldown(account_id, 60)
            print(f"账号 {account_id} 触发速率限制，冷却 60 秒")
        elif error_code == 401:
            # Token 失效，冷却 30 秒（等待刷新）
            account_manager.set_account_cooldown(account_id, 30)
            print(f"账号 {account_id} Token 失效，冷却 30 秒")
        elif error_code == 403:
            # 配额耗尽，冷却 5 分钟
            account_manager.set_account_cooldown(account_id, 300)
            print(f"账号 {account_id} 配额耗尽，冷却 5 分钟")
    
    def mark_account_success(self, account_id: str):
        """标记请求成功"""
        account_manager.update_account_stats(account_id, success=True)


# 全局单例
load_balancer = LoadBalancer()