shuzhiren-comfyui/message-dispatcher/src/md-websocket-client/index.js

import WebSocket from 'ws';
import dotenv from 'dotenv';
import bridgeManager from '../bridge-manager/index.js';
import jwt from 'jsonwebtoken';

dotenv.config();

class MDWebSocketClient {
  constructor() {
    this.ws = null;
    this.connected = false;
    this.reconnectAttempts = 0;
    this.tokenPushInterval = null;
    this.heartbeatInterval = null;
    this.serverUrl = process.env.TASK_QUEUE_WS_URL || 'ws://localhost:8087';
    this.jwtSecret = process.env.JWT_SECRET || 'comfyui-cluster-bridge-secret-key-2024';
    this.taskQueueToken = process.env.TASK_QUEUE_TOKEN || '1Ag9BJJn0rXDnidCyXqu';
  }

  async init() {
    console.log('[MDWebSocketClient] 初始化 WebSocket 客户端');
    bridgeManager.onBridgeChange(() => {
      if (this.connected) {
        this.pushCapacityState();
      }
    });
    await this.connect();
  }

  async connect() {
    return new Promise((resolve, reject) => {
      const urlWithParams = `${this.serverUrl}?token=${this.taskQueueToken}&id=message-dispatcher`;
      console.log(`[MDWebSocketClient] 正在连接到 ${urlWithParams}`);
      
      this.ws = new WebSocket(urlWithParams);

      this.ws.on('open', () => {
        console.log('[MDWebSocketClient] WebSocket 连接已建立');
        this.connected = true;
        this.reconnectAttempts = 0;

        this.pushJwtToken();
        this.pushCapacityState();

        this.startHeartbeat();
        this.startTokenPushTimer();

        resolve();
      });

      this.ws.on('message', (data) => {
        this.handleMessage(data);
      });

      this.ws.on('close', (code, reason) => {
        console.log(`[MDWebSocketClient] WebSocket 连接已关闭 (code: ${code})`);
        this.connected = false;
        this.clearIntervals();
        this.scheduleReconnect();
      });

      this.ws.on('error', (error) => {
        // console.error('[MDWebSocketClient] WebSocket 连接错误:', error);
        this.connected = false;
      });
    });
  }

  disconnect() {
    console.log('[MDWebSocketClient] 断开 WebSocket 连接');
    this.clearIntervals();
    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
    this.connected = false;
  }

  clearIntervals() {
    if (this.tokenPushInterval) {
      clearInterval(this.tokenPushInterval);
      this.tokenPushInterval = null;
    }
    if (this.heartbeatInterval) {
      clearInterval(this.heartbeatInterval);
      this.heartbeatInterval = null;
    }
  }

  scheduleReconnect() {
    const delay = Math.min(1000 * Math.pow(2, this.reconnectAttempts), 30000);
    this.reconnectAttempts++;
    console.log(`[MDWebSocketClient] ${delay}ms 后尝试重连 (第 ${this.reconnectAttempts} 次)`);
    
    setTimeout(() => {
      this.connect().catch(err => {
        console.error('[MDWebSocketClient] 重连失败:', err);
      });
    }, delay);
  }

  pushJwtToken() {
    const token = this.generateJwtToken();
    const expiresAt = new Date(Date.now() + 24 * 60 * 60 * 1000).toISOString();
    
    const message = {
      type: 'JWT_UPDATE',
      data: {
        token,
        expiresAt,
        timestamp: new Date().toISOString()
      }
    };
    
    this.send(message);
    console.log('[MDWebSocketClient] 已推送 JWT Token');
  }

  generateJwtToken() {
    return jwt.sign(
      {
        sub: 'message-dispatcher',
        iss: 'comfyui-bridge',
        iat: Math.floor(Date.now() / 1000)
      },
      this.jwtSecret,
      { expiresIn: '24h' }
    );
  }

  pushCapacityState(pendingCount = 0) {
    const bridges = bridgeManager.getAllBridges();
    const summary = this.calculateCapacitySummary(bridges, pendingCount);
    
    const message = {
      type: 'CAPACITY_UPDATE',
      data: {
        timestamp: new Date().toISOString(),
        bridges: bridges.map(b => ({
          bridgeId: b.id,
          info: b.info
        })),
        summary
      }
    };
    
    this.send(message);
    console.log('[MDWebSocketClient] 已推送算力状态:', summary);
  }

  calculateCapacitySummary(bridges, pendingCount = 0) {
    let totalInstances = 0;
    let onlineInstances = 0;
    let lockedInstances = 0;
    let offlineInstances = 0;

    for (const bridge of bridges) {
      if (bridge.info?.instances) {
        totalInstances += bridge.info.instances.length;
        for (const instance of bridge.info.instances) {
          if (instance.status === 'online') {
            onlineInstances++;
            if (bridgeManager.isInstanceLocked(instance.id)) {
              lockedInstances++;
            }
          } else {
            offlineInstances++;
          }
        }
      }
    }

    const busyInstances = lockedInstances;
    const availableCapacity = Math.max(0, onlineInstances - busyInstances - pendingCount);

    return {
      totalBridges: bridges.length,
      totalInstances,
      onlineInstances,
      busyInstances,
      offlineInstances,
      lockedInstances,
      availableCapacity
    };
  }

  pushInstanceOnline(instanceId, bridgeId) {
    const message = {
      type: 'INSTANCE_ONLINE',
      data: {
        instanceId,
        bridgeId,
        timestamp: new Date().toISOString()
      }
    };
    this.send(message);
    console.log(`[MDWebSocketClient] 已推送实例上线: ${instanceId}`);
  }

  pushInstanceOffline(instanceId, bridgeId) {
    const message = {
      type: 'INSTANCE_OFFLINE',
      data: {
        instanceId,
        bridgeId,
        timestamp: new Date().toISOString()
      }
    };
    this.send(message);
    console.log(`[MDWebSocketClient] 已推送实例下线: ${instanceId}`);
  }

  send(message) {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(message));
    } else {
      console.warn('[MDWebSocketClient] WebSocket 未连接，无法发送消息');
    }
  }

  handleMessage(data) {
    const messageStr = data.toString();
    
    if (messageStr === 'please give me tasks') {
      console.log('[MDWebSocketClient] 收到任务请求');
      return;
    }
    
    if (messageStr === 'ping') {
      this.ws.send('pong');
      return;
    }
    
    if (messageStr === 'pong') {
      return;
    }
    
    try {
      const message = JSON.parse(messageStr);
      // console.log(`[MDWebSocketClient] 收到消息: ${message.type}`);

      switch (message.type) {
        case 'HEARTBEAT_ACK':
          break;
        default:
          console.log('[MDWebSocketClient] 未知消息类型:', message.type);
      }
    } catch (error) {
      console.error('[MDWebSocketClient] 解析消息失败:', error);
    }
  }

  startHeartbeat() {
    this.heartbeatInterval = setInterval(() => {
      const message = {
        type: 'HEARTBEAT',
        data: {
          timestamp: new Date().toISOString()
        }
      };
      this.send(message);
    }, 30000);
  }

  startTokenPushTimer() {
    this.tokenPushInterval = setInterval(() => {
      this.pushJwtToken();
    }, 20 * 60 * 60 * 1000);
  }
}

export default new MDWebSocketClient();