shuzhiren-comfyui/message-dispatcher/src/websocket-server/index.js

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();