creevey/lib/server/master/pool.js

Cross-browser screenshot testing tool for Storybook with fancy UI Runner

"use strict";

Object.defineProperty(exports, "__esModule", {
  value: true
});
exports.default = void 0;

var _cluster = _interopRequireDefault(require("cluster"));

var _events = require("events");

var _types = require("../../types");

var _messages = require("../messages");

var _utils = require("../utils");

function _interopRequireDefault(obj) { return obj && obj.__esModule ? obj : { default: obj }; }

function _defineProperty(obj, key, value) { if (key in obj) { Object.defineProperty(obj, key, { value: value, enumerable: true, configurable: true, writable: true }); } else { obj[key] = value; } return obj; }

const FORK_RETRIES = 5;

class Pool extends _events.EventEmitter {
  get isRunning() {
    return this.workers.length !== this.freeWorkers.length;
  }

  constructor(config, browser) {
    super();
    this.browser = browser;

    _defineProperty(this, "maxRetries", void 0);

    _defineProperty(this, "config", void 0);

    _defineProperty(this, "workers", []);

    _defineProperty(this, "queue", []);

    _defineProperty(this, "forcedStop", false);

    this.maxRetries = config.maxRetries;
    this.config = config.browsers[browser];
  }

  async init() {
    const poolSize = this.config.limit || 1;
    this.workers = (await Promise.all(Array.from({
      length: poolSize
    }).map(() => this.forkWorker()))).filter(workerOrError => workerOrError instanceof _cluster.default.Worker);
    if (this.workers.length != poolSize) throw new Error(`Can't instantiate workers for ${this.browser} due many errors`);
    this.workers.forEach(worker => this.exitHandler(worker));
  }

  start(tests) {
    if (this.isRunning) return false;
    this.queue = tests.map(({
      id,
      path
    }) => ({
      id,
      path,
      retries: 0
    }));
    this.process();
    return true;
  }

  stop() {
    if (!this.isRunning) {
      this.emit('stop');
      return;
    }

    this.forcedStop = true;
    this.queue = [];
  }

  process() {
    const worker = this.getFreeWorker();
    const [test] = this.queue;

    if (this.queue.length == 0 && this.workers.length === this.freeWorkers.length) {
      this.forcedStop = false;
      this.emit('stop');
      return;
    }

    if (!worker || !test) return;
    worker.isRunning = true;
    const {
      id
    } = test;
    this.queue.shift();
    this.sendStatus({
      id,
      status: 'running'
    });
    this.subscribe(worker, test);
    (0, _messages.sendTestMessage)(worker, {
      type: 'start',
      payload: test
    });
    this.process();
  }

  sendStatus(message) {
    this.emit('test', message);
  }

  getFreeWorker() {
    return this.freeWorkers[Math.floor(Math.random() * this.freeWorkers.length)];
  }

  get aliveWorkers() {
    return this.workers.filter(worker => !worker.exitedAfterDisconnect);
  }

  get freeWorkers() {
    return this.aliveWorkers.filter(worker => !worker.isRunning);
  }

  async forkWorker(retry = 0) {
    _cluster.default.setupMaster({
      args: ['--browser', this.browser, ...process.argv.slice(2)]
    });

    const worker = _cluster.default.fork();

    const message = await new Promise(resolve => {
      const readyHandler = message => {
        if (!(0, _types.isWorkerMessage)(message)) return;
        worker.off('message', readyHandler);
        resolve(message);
      };

      worker.on('message', readyHandler);
    });
    if (message.type != 'error') return worker;
    this.gracefullyKill(worker);
    if (retry == FORK_RETRIES) return message.payload;
    return this.forkWorker(retry + 1);
  }

  exitHandler(worker) {
    // eslint-disable-next-line @typescript-eslint/no-misused-promises
    worker.once('exit', async () => {
      if (_utils.isShuttingDown.current) return;
      const workerOrError = await this.forkWorker();
      if (!(workerOrError instanceof _cluster.default.Worker)) throw new Error(`Can't instantiate worker for ${this.browser} due many errors`);
      this.exitHandler(workerOrError);
      this.workers[this.workers.indexOf(worker)] = workerOrError;
      this.process();
    });
  }

  gracefullyKill(worker) {
    const timeout = setTimeout(() => worker.kill(), 10000);
    worker.on('exit', () => clearTimeout(timeout));
    (0, _messages.sendShutdownMessage)(worker);
  }

  shouldRetry(test) {
    return test.retries < this.maxRetries && !this.forcedStop;
  }

  subscribe(worker, test) {
    worker.once('message', message => {
      if (!(0, _types.isWorkerMessage)(message) && !(0, _types.isTestMessage)(message)) return;
      if (message.type != 'end' && message.type != 'error') return;
      let result;

      if (message.type == 'error') {
        this.gracefullyKill(worker);
        result = {
          status: 'failed',
          ...message.payload
        };
      } else {
        result = message.payload;
      }

      const shouldRetry = result.status == 'failed' && this.shouldRetry(test);

      if (shouldRetry) {
        test.retries += 1;
        this.queue.push(test);
      }

      worker.isRunning = false;
      this.sendStatus({
        id: test.id,
        status: result.status,
        result
      });
      this.process();
    });
  }

}

exports.default = Pool;