balancer.js

/*
 * proxy.js: Responsible for proxying across all applications available to haibu.
 *
 * (C) 2010, Nodejitsu Inc.
 *
 */

var util = require('util'),
    path = require('path'),
    fs = require('fs'),
    events = require('events'),
    async = require('async'),
    colors = require('colors'),
    qs = require('querystring'),
    httpProxy = require('http-proxy'),
    haibu = require('haibu');

var matchers = {
  app: /^([\w|\-]+)\.package\.json/,
  pid: /^([\w|\-]+)\.(\d+)\.json/
};

var Balancer = exports.Balancer = function (options) {
  events.EventEmitter.call(this);
  
  var self = this;
  

  this.active    = {};
  this.httpProxy = new httpProxy.HttpProxy(options);
  this.processes = new haibu.ProcessStore(options);
  

  this.processes.on('created', function (file) {
    self._add(file, function () {
      self.emit('created', file);
    });
  });
  

  this.processes.on('removed', function (file) {
    self._remove(file);
    self.emit('removed', file);
  });
  

  this.processes.once('load', function (files) {
    function checkFile (file, next) {
      return files[file].isFile() ? self._add(file, next) : next();
    }

    async.forEach(Object.keys(files), checkFile, function (err) {
      return err ? self.emit('error', err) : self.emit('ready', self.active);
    });
  })
  

  this.processes.monitor();
};

util.inherits(Balancer, events.EventEmitter);

Balancer.prototype.handle = function (req, res) {
  var record = this.findDrone(req), drone;
  
  if (!record) {
    return this.serveText(req, res, {
      code: 400, 
      message: 'Application not found for: ' + req.headers.host
    });
  }
  else if (!record.drones || record.drones.length === 0) {
    return this.serveText(req, res, {
      code: 400, 
      message: 'No drones for: ' + req.headers.host
    });
  }
  
  drone = record.drones.shift();
  this.httpProxy.proxyRequest(req, res, drone);
  record.drones.push(drone);
};

Balancer.prototype.findDrone = function (req) {
  var self = this, 
      host = req.headers.host.split('.'),
      domain = host.slice(-2).join('.');
  
  return Object.keys(this.active).map(function (app) {
    return self.active[app];
  }).filter(function (rec) {
    return rec.app.domain === domain || (rec.app.domains
      && rec.app.domains.indexOf(domain) !== -1);
  })[0];  
};

Balancer.prototype.serveText = function (req, res, data) {
  var text = data.message; 
  
  if (data.meta) {
    text = [message, qs.unescape(qs.stringify(data.meta, ', '))].join(' | ');
  }
  
  res.writeHead(data.code, {
    'Content-Length': data.length,
    'Content-Type': 'text/plain'
  });
  
  if (req.method !== 'HEAD') {
    res.write(text);
  }
  
  res.end();
};

Balancer.prototype.close = function () {
  this.httpProxy.close();
};

Balancer.prototype._add = function (file, callback) {
  var self = this;
  
  fs.readFile(file, function (err, data) {
    if (err) {
      return callback ? callback(err) : null;
    }
    
    try {
      var json = JSON.parse(data.toString()),
          app = self._parseFilename(file);
      
      if (app && app.type) {
        switch (app.type) {
          case 'package':
            self.active[app.name] = self.active[app.name] || {
              drones: []
            };

            self.active[app.name].app = json;
            break;
          case 'pid':
            self.active[app.name] = self.active[app.name] || {
              drones: []
            };

            self.active[app.name].drones.push({
              pid: json.pid,
              host: json.host,
              port: json.port
            });
            break;
        }
      }

      return callback ? callback() : null;
    }
    catch (ex) {
      return callback ? callback(ex) : null;
    }
  });
};

Balancer.prototype._remove = function (file) {
  var app = this._parseFilename(file), index;
  
  if (app && app.type) {
    switch (app.type) {
      case 'package':
        delete this.active[app.name];
        break;
      case 'pid':
        if (this.active[app.name] && this.active[app.name].drones) {
          index = this.active[app.name].drones.map(function (d) { 
            return d.pid;
          }).indexOf(parseInt(app.drone, 10));
          
          this.active[app.name].drones.splice(index, 1);
        }
        break;
    }
  }
};

Balancer.prototype._parseFilename = function (file) {
  var base = path.basename(file),
      match, drone, app;
  
  if (matchers.app.test(base)) {
    return {
      type: 'package',
      name: base.match(matchers.app)[1]
    }
  }
  else if (matchers.pid.test(base)) {
    match = base.match(matchers.pid);
    return {
      type:  'pid',
      name:  match[1],
      drone: match[2]
    };
  }
  
  return null;
};