zhibo/ws-load-test.js

// WebSocket 并发压测（模拟进入直播间+心跳）
const WebSocket = require('ws');
const crypto = require('crypto');

const WS_HOST = process.env.WS_HOST || 'wss://im.fhhx.runtzh.com/ws/app/webSocket'; 
const LIVE_ID = process.env.LIVE_ID;                    // 必填：直播间ID
const USER_TYPE = parseInt(process.env.USER_TYPE || '0', 10);    // 与前端一致，默认0
const HEARTBEAT_MS = parseInt(process.env.HB || '15000', 10);    // 心跳间隔
const RATE = parseInt(process.env.RATE || '80', 10);             // 每秒建链速率
const DURATION_SEC = parseInt(process.env.DURATION || '600', 10);// 运行总时长(秒)

// 新增：可选 Origin 和扫码参数
const ORIGIN = process.env.ORIGIN || 'https://im.fhhx.runtzh.com';
const COMPANY_ID = process.env.COMPANY_ID;
const COMPANY_USER_ID = process.env.COMPANY_USER_ID;

const UID_START = parseInt(process.env.UID_START || '900000000', 10);
const UID_TO = parseInt(process.env.UID_TO || '0', 10);
let COUNT = parseInt(process.env.COUNT || '0', 10);

if (!LIVE_ID) {
  console.error('缺少 LIVE_ID。示例：LIVE_ID=359 node scripts/ws-load-test.js');
  process.exit(1);
}
if (UID_TO > 0 && UID_TO >= UID_START) {
  COUNT = UID_TO - UID_START + 1;
}
if (!COUNT || COUNT <= 0) {
  console.error('缺少 COUNT 或 UID_TO，无法确定连接数量');
  process.exit(1);
}

function buildUrl(userId) {
  const ts = Date.now();
  const message = `${LIVE_ID}${userId}${USER_TYPE}${ts}`;
  const signature = crypto.createHmac('sha256', String(ts)).update(message).digest('hex');
  let url = `${WS_HOST}?userId=${userId}&liveId=${LIVE_ID}&userType=${USER_TYPE}&timestamp=${ts}&signature=${signature}`;
  // 新增：可选扫码来源参数
  if (COMPANY_ID && COMPANY_USER_ID) {
    url += `&companyId=${encodeURIComponent(COMPANY_ID)}&companyUserId=${encodeURIComponent(COMPANY_USER_ID)}`;
  }
  return url;
}

const stats = { open: 0, close: 0, error: 0, msg: 0 };

function startClient(index) {
  const userId = String(UID_START + index);
  const url = buildUrl(userId);
  // 新增：传入 origin 头，避免被服务端拒绝
  const ws = new WebSocket(url, { perMessageDeflate: false, origin: ORIGIN });
  let hbTimer = null;

  ws.on('open', () => {
    stats.open++;
    console.log(`[${userId}] connected`);
    const sendHb = () => {
      const now = Date.now();
      const payload = JSON.stringify({
        cmd: 'heartbeat',
        msg: 'ping',
        userId,
        liveId: LIVE_ID,
        timestamp: now,
        networkType: 'wifi',
      });
      ws.send(payload);
    };
    // 新增：打开即发送一次心跳， subsequent 定时发送
    sendHb();
    hbTimer = setInterval(sendHb, HEARTBEAT_MS);
  });

  // 新增：打印服务端返回的 cmd，观察是否有 heartbeat/entry/out 等
  ws.on('message', (msg) => {
    stats.msg++;
    try {
      const parsed = JSON.parse(msg);
      if (parsed && parsed.data && parsed.data.cmd) {
        console.log(`[${userId}] recv cmd=${parsed.data.cmd}`);
      }
    } catch (_) {}
  });

  ws.on('error', (err) => {
    stats.error++;
    console.error(`[${userId}] error`, err?.message || err);
    if (hbTimer) clearInterval(hbTimer);
  });

  ws.on('close', (code, reason) => {
    stats.close++;
    console.log(`[${userId}] closed code=${code} reason=${reason || ''}`);
    if (hbTimer) clearInterval(hbTimer);
  });
}

async function run() {
  console.log(`开始压测: liveId=${LIVE_ID}, count=${COUNT}, rate=${RATE}/s, hb=${HEARTBEAT_MS}ms, uid=[${UID_START}..${UID_START + COUNT - 1}]`);
  for (let i = 0; i < COUNT; i++) {
    startClient(i);
    const delayMs = Math.max(1, Math.floor(1000 / RATE));
    await new Promise((r) => setTimeout(r, delayMs));
  }
  const endAt = Date.now() + DURATION_SEC * 1000;
  const t = setInterval(() => {
    console.log(`[stats] open=${stats.open} close=${stats.close} error=${stats.error} msg=${stats.msg}`);
    if (Date.now() >= endAt) {
      console.log('压测结束。');
      clearInterval(t);
      process.exit(0);
    }
  }, 2000);
}

run().catch((e) => { console.error('运行异常', e); process.exit(1); });