nodejs webSocket 压测工具

今天用nodejs 提供的WS 模块写了一个简单的压测机器人，支持多进程分配不同账号，留着以后看看

app.js 

/**
 * Created by zheng jinwei
 */
var cluster = require('cluster');
var playerCount = require('./setting.js').playerCount;
var processCount = require('./setting.js').processCount;

if (cluster.isMaster) {
    if (!processCount) {
        processCount = require('os').cpus().length;
    }
    for (var i = 0; i < processCount; i++) {
        cluster.fork();
    }

    cluster.on('fork', function (worker) {
        console.log("[fork]", worker.id);

        var app = require('./run.js');
        app.run(playerCount, worker.id);
        console.log("worker start");
    });


    cluster.on('exit', function (worker, code, signal) {
        console.log('worker ' + worker.process.pid + ' died');
    });
    cluster.on('listening', function (worker, address) {
        console.log("A worker with #" + worker.id + " is now connected to " +
        address.address +
        ":" + address.port);
    });
} else {
    //var app = require('./run.js');
    //app.run(playerCount);
    //console.log("worker start");
}

package.json:

{
  "name": "robotTool",
  "version": "0.0.1",
  "private": true,
  "scripts": {
    "start": "node app.js"
  },
  "dependencies": {
    "underscore": "*",
    "lodash":"*",
    "async": "*",
    "socket.io-client": "*",
    "ws":"*"
  }
}

player.js:

/**
 * Created by zhengjinwei.
 */

var WS = require('./webSocket.js').webSocket;
var accountIdStart = require('./setting.js').accountIdStart;
var _ = require('lodash');


//状态机:消息id
var STATUS_M = {
    STATUS_M_LOGIN_MSG: 10001
};

var Status_A = [10001];

function Player(index) {

    this.accountId = accountIdStart + index;
    this.channelId = 20001;
    this.ticket = index;
    this.imei = 1000 + index;

    this.status = STATUS_M.STATUS_M_LOGIN_MSG;

    this.socket = new WS();

    this.sleepTime = parseInt(Math.random() * 10 * 1000);

    var self = this;
    this.socket.on('data',function(data){
        self.onMessage(data);
    });
}

Player.prototype.onMessage = function(data){
    console.log("get data:"+data);
    this.run();
};

Player.prototype.login = function () {
    var data = JSON.stringify({});//此处填写具体消息
    return data;
};

Player.prototype.run = function () {
    console.log('run called');
    var data = null;
    switch (this.status) {
        case  STATUS_M.STATUS_M_LOGIN_MSG:
            data = this.login();
            break;
        default :
            break;
    }

    this.sleepTime = parseInt(Math.random() * 10 * 1000);//1~10 second

    var self = this;
    setTimeout(function () {
        var res = self.socket.send(data);
        if(!res){
            console.log("failed");
        }
        self.status = _.sample(Status_A);
    }, self.sleepTime);
};

module.exports = {
    Player: Player
};

run.js

/**
 * zheng jinwei
 * @type {exports}
 */
var async = require("async");
var _ = require("lodash");
var Player = require("./player.js").Player;
var playerCount = require('./setting').playerCount;


function run(playerCount,workerId) {
    var arr = [];
    for(var i=0;i < playerCount;i++){
        arr.push(workerId*10000+i);
    }

    console.log(arr);
    var playerMap = [];
    async.each(arr,function(item,cb){
        var role = new Player(item);
        role.run();
        playerMap.push(role);
        cb(null);
    },function(err,resp){
        console.log("all player login called");
    });
}

if (!module.parent) {
    run(playerCount,1);
} else {
    module.exports = {
        run: run
    }
}

setting.js

/**
 * Created by zheng jinwei.
 */
module.exports = {
    webSocket:{
        ip:"127.0.0.1",
        port:9999
    },
    processCount:3,
    playerCount:10,
    accountIdStart:10000
};

webSocket.js

/**
 * Created by zhengjinwei.
 */
var setting = require("./setting");
var address = "ws://"+setting.webSocket.ip + ":" + setting.webSocket.port;
var WebSocket = require('ws');
var util = require("util");
var events = require("events");//EventEmitter通过events模块来访问

function WS(){
    events.EventEmitter.call(this);

    var ws = new WebSocket(address);
    this.ws = ws;

    this.opened = false;

    var self = this;
    this.ws.on('open', function() {
        console.log("open success");
        self.opened = true;
    });
    ws.on('message', function(data, flags) {
        // flags.binary will be set if a binary data is received.
        // flags.masked will be set if the data was masked.
        self.emit('data',data);
    });
}

util.inherits(WS, events.EventEmitter);//使这个类继承EventEmitter

WS.prototype.send = function(msg){
    if(this.opened){
        this.ws.send(msg);
        return true;
    }
    return false;
};

module.exports = {
    webSocket: WS
};