workforce/lib/manager.js

A cluster manager inspired by Learnboost's cluster.

/**
 * Deps.
 */

var cluster = require('cluster')
  , EE = require('events').EventEmitter
  , debug = require('debug')('workforce:master')
  , path = require('path')
  , ms = require('ms')
  , os = require('os');

/**
 * Expose `Manager`.
 */

exports = module.exports = Manager;

/**
 * Defaults.
 */

exports.defaults = {
    'exit timeout': 15000
  , 'restart threshold': 5000
  , 'title': 'workforce'
  , 'working directory': process.cwd()
  , 'silent': false
  , 'workers': os.cpus().length
  , 'signals': {
      'shutdown': 'SIGTERM'
    , 'restart': 'SIGUSR2'
  }
};

/**
 * Worker execution script.
 */

exports.worker = path.join(__dirname, 'worker.js');

/**
 * Initialize a new `Manager`.
 *
 * @param {String} server
 * @api public
 */

function Manager(server){
  if (!(this instanceof Manager)) return new Manager(server);
  if ('string' != typeof server) throw new Error('server must be a path');

  this.settings = {};

  for (var key in exports.defaults) {
    this.settings[key] = this.settings.hasOwnProperty(key)
      ? this.settings[key]
      : exports.defaults[key];
  }

  this.env = process.env.NODE_ENV || 'development';
  this.server = this.resolve(server);
  this.registerSignals();

  process.title = this.get('title').replace(/\s+/g, '-') + '-master';
}

/**
 * Mixin "configurable" for option accessors.
 */

require('configurable')(Manager.prototype);

/**
 * Inherits `EventEmitter.prototype`.
 */

Manager.prototype.__proto__ = EE.prototype;

/**
 * Limits `fn` to be called only when `env` matches.
 *
 * @param {String} env
 * @param {Function} fn
 * @api public
 */

Manager.prototype.configure = function(env, fn){
  if (arguments.length === 1) env();
  if (env === this.env) fn();
  return this;
};

/**
 * Use the `plugin`.
 *
 * @param {Function} plugin
 * @return {Manager}
 * @api public
 */

Manager.prototype.use = function(plugin){
  plugin(this);
  return this;
};


/**
 * Resolves the path to the requireable `server`.
 *
 * @param {String} server
 * @return {String}
 * @api private
 */

Manager.prototype.resolve = function(server){
  var dir = this.dir = path.dirname(process.argv[1]);
  return '/' == server[0]
    ? server
    : dir + '/' + server;
};

/**
 * Starts workers.
 *
 * @param {Number|String} port
 * @param {Function} cb
 * @api public
 */

Manager.prototype.listen = function(port, cb){
  var server = this.server
    , count = this.get('workers')
    , title = this.get('title')
    , dir = this.get('working directory');

  this.on('worker', function(){
    if (--count == 0) {
      this.removeListener('worker', arguments.callee);
      this.emit('listening');
      cb && cb();
    }
  });

  process.chdir(dir);
  
  cluster.setupMaster({
      exec: exports.worker
    , args: [ title, server, port ]
    , silent: this.silent
  });

  this.fork(count);
  return this;
};

/**
 * Works `count` amount of workers.
 *
 * @param {Number} count
 * @private
 */

Manager.prototype.fork = function(count){
  count = count || 1;
  while (--count >= 0) this.createWorker();
};

/**
 * Forks a new worker.
 *
 * Attribution: "uptime" implementation
 * https://github.com/isaacs/cluster-master/blob/master/cluster-master.js#L193
 *
 * @return {cluster.Worker}
 * @api private
 */

Manager.prototype.createWorker = function(){
  var worker = cluster.fork();

  worker.forkedAt = Date.now();
  worker.pid = worker.process.pid;

  Object.defineProperty(worker, 'uptime', {
    get: function(){
      return Date.now() - this.forkedAt;
    }
  });

  worker.on('listening', this.onlistening.bind(this, worker));
  worker.on('exit', this.onexit.bind(this, worker));

  return worker;
};

/**
 * Handle for worker "listening" events.
 *
 * @param {cluster.Worker} worker
 * @api private
 */

Manager.prototype.onlistening = function(worker){
  debug('worker %s listening -- pid=%s', worker.id, worker.pid);
  this.emit('worker', worker);
};

/**
 * Handle for worker "exit" events.
 *
 * @param {cluster.Worker} worker
 * @api private
 */

Manager.prototype.onexit = function(worker){
  var suicide = worker.suicide
    , uptime = worker.uptime
    , threshold = ms(this.get('restart threshold'));

  this.emit('worker exit', worker);

  if (!suicide) {
    if (uptime > threshold) {
      this.fork(1);
    }
  }
};

/**
 * Register kill signals.
 *
 * @api private
 */

Manager.prototype.registerSignals = function(){
  var signals = this.get('signals');

  for (var key in signals) {
    debug('%s with kill -s %s %s', key, signals[key], process.pid);
  }

  process.on(signals.restart, this.restart.bind(this));
  process.on(signals.shutdown, this.shutdown.bind(this));
};

/**
 * Shutdown `worker` gracefully.
 *
 * @param {cluster.Worker} worker
 * @api private
 */

Manager.prototype.quitWorker = function(worker){
  if (!worker.process.connected) return;
  var self = this;

  worker.on('disconnect', function(){
    setTimeout(function(){
      worker.process.kill('SIGKILL');
    }, ms(self.get('exit timeout')));
  });

  worker.disconnect();
};

/**
 * Exit `worker` immediately.
 *
 * TODO: Do we really need to catch here?
 *
 * @param {cluster.Worker} worker
 * @api private
 */

Manager.prototype.destroyWorker = function(worker){
  try {
    debug('force killing worker %s', worker.pid);
    process.kill(worker.pid);
  } catch (e) {
    console.error(e);
  }
};

/**
 * Restarts by spawning a new set of workers and
 * invoking `shutdown` on the "old" workers after
 * the first "new" worker is ready.
 *
 * Attribution: modified version from learnboost/up
 * https://github.com/LearnBoost/up/blob/master/lib/up.js#L108
 *
 * @api public
 */

Manager.prototype.restart = function(cb){
  var old = Object.keys(cluster.workers).map(function(id){
    return cluster.workers[id];
  });

  this.on('worker', function listener(){
    this.removeListener('worker', listener);
    this.emit('restart');
    debug('workers restarted');
    old.forEach(this.quitWorker.bind(this));
    cb && cb();
  });

  debug('restarting workers');
  this.fork(old.length);
};

/**
 * Applys `fn` to each worker.
 *
 * @param {Function} fn
 * @api private
 */

Manager.prototype.each = function(fn){
  Object.keys(cluster.workers).forEach(function(id){
    fn(cluster.workers[id]);
  });
};

/**
 * Graceful shutdown of all workers.
 *
 * @api public
 */

Manager.prototype.shutdown = function(cb){
  var count = this.count()
    , quit = this.quitWorker.bind(this);

  debug('shutdown');

  this.on('worker exit', function(){
    if (--count === 0) {
      this.emit('shutdown');
      cb && cb();
    }
  });

  this.each(quit);
};

/**
 * Hard shutdown.
 *
 * @api public
 */

Manager.prototype.exit = function(cb){
  var count = this.count()
    , destroy = this.destroyWorker.bind(this);

  debug('exit');

  this.on('worker exit', function(){
    if (--count === 0) {
      this.emit('exit');
      cb && cb();
    }
  });

  this.each(destroy);
};

/**
 * Gets the current count of workers.
 *
 * @api private
 */

Manager.prototype.count = function(){
  return Object.keys(cluster.workers).length;
};

/**
 * Gets the pids of the current workers.
 *
 * @api private
 */

Manager.prototype.pids = function(){
  return Object.keys(cluster.workers)
    .map(function(id){
      return cluster.workers[id].process.pid;
    });
};