zsglpt/browser_pool.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""浏览器池管理 - 线程本地存储，每个线程复用自己的浏览器"""

import threading
import time
import nest_asyncio
nest_asyncio.apply()

# 线程本地存储
_thread_local = threading.local()


class BrowserPool:
    """浏览器池 - 使用线程本地存储，每个线程有自己的浏览器"""

    def __init__(self, pool_size=3, log_callback=None):
        self.pool_size = pool_size
        self.log_callback = log_callback
        self.lock = threading.Lock()
        self.all_browsers = []  # 追踪所有浏览器（用于关闭）
        self.initialized = True

    def log(self, message):
        if self.log_callback:
            self.log_callback(message)
        else:
            print(f"[浏览器池] {message}")

    def initialize(self):
        """初始化（线程本地模式下不预热）"""
        self.log(f"浏览器池已就绪（线程本地模式，每线程独立浏览器）")
        self.initialized = True

    def _create_browser(self):
        """创建一个浏览器实例"""
        try:
            from playwright.sync_api import sync_playwright

            playwright = sync_playwright().start()
            browser = playwright.chromium.launch(
                headless=True,
                args=[
                    '--no-sandbox',
                    '--disable-setuid-sandbox',
                    '--disable-dev-shm-usage',
                    '--disable-gpu',
                    '--single-process'
                ]
            )
            instance = {
                'playwright': playwright,
                'browser': browser,
                'thread_id': threading.current_thread().ident,
                'created_at': time.time(),
                'use_count': 0
            }
            with self.lock:
                self.all_browsers.append(instance)
            return instance
        except Exception as e:
            self.log(f"创建浏览器失败: {e}")
            return None

    def acquire(self, timeout=60):
        """获取当前线程的浏览器实例（如果没有则创建）"""
        # 检查当前线程是否已有浏览器
        browser_instance = getattr(_thread_local, 'browser_instance', None)

        if browser_instance:
            # 检查浏览器是否还有效
            try:
                if browser_instance['browser'].is_connected():
                    browser_instance['use_count'] += 1
                    self.log(f"复用线程浏览器（第{browser_instance['use_count']}次使用）")
                    return browser_instance
            except:
                pass
            # 浏览器已失效，清理
            self._close_browser(browser_instance)
            _thread_local.browser_instance = None

        # 为当前线程创建新浏览器
        self.log("为当前线程创建新浏览器...")
        browser_instance = self._create_browser()
        if browser_instance:
            browser_instance['use_count'] = 1
            _thread_local.browser_instance = browser_instance
        return browser_instance

    def release(self, browser_instance):
        """释放浏览器（线程本地模式下保留不关闭）"""
        if browser_instance is None:
            return

        # 检查浏览器是否还有效
        try:
            if browser_instance['browser'].is_connected():
                self.log(f"浏览器保持活跃（已使用{browser_instance['use_count']}次）")
                return
        except:
            pass

        # 浏览器已断开，清理
        self.log("浏览器已断开，清理资源")
        self._close_browser(browser_instance)
        if getattr(_thread_local, 'browser_instance', None) == browser_instance:
            _thread_local.browser_instance = None

    def _close_browser(self, browser_instance):
        """关闭单个浏览器实例"""
        try:
            if browser_instance.get('browser'):
                browser_instance['browser'].close()
            if browser_instance.get('playwright'):
                browser_instance['playwright'].stop()
            with self.lock:
                if browser_instance in self.all_browsers:
                    self.all_browsers.remove(browser_instance)
        except Exception as e:
            self.log(f"关闭浏览器失败: {e}")

    def shutdown(self):
        """关闭所有浏览器"""
        self.log("正在关闭所有浏览器...")
        for browser_instance in list(self.all_browsers):
            self._close_browser(browser_instance)
        self.all_browsers.clear()
        self.initialized = False
        self.log("浏览器池已关闭")

    def get_status(self):
        """获取池状态"""
        return {
            'pool_size': self.pool_size,
            'total_browsers': len(self.all_browsers),
            'initialized': self.initialized,
            'mode': 'thread_local'
        }


# 全局浏览器池实例
_browser_pool = None
_pool_lock = threading.Lock()


def get_browser_pool(pool_size=3, log_callback=None):
    """获取全局浏览器池实例"""
    global _browser_pool
    with _pool_lock:
        if _browser_pool is None:
            _browser_pool = BrowserPool(pool_size=pool_size, log_callback=log_callback)
        return _browser_pool


def init_browser_pool(pool_size=3, log_callback=None):
    """初始化浏览器池"""
    pool = get_browser_pool(pool_size, log_callback)
    pool.initialize()
    return pool