All files / zyspawn zygote-pool.js

86.87% Statements 86/99
48.28% Branches 14/29
90.63% Functions 29/32
88.3% Lines 83/94

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355                                                                                    3x 3x 3x 3x               3x   8x                         16x 16x 16x 16x 16x                   16x   16x 16x 41x 41x 41x 41x 41x       41x         16x 16x 16x                           16x           16x   16x 16x 41x 41x 41x 41x         16x 16x 16x                           16x 16x               4x                               4x                                   733x                   732x         732x 732x 732x 732x       732x 732x   732x                           732x   732x   732x       732x 732x     732x                               733x 733x 733x 733x                     732x 732x 732x             732x 732x 732x                     29276x                           29276x   732x 732x     732x     732x   28544x 28544x                               733x     3x 3x 3x     3x 3x   3x 3x 3x 3x 3x 3x  
/**
 * @fileoverview
 * This module manages a pool of python zygotes and includes API to use them.
 *
 * Usage: // TODO need to be updated
 *  const {ZygotePool} = require('zygote-pool');
 *
 *  var zygotePool = new ZygotePool(10);
 *
 *  var zygoteInterface = zygotePool.request();
 *
 *  zygoteInterface.call(
 *      "file-name",
 *      "function-name",
 *      {"argument1": "value1", "argument2": "value2", "and": "more"},
 *      (err, output)=> {
 *          if (err) {
 *              // failed, need to request another zygote and restart work
 *              // no need to call zygoteInterface.done()
 *              // actually calling it will not make any effect
 *          } else {
 *              use(output.stdout, output.stderr, output.result);
 *              // call zygoteInterface.run() again if needed
 *              // after finishing using zygoteInterface
 *              zygoteInterface.done();
 *              // otherwise this zygote would not be reused
 *          }
 *      }
 *  );
 *
 * Design:
 *  TODO
 *
 * Error handling:
 *  TODO
 *
 * Dependencies:
 *  ./blocking-queue.js
 *  ./zygote-manager.js
 *  ./zygote.py
 *
 */
const _ = require('lodash');
const assert = require('assert');
const BlockingQueue = require('./blocking-queue');
const ZygoteManager = require('./zygote-manager');
const {
    ZyspawnError,
    InternalZyspawnError,
    FileMissingError,
    FunctionMissingError,
    InvalidOperationError,
    TimeoutError
} = require('./error');
 
const DEFAULT_CALLBACK = (err) => { Iif(err) throw err; };
 
/**
 * Manages a pool of zygotes. Users create and use zygotes through this class.
 */
class ZygotePool {
    /**
     * Create a zygote pool.
     * @param {number} zygoteNum The number of zygotes to use
     * @param {function(Error)} callback Called after initialization,
     *      if not specified, errors will be throwed.
     */
    constructor(zygoteNum, callback = DEFAULT_CALLBACK, debugMode = false) {
        this._isShutdown = false;
        this._totalZygoteNum = 0;
        this._zygoteManagerList = []; // TODO health check?
        this._idleZygoteManagerQueue = new BlockingQueue();
        this.addZygote(zygoteNum, callback, debugMode);
    }
 
    /**
     * Add zygotes to the pool.
     * @param {number} num Number of zygotes to add
     * @param {function(Error)} callback Called after zygotes are created,
     *                                   or error happens
     */
    addZygote(num, callback = DEFAULT_CALLBACK, debugMode = false) {
        this._totalZygoteNum += num;
 
        var jobs = [];
        for (let i = 0; i < num; i++) {
            jobs.push(new Promise((resolve) => {
                ZygoteManager.create((err, zygoteManager) => {
                    Eif (!err) {
                        this._zygoteManagerList.push(zygoteManager);
                        this._idleZygoteManagerQueue.put(zygoteManager);
                        // TODO need to consider the case of shutdown before
                        // before creating finished
                    }
                    resolve(err);
                }, "zygote.py", debugMode);
            }));
        }
 
        Promise.all(jobs).then((errs) => {
            _.pull(errs, null);
            callback(errs.length == 0 ? null : errs);
            // TODO
            // need to define error object
            // kill live zygotes when error happens?
        });
    }
 
    /**
     * Remove zygotes from the pool.
     * @param {number} num Number of zygotes to remove
     * @param {function(Error)} callback Called after zygotes are removed,
     *                                   or error happens
     */
    removeZygote(num, callback = DEFAULT_CALLBACK) {
        Iif (num > this._totalZygoteNum) {
            callback(new InvalidOperationError(
                `Trying to remove ${num} zygote(s) while totalZygoteNum is ${this._totalZygoteNum}`
            ));
        }
 
        this._totalZygoteNum -= num;
 
        var jobs = [];
        for (let i = 0; i < num; i++) {
            jobs.push(new Promise((resolve) => {
                this._idleZygoteManagerQueue.get((err, zygoteManager) => {
                    assert(!err); // BlockingQueue.clearWaiting() is never called
                    zygoteManager.shutdown((err) => { resolve(err); });
                });
            }));
        }
 
        Promise.all(jobs).then((errs) => {
            _.pull(errs, null);
            callback(errs.length == 0 ? null : errs);
            // TODO
            // need to define error object
            // kill live zygotes when error happens?
        });
    }
 
    /**
     * Shutdown ZygotePool. Stop allocating idle zygotes but working zygotes
     * won't be interrupted. All zygotes will be shutdown after they finish
     * their work.
     * @param {function(Error)} callback Called after all zygotes are shutdown.
     */
    shutdown(callback = DEFAULT_CALLBACK) {
        this._isShutdown = true;
        this.removeZygote(this._totalZygoteNum, callback);
    }
 
    /**
     * Check if this ZygotePool has been shutdown
     * @return {boolean} True if the ZygotePool has been shutdown
     */
    isShutdown() {
        return this._isShutdown;
    }
 
    /**
     * Get total number of zygotes.
     * @return {number} Total number of zygotes
     */
    totalZygoteNum() {
        return this._totalZygoteNum;
    }
 
    /**
     * Get number of idle zygotes.
     * @return {number} Number of idle zygotes
     */
    idleZygoteNum() {
        return this._idleZygoteManagerQueue.size();
    }
 
    /**
     * Get number of busy zygotes.
     * @return {number} Number of busy zygotes
     */
    busyZygoteNum() {
        return this.totalZygoteNum() - this.idleZygoteNum();
    }
 
    /**
     * Request a ZygoteIterface to use (but Zygote will not be allocated until
     * the first time of calling ZygoteIterface.run()). See implementation of
     * ZygoteInterface and allocateZygoteManager() below.
     * @return {ZygoteInterface} An interface to use the zygote.
     */
    request() {
        return new ZygoteInterface(this);
    }
 
    /**
     * Allocate a ZygoteManager for a ZygoteInterface.
     * @param {ZygoteInterface} zygoteInterface Where to allocate the ZygoteManager
     * @param {function(Error)} callback Called after a ZygoteManager is allocated
     *                                   or error happens
     */
    _allocateZygoteManager(zygoteInterface, callback) {
        Iif (this._isShutdown) {
            callback(new Error()); // TODO error type
            return;
        }
 
        this._idleZygoteManagerQueue.get((err, zygoteManager) => {
            assert(!err); // BlockingQueue.clearWaiting() is never called
            zygoteManager.startWorker((err) => {
                Iif (err) {
                    // TODO create a new Zygote?
                    callback(err); // do we need to pass the error outside
                } else {
                    zygoteInterface._initialize(zygoteManager, (callback) => {
                        this._reclaimZygoteManager(zygoteInterface, callback);
                    });
                    callback(null);
                }
            });
        });
    }
 
    /**
     * Reclaim the ZygoteManager from a ZygoteInterface.
     * @param {ZygoteInterface} zygoteInterface
     * @param {function(Error)} callback Called after the ZygoteManager is reclaimed
     *                                   or error happens
     */
    _reclaimZygoteManager(zygoteInterface, callback) {
        // TODO check if the zygote is still healthy
        var zygoteManager = zygoteInterface._zygoteManager;
        // console.log("Cleaning up! Start killing worker...");
        zygoteManager.killWorker((err) => {
            // console.log("Worker is killed!");
            Iif (err) {
                // TODO create a new Zygote?
                callback(err);
            } else {
                this._idleZygoteManagerQueue.put(zygoteManager);
                callback(null);
            }
        });
        zygoteInterface._finalize();
    }
}
 
 
/**
 * A handle object representing a working zygote and hiding implementation
 * details. It wraps a ZygoteManager internally and is registered with a
 * done() function which releases the resource.
 *
 * Notice:
 * Users must call done() method to release resources after finishing
 * their work.
 */
class ZygoteInterface {
    constructor(zygotePool) {
        this._zygotePool = zygotePool;
        this._zygoteManager = null;
        this._done = (callback) => { callback(null); };
        this._state = ZygoteInterface.UNINITIALIZED;
    }
 
    /**
     * Initialize with a ready ZygoteManager and done function
     * @param {ZygoteManager} zygoteManager A ready ZygoteManager
     * @param {function(function(Error))} done The function to release resource
     *      of which the callback will be called after resource is released or
     *      error happens.
     */
    _initialize(zygoteManager, done) {
        this._zygoteManager = zygoteManager;
        this._done = done;
        this._state = ZygoteInterface.INITIALIZED;
    }
 
    /**
     * Finalize such that this ZygoteInterface is no longer usable.
     */
    _finalize() {
        this._zygoteManager = null;
        this._done = (callback) => { callback(null); };
        this._state = ZygoteInterface.FINALIZED;
    }
 
    /**
     * Get the state of ZygoteInterface
     * @return {Number} One of the following:
     *      ZygoteInterface.UNINITIALIZED (0)
     *      ZygoteInterface.INITIALIZED (1)
     *      ZygoteInterface.FINALIZED (2);
     */
    state() {
        return this._state;
    }
 
    /**
     * Run a function in a python script (See ZygoteManager.run()).
     * @param {String} moduleName The module where the function resides
     * @param {String} functionName The function to run
     * @param {Array} arg Arguments for the function as an array
     * @param {Object} options Include optional cwd (as absolute path), paths and timeout.
     * @param {function(Error, Output)} callback Called when the result is computed
     *      or any error happens. Output contains tree fields: stdout(String),
     *      stderr(String), result(object)
     */
    call(moduleName, functionName, arg, options, callback) {
        switch (this.state()) {
            case ZygoteInterface.UNINITIALIZED:
                this._zygotePool._allocateZygoteManager(this, (err) => {
                    Iif (err) { // Failure in ZygoteManager.startWorker()
                        callback(err);
                    } else {
                        this._zygoteManager.call(moduleName, functionName, arg, options, callback);
                    }
                });
                break;
            case ZygoteInterface.INITIALIZED:
                this._zygoteManager.call(moduleName, functionName, arg, options, callback);
                break;
            case ZygoteInterface.FINALIZED:
                callback(new InvalidOperationError("Calling call() after done() on ZygoteInterface"));
                break;
            default:
                assert(false, "Bad state of ZygoteInterface: " + this.state());
        }
    }
 
    /**
     * Call the registered done function. Must be called after finishing using
     * this zygote, except that any error happens in use.
     * @param {function(Error)} callback Called after the zygote is released
     *      or error happens. If not specified, errors will be throwed.
     */
    done(callback = DEFAULT_CALLBACK) {
        this._done(callback);
    }
}
ZygoteInterface.UNINITIALIZED = 0;
ZygoteInterface.INITIALIZED = 1;
ZygoteInterface.FINALIZED = 2;
 
 
module.exports.ZygotePool = ZygotePool;
module.exports.ZygoteInterface = ZygoteInterface;
 
module.exports.ZyspawnError = ZyspawnError
module.exports.InternalZyspawnError = InternalZyspawnError;
module.exports.FileMissingError = FileMissingError;
module.exports.FunctionMissingError = FunctionMissingError;
module.exports.InvalidOperationError = InvalidOperationError;
module.exports.TimeoutError = TimeoutError;