context-local.js

// Copyright 2016 Luca-SAS, licensed under the Apache License 2.0

'use strict';

var child_process = require('child_process');
var fs = require('fs');
var os = require('os');
var url = require('url');
var zlib = require('zlib');

var mkdirp = require('mkdirp');
var rimraf = require('rimraf');
var uuid = require('uuid');

var dataset = require('./dataset.js');
var Task = require('./task.js');

var start = process.hrtime();
var memory = Number(process.env.SKALE_MEMORY);

module.exports = Context;

function Context(args) {
  if (!(this instanceof Context))
    return new Context(args);
  this.contextId = uuid.v4();
  var nworker = Number(process.env.SKALE_WORKERS) || (os.cpus().length - 1) || 1;
  var tmp = process.env.SKALE_TMP || '/tmp';
  var self = this;
  log('workers:', nworker);
  this.env = {};
  this.worker = new Array(nworker);
  for (var i = 0; i < nworker; i++) {
    this.worker[i] = new Worker(i);
  }
  this.jobId = 0;

  process.title = 'skale-master';

  this.basedir = tmp + '/skale/' + this.contextId + '/';
  mkdirp.sync(this.basedir + 'tmp');
  mkdirp.sync(this.basedir + 'stream');

  this.datasetIdCounter = 0;  // global dataset id counter

  this.parallelize = function (localArray, nPartitions) {return dataset.parallelize(this, localArray, nPartitions);};
  this.range = function (start, end, step, nPartitions) {return dataset.range(this, start, end, step, nPartitions);};
  this.lineStream = function (stream, config) {return new dataset.Stream(this, stream, 'line', config);};
  this.objectStream = function (stream, config) {return new dataset.Stream(this, stream, 'object', config);};
  this.log = log;

  this.end = function () {
    rimraf(self.basedir, function (err) {
      if (err) log('remove', err);
    });
    for (var i = 0; i < self.worker.length; i++) {
      try {
        self.worker[i].child.disconnect();
      } catch(err) {
        console.log(err);
      }
    }
  };

  this.getReadStream = function (fileObj, opt, done) {
    done(null, fs.createReadStream(fileObj.path, opt));
  };

  this.getReadStreamSync = function (fileObj, opt) {
    return fs.createReadStream(fileObj.path, opt);
  };

  this.textFile = function (file, opt /*, nPartitions*/) {
    opt = opt || {};
    var u = url.parse(file);

    if (u.protocol === 's3:')
      return new dataset.TextS3(this, file.slice(5), opt);
    if (u.protocol === 'wasb:')
      return new dataset.TextAzure(this, file.slice(7), opt);
    return new dataset.TextLocal(this, file, opt);
  };

  this.runTask = function (task, callback) {
    function getLeastBusyWorkerId(/* preferredLocation */) {
      var wid, ntask;
      for (var i = 0; i < self.worker.length; i++) {
        if ((ntask == undefined) || (ntask > self.worker[i].ntask)) {
          ntask = self.worker[i].ntask;
          wid = i;
        }
      }
      return wid;
    }

    function serialize(task) {
      var pleft;
      var pright;
      var nodeId;
      var p = task.pid;
      var node = task.nodes[task.datasetId];
      var part = node.shufflePartitions ? node.shufflePartitions[p] : node.partitions[p];
      var pindex = {};

      // Walk through dataset ancestors to keep partition dependencies
      while (part) {
        pindex[part.datasetId] = p;
        node = task.nodes[part.parentDatasetId];
        if (!node) break;
        p = part.parentPartitionIndex;
        part = node.shufflePartitions ? node.shufflePartitions[p] : node.partitions[p];
      }

      return str = JSON.stringify(task, function(key, value) {
        if (key == 'sc') return undefined;
        if (key == 'dependencies') {
          var dep = [];
          for (var i = 0; i < value.length; i++) dep[i] = value[i].id;
          return dep;
        }
        if (key === 'pleft') pleft = value;
        else if (key === 'pright') pright = value;
        else if (key === 'id') nodeId = value;

        // For shufflePartitions (not cartesian), return only the ones used by the task.
        if (key === 'files' && ! value.path) {
          var v = {};
          v[pindex[nodeId]] = value[pindex[nodeId]];
          return v;
        }

        // For cartesian shufflePartitions, return only the ones used by the task
        if (key === 'shufflePartitions' && value[0] && value[0].files && value[0].files.path) {
          var p1 = Math.floor(task.pid / pright);
          var p2 = task.pid % pright + pleft;
          v = {};
          v[task.pid] = value[task.pid];
          v[p1] = value[p1];
          v[p2] = value[p2];
          return v;
        }
        return (typeof value === 'function' ) ? value.toString() : value;
      });
    }

    var wid = getLeastBusyWorkerId(task.nodes[task.datasetId].getPreferedLocation(task.pid));

    // Init some environment on the worker the first time we send it a task
    if (!this.worker[wid].init) {
      this.worker[wid].init = true;
      task.env = this.env;
    }

    var str = serialize(task);
    this.worker[wid].ntask++;
    //log('task size for worker ' + wid + ':', str.length);
    //log('task', str);
    if (str.length > 1000000) {
      zlib.gzip(str, {chunkSize: 65536}, function (err, res) {
        var filename = task.basedir + 'task-' + uuid.v4() + '.gz';
        fs.writeFile(filename, res, function (err) {
          if (err) throw new Error(err);
          rpc('runztask', wid, filename, function (err, res) {
            self.worker[wid].ntask--;
            callback(err, res, task);
          });
        });
      });
    } else {
      rpc('runTask', wid, str, function(err, res) {
        self.worker[wid].ntask--;
        callback(err, res, task);
      });
    }
  };

  this.runJob = function (opt, root, action, callback) {
    var jobId = this.jobId++;

    //this.getWorkers(function () {
    findShuffleStages(function(shuffleStages) {
      if (shuffleStages.length == 0) runResultStage();
      else {
        var cnt = 0;
        runShuffleStage(shuffleStages[cnt], shuffleDone);
      }
      function shuffleDone() {
        if (++cnt < shuffleStages.length) {
          runShuffleStage(shuffleStages[cnt], shuffleDone);
        } else runResultStage();
      }
    });
    //});

    function runShuffleStage(stage, done) {
      findNodes(stage, function(nodes) {
        var pid = 0, tasks = [], i, j, node;
        var index = 0, busy = 0, complete = 0, totalFiles = 0, totalSize = 0;
        stage.shufflePartitions = {};

        for (i = 0; i < stage.dependencies.length; i++) {
          node = stage.dependencies[i];
          for (j = 0; j < node.nPartitions; j++) {
            stage.shufflePartitions[pid++] = new dataset.Partition(stage.id, pid, node.id, node.partitions[j].partitionIndex);
          }
        }
        stage.nShufflePartitions = pid;

        for (i = 0; i < stage.nShufflePartitions; i++) {
          tasks.push(new Task({
            basedir: self.basedir,
            jobId: jobId,
            nodes: nodes,
            datasetId: stage.id,
            pid: i
          }));
        }

        function runNext() {
          while (busy < nworker && index < tasks.length) {
            busy++;
            self.runTask(tasks[index++], function (err, res, task) {
              stage.shufflePartitions[res.pid].files = res.files;
              busy--;
              complete++;
              var n = 0, size = 0;
              for (var f in res.files) {
                n++;
                size += res.files[f].size;
              }
              totalFiles += n;
              totalSize += size;
              log('part', task.pid, 'from worker-' + res.workerId,
                '(' + complete + '/' + tasks.length + '), shuffle out:', n, 'files,', (size/(1<<20)).toFixed(3), 'MB');
              if (complete === tasks.length) {
                log('pre-shuffle stage done, output:', totalFiles, 'files,', (totalSize / (1 << 20)).toFixed(3), 'MB');
                stage.executed = true;
                return done();
              }
              runNext();
            });
          }
        }

        log('start shuffle stage, partitions:', stage.nShufflePartitions);
        runNext();
      });
    }

    function runResultStage() {
      findNodes(root, function(nodes) {
        var tasks = [];
        log('start result stage, partitions:', root.nPartitions);
        for (var i = 0; i < root.nPartitions; i++) {
          tasks.push(new Task({
            basedir: self.basedir,
            jobId: jobId,
            nodes: nodes,
            datasetId: root.id,
            pid: i,
            action: action
          }));
        }
        callback({id: jobId}, tasks);
      });
    }

    function findNodes(node, done) {
      var nodes = {};
      interruptibleTreewalk(node, function cin(n, done) {
        done(n.shuffling && n.executed);
      }, function cout(n, done) {
        n.getPartitions(function() {
          if (nodes[n.id] == undefined) nodes[n.id] = n;
          done();
        });
      }, function() {done(nodes);});
    }

    function findShuffleStages(callback) {
      var stages = [];
      interruptibleTreewalk(root, function cin(n, done) {
        if (n.shuffling && !n.executed) stages.unshift(n);
        done(n.shuffling && n.executed);  // stage boundary are shuffle nodes
      }, function cout(n, done) {done();}, function() {callback(stages);});
    }

    function interruptibleTreewalk(n, cin, cout, done) {
      cin(n, function(uturn) { // if uturn equals true the subtree under node won't be treewalked
        if (!uturn) {
          var nDependencies = 0;
          for (var i = 0; i < n.dependencies.length; i++)
            interruptibleTreewalk(n.dependencies[i], cin, cout, function() {
              if (++nDependencies == n.dependencies.length) cout(n, done);
            });
          if (n.dependencies.length === 0) cout(n, done);
        } else cout(n, done);
      });
    }
  };

  function rpc(cmd, workerNum, args, callback) {
    self.worker[workerNum].request({cmd: cmd, args: args}, callback);
  }
}

function Worker(index) {
  this.index = index;
  this.reqid = 1;
  this.reqcb = {};
  this.ntask = 0;
  this.child = child_process.fork(__dirname + '/worker-local.js', [index, memory], {
    execArgv: memory ? ['--max_old_space_size=' + memory] : undefined
  });
  var self = this;

  this.child.on('message', function(msg) {
    if (self.reqcb[msg.id]) {
      self.reqcb[msg.id](msg.error, msg.result);
      delete self.reqcb[msg.id];
    }
  });
}

Worker.prototype.send = function(msg) {
  this.child.send(msg);
};

Worker.prototype.request = function (req, done) {
  this.reqcb[this.reqid] = done;
  this.child.send({id: this.reqid++, req: req});
};

if (process.env.SKALE_DEBUG) {
  var log =  function() {
    var args = Array.prototype.slice.call(arguments);
    var elapsed = process.hrtime(start);
    args.unshift('[master ' + (elapsed[0] + elapsed[1] / 1e9).toFixed(3) + 's]');
    console.error.apply(null, args);
  };
} else {
  log = function () {};
}