/** * @fileoverview 任务队列管理模块 * @description 负责请求队列、并发控制和心跳机制 */ import { logger } from '../utils/logger.js'; import { sendJson, sendSse, sendSseDone, sendHeartbeat, sendApiError, buildChatCompletion, buildChatCompletionChunk } from './respond.js'; /** * @typedef {object} TaskContext * @property {import('http').IncomingMessage} req - HTTP 请求对象 * @property {import('http').ServerResponse} res - HTTP 响应对象 * @property {string} prompt - 用户提示词 * @property {string[]} imagePaths - 图片路径列表 * @property {string|null} modelId - 模型 ID * @property {string|null} modelName - 模型名称 * @property {string} id - 请求唯一标识 * @property {boolean} isStreaming - 是否流式请求 */ /** * @typedef {object} QueueConfig * @property {number} maxConcurrent - 最大并发数 * @property {number} maxQueueSize - 最大队列大小 * @property {string} keepaliveMode - 心跳模式 ('comment' | 'content') */ /** * @typedef {object} BrowserContext * @property {object} browser - Playwright 浏览器实例 * @property {object} page - Playwright 页面实例 * @property {object} config - 配置对象 */ /** * 创建任务队列管理器 * @param {QueueConfig} queueConfig - 队列配置 * @param {object} callbacks - 回调函数 * @param {Function} callbacks.initBrowser - 初始化浏览器函数 * @param {Function} callbacks.generateImage - 生成图片函数 * @param {object} callbacks.config - 配置对象 * @returns {object} 队列管理器 */ export function createQueueManager(queueConfig, callbacks) { const { maxConcurrent, maxQueueSize, keepaliveMode } = queueConfig; const { initBrowser, generateImage, config } = callbacks; /** @type {TaskContext[]} */ const queue = []; /** @type {number} */ let processingCount = 0; /** @type {BrowserContext|null} */ let browserContext = null; /** * 清理任务临时文件 * @param {TaskContext} task - 任务上下文 */ async function cleanupTask(task) { if (task?.imagePaths) { const fs = await import('fs'); for (const p of task.imagePaths) { try { fs.unlinkSync(p); } catch (e) { /* ignore */ } } } } /** * 处理单个任务 * @param {TaskContext} task - 任务上下文 */ async function processTask(task) { const { res, prompt, imagePaths, modelId, modelName, id, isStreaming } = task; logger.info('服务器', '[队列] 开始处理任务', { id, remaining: queue.length }); // 启动心跳(流式请求) let heartbeatInterval = null; if (isStreaming) { heartbeatInterval = setInterval(() => { if (res.writableEnded) { clearInterval(heartbeatInterval); return; } sendHeartbeat(res, keepaliveMode, modelName); }, 3000); } try { // 确保浏览器已初始化 if (!browserContext) { browserContext = await initBrowser(config); } // 调用核心生图逻辑 const result = await generateImage(browserContext, prompt, imagePaths, modelId, { id }); // 清除心跳 if (heartbeatInterval) clearInterval(heartbeatInterval); // 处理结果 let finalContent = ''; if (result.error) { // 适配器层已归一化错误,直接构造错误响应 finalContent = `[生成错误] ${result.error}`; } else if (result.image) { finalContent = `![generated](${result.image})`; logger.info('服务器', '图片已准备就绪 (Base64)', { id }); } else { finalContent = result.text || '生成失败'; } // 发送响应 if (isStreaming) { const chunk = buildChatCompletionChunk(finalContent, modelName); sendSse(res, chunk); sendSseDone(res); } else { const response = buildChatCompletion(finalContent, modelName); sendJson(res, 200, response); } } catch (err) { // 清除心跳 if (heartbeatInterval) clearInterval(heartbeatInterval); logger.error('服务器', '任务处理失败', { id, error: err.message }); sendApiError(res, { code: ERROR_CODES.INTERNAL_ERROR, error: err.message, isStreaming }); } } /** * 处理队列中的任务 */ async function processQueue() { // 如果正在处理的任务已满,或队列为空,则停止 if (processingCount >= maxConcurrent || queue.length === 0) return; // 取出下一个任务 const task = queue.shift(); processingCount++; try { await processTask(task); } finally { // 清理临时文件 cleanupTask(task); processingCount--; // 递归处理下一个任务 processQueue(); } } /** * 添加任务到队列 * @param {TaskContext} task - 任务上下文 */ function addTask(task) { queue.push(task); processQueue(); } /** * 获取当前队列状态 * @returns {{queueLength: number, processing: number, total: number}} */ function getStatus() { return { queueLength: queue.length, processing: processingCount, total: processingCount + queue.length }; } /** * 检查是否可以接受新请求(非流式) * @returns {boolean} */ function canAcceptNonStreaming() { return processingCount + queue.length < maxQueueSize; } /** * 初始化浏览器 * @returns {Promise} */ async function initializeBrowser() { browserContext = await initBrowser(config); return browserContext; } /** * 获取浏览器上下文 * @returns {BrowserContext|null} */ function getBrowserContext() { return browserContext; } return { addTask, getStatus, canAcceptNonStreaming, initializeBrowser, getBrowserContext, maxQueueSize }; }