wsst.js

/*
Copyright (c) 2011 Bogdan Tkachenko <bogus.weber@gmail.com>
ts
Tool written in NodeJS that allows to make a stress test for
your application that uses WebSockets

This file is part of WebSockets Stress Test.

WebSockets Stress Test is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

/**
 * WebSockets Stress Test
 *
 * Tool written in NodeJS that allows to make a stress test for
 * your application that uses WebSockets. You can create behavior
 * scenarios that tool will run on every connection in test.
 *
 * @author Bogdan Tkachenko <bogus.weber@gmail.com>
 * @version 0.1
 */

var
    cli = require('cli'),
    WebSocket = require('ws'),
    test,
    multipleTest,
    writeJson;

var Promise = require('bluebird');


var logger = require('./logger.js');
var pid = process.pid;



var httpCalls = function(countConnections, scenario) {
    Promise.promisifyAll(scenario);
    return new Promise(function(resolve, reject) {

        var httpReq = function(x) {
            if (x < countConnections) {
                scenario.executeHTTPRequestsAsync().then(function() {
                    httpReq(x + 1);
                    if (x + 1 == countConnections) {
                        resolve(true);
                    }
                });
            }
        };

        httpReq(0);

    });

};

/*
 * add log  
 */
// winston.add(winston.transports.File, { filename: pid+'-log.log' });
// winston.remove(winston.transports.Console);


/**
 * Run single test for given scenario on given URL.
 *
 * Result passed to callback
 *
 * @param webSocketUrl
 * @param scenarioName
 * @param countConnections
 * @param cli
 * @param callback
 */
test = function(webSocketUrl, scenarioName, countConnections, cli, callback) {
    var
        i,
        startTime = (new Date()).getTime(),
        endTime = null,
        connections = [],
        scenario = require(scenarioName[0] === "/" ? scenarioName : './' + scenarioName),
        url = webSocketUrl + (scenario.path ? scenario.path : '');
    countOpened = 0;

    cli.info('Scenario: ' + scenario.name);
    cli.info(scenario.description);
    cli.info('-----------------------------\n');
    cli.info('Starting test for ' + countConnections + ' connections...');

    /*
    * first execute http reqeusts
    */

    httpCalls(countConnections, scenario)
        .then(function(rr) {
            /*
            * sockets 
            */
            for (i = 0; i < countConnections; i++) {
                (function(index) {

                    connections[index] = {
                        socket: new WebSocket(url),
                        checkpoints: []
                    };

                    var api = {
                        /**
                         * Create checkpoint in scenario
                         *
                         * @param text
                         */
                        checkpoint: function(text) {
                            var time = (new Date()).getTime(), count;

                            count = connections[index].checkpoints.length;

                            if (count > 0) {
                                connections[index].checkpoints[count - 1].end = time;
                                connections[index].checkpoints[count - 1].total = time;
                                connections[index].checkpoints[count - 1].total -= connections[index].checkpoints[count - 1].start;
                            }

                            connections[index].checkpoints.push({
                                text: text,
                                start: time,
                                end: time,
                                total: 0
                            });

                            cli.debug('Checkpoint: ' + text);
                        }
                    };

                    connections[index].socket.on('open', function() {
                        countOpened++;

                        // Add default checkpoint when connection opens

                        api.checkpoint('Connection opened ' + countOpened);
                        //   winston.info('Connection opened ' + countOpened);
                        // And run scenario on this connection
                        scenario.init(connections[index].socket, api);
                    });

                    connections[index].socket.on('close', function() {
                        var i, j, result, connectionTime;

                        // Add default checkoint when connection closed
                        //   winston.info('Connection closed');
                        api.checkpoint('Connection closed: ' + index + 1);
                        countOpened--;

                        // If we haven't any another connections
                        if (countOpened === 0) {
                            // Save test end time
                            endTime = (new Date()).getTime();

                            // Prepare result
                            result = {
                                connections: countConnections,
                                scenarioName: scenarioName,
                                url: url,
                                total: (endTime - startTime),
                                avg: 0,
                                min: null,
                                max: null,
                                checkpoints: []
                            };

                            for (i in connections) {
                                // Calculate total connection time
                                connectionTime = connections[i].checkpoints[connections[i].checkpoints.length - 1].end - connections[i].checkpoints[0].start;

                                result.avg += connectionTime;

                                if (connectionTime < result.min || result.min === null) {
                                    result.min = connectionTime;
                                }

                                if (connectionTime > result.max || result.max === null) {
                                    result.max = connectionTime;
                                }

                                // And now for every checkpoint
                                for (j in connections[i].checkpoints) {
                                    if (!result.checkpoints[j]) {
                                        result.checkpoints[j] = {
                                            text: connections[i].checkpoints[j].text,
                                            avg: 0,
                                            min: null,
                                            max: null
                                        };
                                    }

                                    result.checkpoints[j].avg += connections[i].checkpoints[j].total;

                                    if (result.checkpoints[j].min > connections[i].checkpoints[j].total || result.checkpoints[j].min === null) {
                                        result.checkpoints[j].min = connections[i].checkpoints[j].total;
                                    }

                                    if (result.checkpoints[j].max < connections[i].checkpoints[j].total || result.checkpoints[j].max === null) {
                                        result.checkpoints[j].max = connections[i].checkpoints[j].total;
                                    }
                                }
                            }

                            result.avg /= countConnections;
                            for (j in result.checkpoints) {
                                result.checkpoints[j].avg /= countConnections;
                            }

                            cli.ok('');
                            cli.ok('Test completed!');
                            cli.ok('--------------------------------------------------');
                            cli.ok('Total test time:             ' + result.total + ' ms.');
                            cli.ok('Average time per connection: ' + result.avg + ' ms.');
                            cli.ok('Minimum connection time:     ' + result.min + ' ms.');
                            cli.ok('Maximum connection time:     ' + result.max + ' ms.');
                            cli.ok('--------------------------------------------------\n');

                            cli.ok('Time profiler:');
                            cli.ok('----------------------------------------------------------------------------------');
                            cli.ok('| #\t| Average\t| Minimum\t| Maximum\t| Name')
                            cli.ok('----------------------------------------------------------------------------------');
                            for (j in result.checkpoints) {
                                cli.ok('| ' + j + '\t| '
                                    + result.checkpoints[j].avg + '\t\t| '
                                    + result.checkpoints[j].min + '\t\t| '
                                    + result.checkpoints[j].max + '\t\t| '
                                    + result.checkpoints[j].text);
                            }

                            cli.ok('----------------------------------------------------------------------------------');

                            usage.lookup(pid, options, function(err, result) {
                                //     winston.info(result);
                            });

                            if (typeof callback === 'function') {
                                callback.call(cli, result);
                            }
                        }
                    });


                })(i);
            }

        },
        function name(params) {
            cli.info(params);
        });

};

/**
 * Write test results to file in JSON format
 *
 * @param fileName
 * @param data
 */
writeJson = function(fileName, data) {
    var fs = require('fs');

    // TODO: getting error on this line. don't know how to fix
    //fs.writeFile(fileName, data);
};

/**
 * Run several tests for scenario on given URL
 *
 * @param webSocketUrl
 * @param scenarioName
 * @param countConnections
 * @param cli
 * @param callback
 */
multipleTest = function(webSocketUrl, scenarioName, countConnections, cli, callback) {
    var
        i = 0, results = [];

    var singleTest = function(result) {
        if (result) {
            results.push(result);
        }

        if (countConnections[i]) {
            test(webSocketUrl, scenarioName, countConnections[i], cli, singleTest);
            i++;
        } else {
            callback.call(cli, results);
        }
    };

    singleTest();
};

cli.setUsage(
    cli.app + ' [OPTIONS] <URL> <scenario>\n\n' +
    '\x1b[1mExample\x1b[0m:\n ' +
    ' ' + cli.app + ' --connections 100 ws://localhost:8080 myScenario.js'
);

cli.parse({
    connections: ['c', 'Single test for specified count of connections', 'int', '100'],
    connectionsList: ['l', 'Multiple tests for specified list count of connections (-l 1,10,100,1000)', 'string'],
    output: ['o', 'File to save JSON result', 'file']
});



cli.main(function(args, options) {
    if (args.length !== 2) {
        cli.fatal('Wrong number of arguments. Must be exactly 2 arguments! See `' + cli.app + ' -h` for details');
    }



    var connections;

    if (options.connectionsList) {
        connections = options.connectionsList.split(',');

        multipleTest(args[0], args[1], connections, cli, function(result) {
            if (options.output) {
                writeJson(result);
            }
        });
    } else {
        test(args[0], args[1], options.connections, cli, function(result) {
            if (options.output) {
                writeJson([result]);
            }
        });
    }
});