elwt/lib/pool.js

Worker threads pool manager

const { MessageChannel } = require('worker_threads');
const EventEmitter = require('events');
const { actions, moveSharedToRaw, roundRobin, serialize, deserialize } = require('./tools');
const defaults = {
  size: require('os').cpus().length,
  Storage: require('./storage'),
  PoolWorker: require('./worker'),
  templater: require('./templater'),
  roundRobin,
};

class Pool extends EventEmitter {
  constructor ({ queue, free, active, cache, template, PoolWorker, unitProps, autoRespawn } = {}) {
    super();
    this._size = 0;
    this.queue = queue;
    this.free = free;
    this.active = active;
    this.cache = cache;
    this.template = template;
    this.PoolWorker = PoolWorker;
    this.unitProps = unitProps;
    this.autoRespawn = autoRespawn;
    if (this.cache) {
      this.exec = this._execCached;
    };
    this.on(actions.FREE, this.next);
  }

  static async spawnPool ({
    size = defaults.size,
    queue = new defaults.Storage(),
    free = new defaults.Storage(),
    active = new defaults.Storage(),
    cache = false,
    PoolWorker = defaults.PoolWorker,
    templater = defaults.templater,
    roundRobin = defaults.roundRobin,
    unitProps = {},
    autoRespawn = true,
  } = {}) {
    let template = templater();
    await free.clear();
    await active.clear();
    let pool = new this({ queue, free, active, cache, template, PoolWorker, unitProps, autoRespawn });
    pool.rr = {};
    pool.rr.queue = roundRobin(pool.queue);
    pool.rr.free = roundRobin(pool.free);
    pool.rr.active = roundRobin(pool.active);
    await pool.toSize(size);
    return pool;
  }

  async addUnit () {
    let unit = new this.PoolWorker({ template: this.template, props: this.unitProps });
    unit.on('online', async () => {
      this.emit(actions.WORKER_ONLINE, { threadId: unit.threadId });
    });
    unit.on('error', this.autoRespawn ? async (e) => {
      let threadId = unit.threadId;
      this.emit(actions.WORKER_ERROR, { threadId, e });
      if (await this.free.delete(threadId) || await this.active.delete(threadId)) {
        let [ respawnedThreadId, respawnedUnit ] = await this.addUnit();
        this.emit(actions.FREE, { target: respawnedUnit });
      } else {
        this.emit(actions.FREE);
      };
    } : async (e) => {
      let threadId = unit.threadId;
      this.emit(actions.WORKER_ERROR, { threadId, e });
      await this.free.delete(threadId);
      await this.active.delete(threadId);
    });
    unit.on('exit', async (exitCode) => {
      unit.unref();
      this.emit(actions.WORKER_TERMINATED, { threadId: unit.threadId, exitCode });
    });
    await this.free.set(unit.threadId, unit);
    return [ unit.threadId, unit ];
  }

  async size () { return this._size; }

  async toSize (value) {
    try {
      if (value >= Number.MAX_SAFE_INTEGER) value = Number.MAX_SAFE_INTEGER - 1;
      let old = this._size;
      if (value === old) return true;
      if (value < 0) value = 0;
      if (value === 0) {
        this._size = 0;
        await this.free.clear();
        return true;
      };
      let active = this.active.size;
      if (value < old) {
        if (value <= active) {
          this._size = value;
          await this.free.clear();
          return true;
        } else {
          let target = value - active;
          for (let key of this.free.keys()) {
            if (this.free.size <= target) break;
            await this.free.delete(key);
          };
          this._size = value;
          return true;
        };
      } else {
        this._size = value;
        for (let i = 0; i < value - old; i++) {
          let [ threadId, unit ] = await this.addUnit();
          this.emit(actions.FREE, { target: unit });
        };
        return true;
      };
    } catch (e) {
      throw e;
    };
  }

  async activateUnit (unit) {
    await this.free.delete(unit.threadId);
    await this.active.set(unit.threadId, unit);
    return unit;
  }

  async loadUnit (target = null) {
    if (target) return await this.activateUnit(target);
    let item = this.rr.free.next();
    if (item.done) item = this.rr.active.next();
    let [ threadId, unit ] = item.done ? await this.addUnit() : item.value;
    return await this.activateUnit(unit);
  }

  async loadTask () {
    if (! this.queue.size) return null;
    let taskIterator = this.rr.queue.next();
    if (taskIterator.done) return null;
    return taskIterator.value;
  }

  async next ({ task = null, target = null } = {}) {
    let loaded = task || await this.loadTask();
    if (! loaded) return true;
    let [ fn, { data, go, stop, respawn } ] = loaded;
    if (! await this.queue.delete(fn)) return true;
    let unit = await this.loadUnit(target);
    let threadId = unit.threadId;
    let raw = {};
    await moveSharedToRaw(data, raw);
    let channel = new MessageChannel();
    channel.port2.on('close', async () => {
      channel.port2.unref();
    });
    channel.port2.on('message', async ({ action, result }) => {
      if (await this.active.has(threadId)) await this.active.delete(threadId);
      if (this.cache) await this.cache.engine.set(await this.cache.keygen(fn, data), result);
      channel.port2.close();
      await this.free.set(threadId, unit);
      this.next();
      if (action === actions.ERROR) {
        if (respawn) return go(await this.exec(fn, data, { respawn: respawn - 1 }));
        let errorDeserialized = await deserialize(result)
        console.log('ERROR! ', errorDeserialized)
        return stop(new Error(errorDeserialized));
      } else if (action === actions.DONE) {
        return go(await deserialize(result));
      };
    });
    unit.postMessage({
      action: actions.RUN,
      port: channel.port1,
      fn: await serialize(fn),
      data: data instanceof SharedArrayBuffer || ! data ? data : await serialize(data) || {},
      raw,
    }, [channel.port1]);
  }

  async exec (fn, data, { respawn = 0 } = {}) {
    return new Promise(async (go, stop) => {
      let task = [ fn, { data, go, stop, respawn: Math.abs(respawn) } ];
      await this.queue.set(...task);
      await this.next({ task });
    });
  }
  
  async _execCached (fn, data, { respawn = 0 } = {}) {
    return new Promise(async (go, stop) => {
      let ckey = await this.cache.keygen(fn, data);
      if (await this.cache.has(ckey)) return go(await deserialize(await this.cache.engine.read(ckey)));
      let task = [ fn, { data, go, stop, respawn: Math.abs(respawn) } ];
      await this.queue.set(...task);
      await this.next({ task });
    });
  }
};

module.exports = Pool