import { WebSocketServer as WSServer } from 'ws'; import logger from '../logger/index.js'; import bridgeManager from '../bridge-manager/index.js'; import { v4 as uuidv4 } from 'uuid'; class WebSocketServer { constructor() { this.wss = null; this.pendingRequests = new Map(); } start(server) { this.wss = new WSServer({ server, keepalive: true }); logger.info('WebSocket服务器已启动'); this.wss.on('connection', (ws) => { this.handleConnection(ws); }); } handleConnection(ws) { let bridgeId = null; let pingInterval = null; let pongTimeout = null; logger.info('新的WebSocket连接已建立'); const PING_INTERVAL = 30000; const PONG_TIMEOUT = 10000; const sendPing = () => { if (ws.readyState !== WSServer.OPEN) { clearInterval(pingInterval); return; } ws.ping(); pongTimeout = setTimeout(() => { logger.warn('PONG响应超时,关闭连接'); ws.terminate(); }, PONG_TIMEOUT); }; pingInterval = setInterval(sendPing, PING_INTERVAL); ws.on('message', (data) => { this.handleMessage(ws, data, (id) => { bridgeId = id; }); }); ws.on('pong', () => { if (pongTimeout) { clearTimeout(pongTimeout); pongTimeout = null; } }); ws.on('close', (code, reason) => { clearInterval(pingInterval); if (pongTimeout) { clearTimeout(pongTimeout); } if (bridgeId) { bridgeManager.unregisterBridge(bridgeId); this.cleanupPendingRequests(bridgeId); } logger.info(`WebSocket连接已关闭 (code: ${code})`); }); ws.on('error', (error) => { clearInterval(pingInterval); if (pongTimeout) { clearTimeout(pongTimeout); } logger.error('WebSocket连接错误:', error); }); } handleMessage(ws, data, setBridgeId) { try { const message = JSON.parse(data.toString()); logger.debug(`收到消息: ${message.type}`); switch (message.type) { case 'REGISTER': this.handleRegister(ws, message, setBridgeId); break; case 'HEARTBEAT': this.handleHeartbeat(message); break; case 'TASK_ACK': case 'TASK_END': case 'INSTANCE_CHECK_ACK': this.handleBridgeResponse(message); break; case 'PONG': break; default: logger.debug('未知消息类型:', message.type); } } catch (error) { logger.error('解析消息失败:', error); } } handleRegister(ws, message, setBridgeId) { const bridgeId = message.data?.bridgeId || uuidv4(); setBridgeId(bridgeId); bridgeManager.registerBridge(bridgeId, ws, message.data); const response = { type: 'REGISTER_ACK', data: { bridgeId, timestamp: new Date().toISOString() } }; ws.send(JSON.stringify(response)); } handleHeartbeat(message) { const bridgeId = message.data?.bridgeId; if (bridgeId) { bridgeManager.updateHeartbeat(bridgeId); } } handleBridgeResponse(message) { const requestId = message.data?.requestId; if (requestId && this.pendingRequests.has(requestId)) { const pending = this.pendingRequests.get(requestId); pending.resolve(message); this.pendingRequests.delete(requestId); } } sendTaskToBridge(bridgeId, taskData, requestId) { return new Promise((resolve, reject) => { const message = { type: 'TASK_ASSIGN', data: { ...taskData, requestId } }; const success = bridgeManager.sendToBridge(bridgeId, message); if (!success) { reject(new Error('发送任务失败')); return; } const timeout = setTimeout(() => { if (this.pendingRequests.has(requestId)) { this.pendingRequests.delete(requestId); reject(new Error('任务执行超时')); } }, 5 * 60 * 1000); this.pendingRequests.set(requestId, { resolve, reject, timeout, bridgeId }); }); } sendInstanceCheckToBridge(bridgeId, checkType, instanceId, requestId) { return new Promise((resolve, reject) => { const message = { type: 'INSTANCE_CHECK', data: { checkType, instanceId, requestId } }; const success = bridgeManager.sendToBridge(bridgeId, message); if (!success) { reject(new Error('发送实例检查请求失败')); return; } const timeout = setTimeout(() => { if (this.pendingRequests.has(requestId)) { this.pendingRequests.delete(requestId); reject(new Error('实例检查超时')); } }, 30000); this.pendingRequests.set(requestId, { resolve, reject, timeout, bridgeId }); }); } cleanupPendingRequests(bridgeId) { for (const [requestId, pending] of this.pendingRequests) { if (pending.bridgeId === bridgeId) { clearTimeout(pending.timeout); pending.reject(new Error('桥接器连接已断开')); this.pendingRequests.delete(requestId); } } } } export default new WebSocketServer();