369 lines
12 KiB
JavaScript
369 lines
12 KiB
JavaScript
import express, { json, urlencoded } from 'express';
|
||
import cors from 'cors';
|
||
import dotenv from 'dotenv';
|
||
import WebSocket, { WebSocketServer } from 'ws';
|
||
import { Worker } from 'worker_threads';
|
||
import { checkUsertoken } from './school/api.js';
|
||
import redis from './redis/index.js';
|
||
import initQueue from './redis/initQueue.js';
|
||
import messagePersistence from './redis/messagePersistence.js';
|
||
import code from './config/code.json' with { type: 'json' };
|
||
import fileRouter from './upload/index.js';
|
||
import recordRouter from './outside/callback.js';
|
||
import mdWebSocketServer from './utils/mdWebSocketServer.js';
|
||
|
||
dotenv.config();
|
||
|
||
const logger = {
|
||
info: (message) => {
|
||
const timestamp = new Date().toISOString();
|
||
console.log(`[${timestamp}] INFO: ${message}`);
|
||
},
|
||
error: (message, error) => {
|
||
const timestamp = new Date().toISOString();
|
||
console.error(`[${timestamp}] ERROR: ${message}`, error || '');
|
||
},
|
||
debug: (message) => {
|
||
const timestamp = new Date().toISOString();
|
||
console.debug(`[${timestamp}] DEBUG: ${message}`);
|
||
}
|
||
};
|
||
|
||
let wss = null;
|
||
const workers = [];
|
||
const socketMap = new Map();
|
||
|
||
function createWorker(scriptPath) {
|
||
const worker = new Worker(scriptPath);
|
||
worker.setMaxListeners(20);
|
||
worker.on('error', (error) => {
|
||
logger.error(`Worker ${scriptPath} 错误:`, error);
|
||
});
|
||
worker.on('exit', (code) => {
|
||
if (code !== 0) {
|
||
logger.error(`Worker ${scriptPath} 异常退出，退出码: ${code}`);
|
||
}
|
||
});
|
||
workers.push(worker);
|
||
return worker;
|
||
}
|
||
|
||
const assessment = createWorker('./worker_threads/assessment/assessment.js');
|
||
const wait = createWorker('./worker_threads/wait/waiting.js');
|
||
const polling = createWorker('./worker_threads/process/process.js');
|
||
const result = createWorker('./worker_threads/result/result.js');
|
||
const callback_timeout = createWorker('./worker_threads/callback_timeout/callbackTimeout.js');
|
||
const callback_result = createWorker('./worker_threads/callback_result/result.js');
|
||
const error = createWorker('./worker_threads/error/error.js');
|
||
|
||
async function sendMessageToClient(id, message, close = false, closeCode = 1000, closeReason = '') {
|
||
let socket;
|
||
if (typeof id === 'string' && id) {
|
||
socket = socketMap.get(id);
|
||
}
|
||
if (socket && socket.readyState === WebSocket.OPEN && message) {
|
||
try {
|
||
socket.send(message);
|
||
const messagePreview = typeof message === 'string' ? message.slice(0, 50) : JSON.stringify(message).slice(0, 50);
|
||
logger.debug(`成功发送消息到客户端，id: ${id}, 消息: ${messagePreview}...`);
|
||
if (close) {
|
||
socket.close(closeCode, closeReason);
|
||
}
|
||
return true;
|
||
} catch (error) {
|
||
logger.error(`发送消息给客户端失败，id: ${id}`, error);
|
||
return false;
|
||
}
|
||
} else {
|
||
if (!message) {
|
||
logger.debug(`消息为空，无法发送，id: ${id}`);
|
||
return false;
|
||
} else {
|
||
logger.debug(`未找到目标客户端或连接已关闭，保存消息到待发送队列，id: ${id}`);
|
||
try {
|
||
await messagePersistence.savePendingMessage(id, message);
|
||
logger.info(`消息已保存到待发送队列，等待重试: backendId=${id}`);
|
||
return false;
|
||
} catch (error) {
|
||
logger.error(`保存待发送消息失败: backendId=${id}`, error);
|
||
return false;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
function createWebSocketServer(server) {
|
||
wss = new WebSocketServer({
|
||
server,
|
||
verifyClient: async (info, callback) => {
|
||
try {
|
||
const urlParams = new URLSearchParams(info.req.url.split('?')[1]);
|
||
const token = urlParams.get('token');
|
||
const id = urlParams.get('id');
|
||
|
||
if (!token) {
|
||
logger.info('缺少令牌');
|
||
callback(false, 401, '缺少令牌');
|
||
return;
|
||
} else if (token !== process.env.TOKEN_SECRET){
|
||
logger.info('验证后端失败');
|
||
callback(false, 401, 'Token is invalid');
|
||
return;
|
||
}
|
||
|
||
info.req.id = id;
|
||
logger.info(`用户ID: token 验证成功`);
|
||
callback(true);
|
||
} catch (error) {
|
||
logger.error('验证后端失败:', error);
|
||
callback(false, 401, 'Token is invalid');
|
||
}
|
||
}
|
||
});
|
||
|
||
logger.info(`WebSocket server is running on the same port as HTTP server`);
|
||
|
||
wss.on('error', (error) => {
|
||
logger.error('WebSocket服务器错误:', error);
|
||
});
|
||
|
||
wss.on('connection', async (socket, req) => {
|
||
const id = req.id;
|
||
logger.info(`${id}号后端 连接成功`);
|
||
socketMap.set(id, socket);
|
||
socket.send('please give me tasks');
|
||
|
||
try {
|
||
const pendingMessages = await messagePersistence.getPendingMessages(id);
|
||
if (pendingMessages.length > 0) {
|
||
logger.info(`${id}号后端 发现 ${pendingMessages.length} 条待发送消息，开始重试发送`);
|
||
|
||
for (const pendingMsg of pendingMessages) {
|
||
try {
|
||
socket.send(pendingMsg.message);
|
||
await messagePersistence.removePendingMessage(pendingMsg.key);
|
||
logger.debug(`成功重试发送消息: backendId=${id}, messageKey=${pendingMsg.key}`);
|
||
} catch (error) {
|
||
logger.error(`重试发送消息失败: backendId=${id}, messageKey=${pendingMsg.key}`, error);
|
||
await messagePersistence.incrementRetryCount(pendingMsg.key);
|
||
}
|
||
}
|
||
|
||
logger.info(`${id}号后端 待发送消息重试完成`);
|
||
}
|
||
} catch (error) {
|
||
logger.error(`获取或发送待发送消息失败: backendId=${id}`, error);
|
||
}
|
||
|
||
socket.on('message', (message) => {
|
||
const messageStr = typeof message === 'string' ? message : message.toString();
|
||
|
||
if (messageStr === 'ping') {
|
||
socket.send('pong');
|
||
return;
|
||
}
|
||
|
||
if (messageStr === 'pong') {
|
||
return;
|
||
}
|
||
|
||
if (messageStr === 'please give me tasks') {
|
||
return;
|
||
}
|
||
|
||
try {
|
||
const msg = JSON.parse(messageStr);
|
||
|
||
if (msg.type === 'JWT_UPDATE') {
|
||
logger.info(`收到 JWT_UPDATE 消息`);
|
||
return;
|
||
}
|
||
|
||
if (msg.type === 'CAPACITY_UPDATE') {
|
||
logger.debug(`收到算力状态更新: 可用容量 = ${msg.data?.summary?.availableCapacity || 0}`);
|
||
return;
|
||
}
|
||
|
||
if (msg.type === 'INSTANCE_ONLINE') {
|
||
logger.debug(`收到实例上线: ${msg.data?.instanceId}`);
|
||
return;
|
||
}
|
||
|
||
if (msg.type === 'INSTANCE_OFFLINE') {
|
||
logger.debug(`收到实例下线: ${msg.data?.instanceId}`);
|
||
return;
|
||
}
|
||
|
||
if (msg.type === 'HEARTBEAT') {
|
||
socket.send(JSON.stringify({
|
||
type: 'HEARTBEAT_ACK',
|
||
data: { timestamp: new Date().toISOString() }
|
||
}));
|
||
return;
|
||
}
|
||
|
||
const prefix = messageStr.slice(0, 50);
|
||
if (prefix.includes('"type":"generate"') || prefix.includes("'type':'generate'")) {
|
||
assessment.postMessage({
|
||
type: 'submit',
|
||
data: messageStr
|
||
});
|
||
} else {
|
||
logger.debug(`收到未知消息类型: ${prefix}`);
|
||
}
|
||
} catch (e) {
|
||
logger.error('处理消息出错:', e);
|
||
socket.send(JSON.stringify({
|
||
error: '处理消息出错',
|
||
details: e.message
|
||
}));
|
||
}
|
||
});
|
||
|
||
const heartbeatInterval = setInterval(() => {
|
||
if (socket.readyState === WebSocket.OPEN) {
|
||
socket.send('ping');
|
||
// logger.debug(`向 ${id} 号后端发送心跳`);
|
||
}
|
||
}, 30000);
|
||
|
||
socket.on('close', (code, reason) => {
|
||
clearInterval(heartbeatInterval);
|
||
logger.info(`${id}号后端 连接关闭，关闭码: ${code}，原因: ${reason}`);
|
||
});
|
||
|
||
socket.on('error', (error) => {
|
||
logger.error(`${id}号后端 连接错误:`, error);
|
||
});
|
||
});
|
||
|
||
assessment.on('message', async (message) => {
|
||
logger.debug(`收到assessment worker消息: ${JSON.stringify(message)}`);
|
||
if (message.type === 'AssessmentSuccess') {
|
||
await sendMessageToClient(message.backendId, code.SUCCESS[message.type]);
|
||
} else {
|
||
await sendMessageToClient(message.backendId, code.ERROR[message.type], false, 4401, code.ERROR[message.type]);
|
||
}
|
||
});
|
||
|
||
result.on('message', async (message) => {
|
||
logger.debug(`收到result worker消息: ${JSON.stringify(message)}`);
|
||
if (message.type === 'success') {
|
||
await sendMessageToClient(message.backendId, message.message, false, 1000, 'success');
|
||
} else {
|
||
await sendMessageToClient(message.backendId, '获取结果失败，可在历史记录区刷新查看结果', false, 4401, code.ERROR[message.type]);
|
||
}
|
||
});
|
||
|
||
callback_result.on('message', async (message) => {
|
||
logger.debug(`收到callback_result worker消息: ${JSON.stringify(message)}`);
|
||
if (message.type === 'success') {
|
||
await sendMessageToClient(message.backendId, message.message, false, 1000, 'success');
|
||
} else {
|
||
await sendMessageToClient(message.backendId, '获取结果失败，可在历史记录区刷新查看结果', false, 4401);
|
||
}
|
||
});
|
||
|
||
error.on('message', async (message) => {
|
||
logger.debug(`收到error worker消息: ${JSON.stringify(message)}`);
|
||
await sendMessageToClient(message.backendId, message.message, false, 4402, 'false');
|
||
});
|
||
}
|
||
|
||
function gracefulShutdown() {
|
||
logger.info('开始优雅关闭...');
|
||
|
||
if (wss) {
|
||
wss.close(() => {
|
||
logger.info('WebSocket服务器已关闭');
|
||
});
|
||
|
||
wss.clients.forEach((client) => {
|
||
client.close(1001, '服务器正在关闭');
|
||
});
|
||
}
|
||
|
||
workers.forEach((worker, index) => {
|
||
logger.info(`终止worker线程 ${index}`);
|
||
worker.terminate();
|
||
});
|
||
|
||
if (redis.isOpen) {
|
||
redis.disconnect()
|
||
.then(() => {
|
||
logger.info('Redis连接已关闭');
|
||
process.exit(0);
|
||
})
|
||
.catch((error) => {
|
||
logger.error('关闭Redis连接失败:', error);
|
||
process.exit(1);
|
||
});
|
||
} else {
|
||
process.exit(0);
|
||
}
|
||
|
||
setTimeout(() => {
|
||
logger.info('强制关闭...');
|
||
process.exit(1);
|
||
}, 5000);
|
||
}
|
||
|
||
async function initialize() {
|
||
logger.info('***************初始化队列开始***************');
|
||
try {
|
||
if (!redis.isOpen) {
|
||
await redis.connect();
|
||
logger.info('Redis 连接成功');
|
||
}
|
||
|
||
await initQueue.init();
|
||
logger.info('***************初始化队列完成***************');
|
||
|
||
setInterval(async () => {
|
||
try {
|
||
await messagePersistence.cleanupOldMessages(2 * 24 * 60 * 60 * 1000);
|
||
} catch (error) {
|
||
logger.error('定期清理过期消息失败:', error);
|
||
}
|
||
}, 24 * 60 * 60 * 1000);
|
||
|
||
} catch (err) {
|
||
logger.error('初始化失败:', err);
|
||
process.exit(1);
|
||
}
|
||
}
|
||
|
||
const app = express();
|
||
const hostname = '0.0.0.0';
|
||
const port = process.env.WS_PORT || 8087;
|
||
|
||
app.use(cors());
|
||
app.use('/workflow/uploads', express.static('uploads'));
|
||
app.use(json());
|
||
app.use(urlencoded({ extended: true }));
|
||
app.use('/workflow/file', fileRouter);
|
||
app.use('/callback', recordRouter);
|
||
|
||
app.get('/', (req, res) => {
|
||
res.json({
|
||
name: 'ComfyUI Task Queue Backend',
|
||
version: '1.0.0',
|
||
status: 'running',
|
||
timestamp: new Date().toISOString()
|
||
});
|
||
});
|
||
|
||
process.on('SIGINT', gracefulShutdown);
|
||
process.on('SIGTERM', gracefulShutdown);
|
||
|
||
initialize().then(() => {
|
||
const server = app.listen(port, hostname, () => {
|
||
console.log(`========================================`);
|
||
console.log(`HTTP & WebSocket 服务器已启动`);
|
||
console.log(`服务地址: http://${hostname}:${port}/`);
|
||
console.log(`========================================`);
|
||
|
||
createWebSocketServer(server);
|
||
});
|
||
});
|