WebAI2API/src/backend/pool/PoolManager.js

/**
 * @fileoverview PoolManager 类
 * @description 管理 Worker 池，负责初始化、任务分发和故障转移
 */

import { logger } from '../../utils/logger.js';
import { registry } from '../registry.js';
import { createStrategySelector } from '../strategies/index.js';
import { executeWithFailover } from '../strategies/failover.js';
import { normalizeError } from '../utils/error.js';
import { Worker } from './Worker.js';

/**
 * PoolManager 类 - 管理 Worker 池
 */
export class PoolManager {
    /**
     * @param {object} config - 全局配置
     */
    constructor(config) {
        this.config = config;
        this.workers = [];
        this.strategy = config.backend.pool.strategy || 'least_busy';
        this.strategySelector = createStrategySelector(this.strategy);
        this.initialized = false;
        this.roundRobinIndex = 0;
    }

    /**
     * 初始化所有 Worker
     */
    async initAll() {
        if (this.initialized) return;

        // 先加载所有适配器
        await registry.loadAll();

        // 注入适配器配置（用于模型过滤）
        const adapterConfig = this.config.backend?.adapter || {};
        registry.setAdapterConfig(adapterConfig);

        // 解析登录模式参数
        let loginWorkerName = null;
        const loginArg = process.argv.find(arg => arg.startsWith('-login'));
        const isLoginMode = !!loginArg;
        if (loginArg && loginArg.includes('=')) {
            loginWorkerName = loginArg.split('=')[1];
            logger.info('工作池', `登录模式: 仅初始化 Worker "${loginWorkerName}"`);
        } else if (isLoginMode) {
            loginWorkerName = this.config.backend.pool.workers[0]?.name || null;
            logger.info('工作池', `登录模式: 仅初始化第一个 Worker "${loginWorkerName}"`);
        }

        const workerConfigs = this.config.backend.pool.workers;

        if (isLoginMode) {
            logger.info('工作池', `登录模式: 从 ${workerConfigs.length} 个 Worker 中筛选...`);
        } else {
            logger.info('工作池', `正在初始化 ${workerConfigs.length} 个 Worker...`);
        }

        // 过滤并创建 Worker 实例
        const validWorkers = [];
        for (const workerConfig of workerConfigs) {
            if (isLoginMode && workerConfig.name !== loginWorkerName) {
                logger.debug('工作池', `[${workerConfig.name}] 跳过 (不匹配登录目标)`);
                continue;
            }

            if (workerConfig.type !== 'merge' && !registry.hasAdapter(workerConfig.type)) {
                logger.error('工作池', `Worker [${workerConfig.name}] 的类型 "${workerConfig.type}" 无对应适配器，跳过`);
                continue;
            }

            if (workerConfig.type === 'merge') {
                const invalidTypes = (workerConfig.mergeTypes || []).filter(t => !registry.hasAdapter(t));
                if (invalidTypes.length > 0) {
                    logger.error('工作池', `Worker [${workerConfig.name}] 的 mergeTypes 包含无效类型: ${invalidTypes.join(', ')}`);
                    continue;
                }
            }

            validWorkers.push(new Worker(this.config, workerConfig));
        }

        if (isLoginMode && validWorkers.length === 0) {
            const availableNames = workerConfigs.map(w => w.name).join(', ');
            throw new Error(`登录模式未找到 Worker "${loginWorkerName}"。可用的 Worker: ${availableNames}`);
        }

        // 按 userDataDir 分组
        const browserMap = new Map();

        for (const worker of validWorkers) {
            try {
                const existing = browserMap.get(worker.userDataDir);

                if (existing) {
                    const workerProxy = JSON.stringify(worker.proxyConfig || null);
                    const existingProxy = JSON.stringify(existing.proxyConfig || null);
                    if (workerProxy !== existingProxy) {
                        logger.warn('工作池', `[${worker.name}] 代理配置与 [${existing.firstWorkerName}] 不一致，将使用后者的配置`);
                    }

                    logger.debug('工作池', `[${worker.name}] 将与其他 Worker 共享浏览器 (${worker.userDataDir})`);
                    await worker.init(existing.browser);

                    // 建立共享关系：设置所有者引用，并添加到所有者的共享列表
                    worker._browserOwner = existing.ownerWorker;
                    existing.ownerWorker._sharedWorkers.push(worker);
                } else {
                    await worker.init();
                    browserMap.set(worker.userDataDir, {
                        browser: worker.browser,
                        proxyConfig: worker.proxyConfig,
                        firstWorkerName: worker.name,
                        ownerWorker: worker  // 保存所有者 Worker 引用
                    });
                }

                this.workers.push(worker);
            } catch (e) {
                logger.error('工作池', `[${worker.name}] 初始化失败，跳过该 Worker`, { error: e.message });
            }
        }

        if (this.workers.length === 0) {
            throw new Error('所有 Worker 初始化都失败了，无法启动服务');
        }

        this.initialized = true;
        logger.info('工作池', `工作池初始化完成，共 ${this.workers.length} 个 Worker 就绪 (${browserMap.size} 个浏览器实例)`);
    }

    /**
     * 根据模型选择 Worker
     */
    selectWorker(modelId) {
        const candidates = this.workers.filter(w => w.supports(modelId));

        if (candidates.length === 0) {
            throw new Error(`没有 Worker 支持模型: ${modelId}`);
        }

        if (candidates.length === 1) {
            return candidates[0];
        }

        switch (this.strategy) {
            case 'round_robin': {
                const idx = this.roundRobinIndex % candidates.length;
                this.roundRobinIndex++;
                return candidates[idx];
            }
            case 'random': {
                const idx = Math.floor(Math.random() * candidates.length);
                return candidates[idx];
            }
            case 'least_busy':
            default: {
                return candidates.reduce((min, w) => w.busyCount < min.busyCount ? w : min, candidates[0]);
            }
        }
    }

    /**
     * 分发生图任务（支持故障转移）
     */
    async generate(ctx, prompt, paths, modelId, meta) {
        const failoverConfig = this.config.backend?.pool?.failover || {};
        const failoverEnabled = failoverConfig.enabled !== false;
        const maxRetries = failoverConfig.maxRetries || 2;

        let candidates = this.workers.filter(w => w.supports(modelId));

        if (candidates.length === 0) {
            return { error: `没有 Worker 支持模型: ${modelId}` };
        }

        // 如果请求包含图片，优先选择 imagePolicy 为 optional 的 Worker
        const hasImages = paths && paths.length > 0;
        if (hasImages && candidates.length > 1) {
            const optionalCandidates = candidates.filter(w => {
                const policy = w.getImagePolicy(modelId);
                return policy === 'optional' || policy === 'required';
            });

            if (optionalCandidates.length > 0) {
                logger.debug('工作池', `请求包含图片，优先选择支持图片的 Worker (${optionalCandidates.length}/${candidates.length} 个)`);
                candidates = optionalCandidates;
            } else {
                logger.warn('工作池', `请求包含图片，但没有 Worker 的 imagePolicy 为 optional`);
            }
        }

        const sortedCandidates = this.strategySelector.sort(candidates);

        if (!failoverEnabled) {
            const worker = sortedCandidates[0];
            logger.debug('工作池', `任务分发至: ${worker.name} (busy: ${worker.busyCount})`);
            return await this._safeExecuteWorker(worker, ctx, prompt, paths, modelId, meta);
        }

        return await executeWithFailover(
            sortedCandidates,
            async (worker) => {
                logger.debug('工作池', `任务分发至: ${worker.name} (busy: ${worker.busyCount})`);
                return await this._safeExecuteWorker(worker, ctx, prompt, paths, modelId, meta);
            },
            {
                maxRetries,
                meta,
                onRetry: (worker, error) => {
                    logger.warn('工作池', `[${worker.name}] 失败，尝试下一个 Worker...`, { error, ...meta });
                }
            }
        );
    }

    /**
     * 安全执行 Worker（带错误边界）
     * @private
     */
    async _safeExecuteWorker(worker, ctx, prompt, paths, modelId, meta) {
        try {
            return await worker.generate(ctx, prompt, paths, modelId, meta);
        } catch (err) {
            logger.error('工作池', `[${worker.name}] 执行异常`, { error: err.message, ...meta });
            return normalizeError(err.message || '执行异常');
        }
    }

    /**
     * 获取所有模型列表
     */
    getModels() {
        const allModels = [];
        const seenIds = new Set();

        for (const worker of this.workers) {
            const models = worker.getModels();
            for (const m of models) {
                if (!seenIds.has(m.id)) {
                    seenIds.add(m.id);
                    allModels.push(m);
                }
            }
        }

        return { object: 'list', data: allModels };
    }

    /**
     * 获取图片策略（宽松策略：只要有一个 Worker 支持 optional 就返回 optional）
     */
    getImagePolicy(modelKey) {
        const policies = new Set();

        for (const worker of this.workers) {
            if (worker.supports(modelKey)) {
                policies.add(worker.getImagePolicy(modelKey));
            }
        }

        // 宽松策略：只要有一个 optional 就返回 optional
        if (policies.has('optional')) return 'optional';
        if (policies.has('required')) return 'required';
        if (policies.has('forbidden')) return 'forbidden';
        return 'optional';
    }

    /**
     * 获取模型类型
     */
    getModelType(modelKey) {
        for (const worker of this.workers) {
            if (worker.supports(modelKey)) {
                return worker.getModelType(modelKey);
            }
        }
        return 'image';
    }

    /**
     * 获取指定实例的 Cookies
     */
    async getCookies(instanceName, domain) {
        let worker;
        if (instanceName) {
            worker = this.workers.find(w => w.instanceName === instanceName);
            if (!worker) {
                throw new Error(`浏览器实例不存在: ${instanceName}`);
            }
        } else {
            worker = this.workers[0];
            if (!worker) {
                throw new Error('工作池中没有可用的 Worker');
            }
        }

        const cookies = await worker.getCookies(domain);
        return { instance: worker.instanceName, cookies };
    }

    /**
     * 触发所有 merge Worker 的监控导航
     */
    async navigateToMonitor() {
        for (const worker of this.workers) {
            if (worker.type === 'merge' && worker.busyCount === 0) {
                await worker.navigateToMonitor();
            }
        }
    }

    /**
     * 获取第一个 Worker 的 page
     */
    getFirstPage() {
        return this.workers[0]?.page || null;
    }
}