live-hub-py/core/douyu/proxy_manager.py

"""代理池管理。"""

import threading
import time
from typing import Optional

import requests
from loguru import logger

from .proxy_resolver import ProxyResolver
from .proxy_whitelist import DouyuWhitelistSyncer


class ProxyManager:
    """代理管理器（带已验证代理池缓存，批次内共享复用）"""

    # 代理失败软隔离配置
    _COOLDOWN_BASE = 30       # 基础冷却秒数（失败1次）
    _COOLDOWN_MAX = 120       # 最大冷却秒数
    _MAX_FAIL_COUNT = 3       # 连续失败此次数后永久移出池

    def __init__(self, api_url: str = "", whitelist_uid: str = "", whitelist_ukey: str = ""):
        self.api_url = api_url
        self.whitelist_uid = whitelist_uid
        self.whitelist_ukey = whitelist_ukey
        self.current_proxy: Optional[str] = None
        self._cond = threading.Condition()
        # 已验证可用的代理池: {proxy_url: validated_timestamp}
        self._verified_pool: dict[str, float] = {}
        # 正在使用中的代理（取走但未归还），避免并发账号用同一个代理
        self._in_use: set[str] = set()
        # 代理失败记录: {proxy_url: {"count": 失败次数, "cooldown_until": 冷却到期时间戳}}
        self._fail_records: dict[str, dict] = {}
        self._pool_ttl = 90
        self._fetching = False
        self._whitelist_syncer = (
            DouyuWhitelistSyncer(whitelist_uid, whitelist_ukey)
            if whitelist_uid and whitelist_ukey
            else None
        )

    def _is_cooling_down(self, proxy_url: str) -> bool:
        """检查代理是否在冷却期内（调用前需持有锁）。"""
        record = self._fail_records.get(proxy_url)
        if not record:
            return False
        return time.time() < record.get("cooldown_until", 0)

    def _pick_from_pool_locked(self) -> Optional[str]:
        """从池中取一个未过期、未冷却、未在使用的代理（调用前需持有锁）。"""
        now = time.time()
        # 清理过期代理
        expired = [proxy for proxy, ts in self._verified_pool.items() if now - ts > self._pool_ttl]
        for proxy in expired:
            del self._verified_pool[proxy]
            self._in_use.discard(proxy)

        # 优先选未冷却、未在用的代理
        for proxy in self._verified_pool:
            if proxy not in self._in_use and not self._is_cooling_down(proxy):
                self._in_use.add(proxy)
                self.current_proxy = proxy
                return proxy

        # 退而求其次：所有未在用的（含冷却中的），也比没有强
        for proxy in self._verified_pool:
            if proxy not in self._in_use:
                self._in_use.add(proxy)
                self.current_proxy = proxy
                return proxy
        return None

    def get_proxy(self, max_attempts: int = 5) -> Optional[str]:
        """
        获取可用代理 IP，优先从已验证代理池复用。

        线程安全：池空时只有一个线程调 API 获取并验证所有代理入池，
        其他线程等待后复用。
        """
        with self._cond:
            proxy = self._pick_from_pool_locked()
            if proxy:
                logger.debug(f"从代理池复用: {proxy} (池剩余 {len(self._verified_pool)})")
                return proxy

            if self._fetching:
                logger.debug("代理池空，等待其他线程获取...")
                self._cond.wait(timeout=60)
                proxy = self._pick_from_pool_locked()
                if proxy:
                    logger.debug(f"等待后从代理池复用: {proxy} (池剩余 {len(self._verified_pool)})")
                    return proxy
                return None

            self._fetching = True

        try:
            proxies_all, msg = self._fetch_and_verify_all(max_attempts)
        finally:
            with self._cond:
                self._fetching = False
                self._cond.notify_all()

        if proxies_all:
            with self._cond:
                now = time.time()
                for proxy in proxies_all:
                    self._verified_pool[proxy] = now
                first = proxies_all[0]
                self._in_use.add(first)
                self.current_proxy = first
                logger.info(f"代理池充盈 {len(proxies_all)} 个可用代理: {first} ...")
                return first

        logger.warning(f"获取代理失败: {msg}")
        return None

    def _fetch_and_verify_all(self, max_attempts: int) -> tuple[Optional[list[str]], str]:
        """调代理 API 获取一批代理，并发验证所有可用代理。"""
        resolver = ProxyResolver(
            api_url=self.api_url,
            whitelist_syncer=self._whitelist_syncer,
            sync_local_exit_ip=bool(self._whitelist_syncer),
            sync_whitelist_once=False,
        )
        available, msg = resolver.fetch_verified(max_attempts=max_attempts, return_all=True)
        if isinstance(available, list):
            return available, msg
        if isinstance(available, str):
            return [available], msg
        return None, msg

    def mark_bad(self, proxy_url: str) -> None:
        """
        标记代理失败：软隔离（冷却期），而非永久移除。

        - 失败 < _MAX_FAIL_COUNT 次：进入冷却期，冷却后可重新入池
        - 连续失败 >= _MAX_FAIL_COUNT 次：永久移出池
        """
        with self._cond:
            self._in_use.discard(proxy_url)
            if self.current_proxy == proxy_url:
                self.current_proxy = None

            record = self._fail_records.get(proxy_url, {"count": 0, "cooldown_until": 0})
            record["count"] = record.get("count", 0) + 1
            fail_count = record["count"]

            if fail_count >= self._MAX_FAIL_COUNT:
                # 连续失败超限，永久移出
                self._verified_pool.pop(proxy_url, None)
                self._fail_records.pop(proxy_url, None)
                logger.info(
                    f"代理连续失败 {fail_count} 次，永久移出池: {proxy_url} "
                    f"(池剩余 {len(self._verified_pool)})"
                )
            else:
                # 软隔离：指数退避冷却
                cooldown = min(self._COOLDOWN_BASE * (2 ** (fail_count - 1)), self._COOLDOWN_MAX)
                record["cooldown_until"] = time.time() + cooldown
                self._fail_records[proxy_url] = record
                logger.info(
                    f"代理失败 {fail_count} 次，冷却 {cooldown}s: {proxy_url} "
                    f"(池剩余 {len(self._verified_pool)})"
                )

            self._cond.notify_all()

    def release_proxy(self, proxy_url: str) -> None:
        """归还代理到池（登录完成后调用，让其他账号可以复用）。"""
        with self._cond:
            self._in_use.discard(proxy_url)
            # 成功归还时清除失败记录
            self._fail_records.pop(proxy_url, None)
            self._cond.notify_all()

    def get_proxies_dict(self, proxy: str = None) -> dict:
        """获取 requests 使用的 proxies 字典。"""
        proxy = proxy or self.current_proxy
        if proxy:
            return {'http': proxy, 'https': proxy}
        return {}

    def verify_proxy(self, proxy: str = None) -> bool:
        """验证当前代理是否可用。"""
        proxy = proxy or self.current_proxy
        if not proxy:
            return False
        try:
            response = requests.get(
                'https://httpbin.org/ip',
                proxies={'http': proxy, 'https': proxy},
                timeout=10,
            )
            if response.status_code == 200:
                data = response.json()
                logger.info(f"代理验证成功，当前IP: {data.get('origin')}")
                return True
            return False
        except Exception as exc:
            logger.error(f"代理验证失败: {exc}")
            return False


def get_proxy_manager(
    api_url: str = "",
    whitelist_uid: str = "",
    whitelist_ukey: str = "",
) -> ProxyManager:
    """获取代理管理器实例。"""
    return ProxyManager(api_url, whitelist_uid=whitelist_uid, whitelist_ukey=whitelist_ukey)