@jshaftoe/infra/src/ThreadPool.js

// @flow

import EventEmitter from 'events'
import autoBind from 'auto-bind'

import AsyncQueue, { QueueClosedError } from './AsyncQueue'
import { OneTimeBroadcastEvent, Semaphore } from './asyncSyncUtils'
import { ExtendedError } from './errorUtils'


type ThreadPoolOptions<T, R> = {
  name?: ?string,
  threads?: ?number,
  items?: ?Array<T>,
  task: (item: T) => Promise<R>,
  errorHandler?: ?(err: Error) => void, 
  queueMaxSize?: ?number,
}

type Task<R> = { 
  func: () => Promise<R>, 
  index: number,
}

/**
 * A thread-pool abstraction for ES6 async operations",
 * 
 * @export
 * @class ThreadPool
 * @extends {EventEmitter}
 */
export default class ThreadPool<T, R> extends EventEmitter {
  
  queuedCount: number = 0
  startedCount: number = 0
  endedCount: number = 0

  _options: ThreadPoolOptions<T, R>
  _errorHandler: (err: Error) => void
  _uncaughtErrors: Array<Error> = []
  _closed: boolean = false
  _queuedTasks: AsyncQueue<{ func: () => Promise<R>, index: number }>
  _threadsSemaphore: Semaphore
  _completeEvent: OneTimeBroadcastEvent
  _allTasksCompleteOrSomeFailedEvent: OneTimeBroadcastEvent
  _res: Array<R> = []


  static async run(options: ThreadPoolOptions<T, R>): Promise<Array<R>> {
    options.errorHandler = options.errorHandler || ((err) => { throw err })
    const tp = new ThreadPool(options)
    if (options.items) {
      await tp.queueItems(options.items)
    }
    return await tp.runAllQueued()
  }

  static async all(items: Array<T>, task: (item: T) => Promise<R>, threads: ?number): Promise<Array<R>> {
    const tp = new ThreadPool({
      task: task,
      threads: threads,
      items: items,
    })
    if (items) {
      await tp.queueItems(items)
    }
    return await tp.runAllQueued()
  }


  /**
   * Creates an instance of ThreadPool
   */
  constructor(options: ThreadPoolOptions<T, R>) {
    super()

    this._options = options
    this._options.threads = this._options.threads || Infinity

    this._errorHandler = this._options.errorHandler || ((err) => { throw err }) // eslint-disable-line no-console

    this._queuedTasks = new AsyncQueue({ 
      name: options.name,
      maxSize: options.queueMaxSize,
    })
    this._threadsSemaphore = new Semaphore(this._options.threads)
    this._completeEvent = new OneTimeBroadcastEvent(false)
    this._allTasksCompleteOrSomeFailedEvent = new OneTimeBroadcastEvent(false)

    autoBind(this)
  }

  /**
   * Queues an itme from the ThreadPool
   */
  async queueItem(item: T): Promise<void> {
    if (this._closed)
      throw new Error(`Trying to queue a job to a closed ThreadPool`)

    const index = this.queuedCount++   
    await this._queuedTasks.enqueue({ 
      func: async () => this._res[index] = await this._options.task(item), 
      index: index,
    })
  }

  /**
   * Queues an itme from the ThreadPool
   */
  async queueItems(queueItem: Array<T>): Promise<void> {
    for (const item of queueItem) {
      await this.queueItem(item)
    }
  }
  
  /**
   * Starts executing all queued tasks
   * 
   * This function should be awaited - it will return after the ThreadPool has been closed and all it's tasks completed, or after a task threw an error.
   */
  async run() {
    try {
      while (true) {

        // wait for an available task
        // TODO: also wait for errors here?
        let task
        try {
          task = await this._queuedTasks.dequeue()
        } catch (err) {
          if (err instanceof QueueClosedError)
            break
          throw err
        }
  
        // wait for an available 'thread'
        await this._threadsSemaphore.enter()
        if (this._uncaughtErrors.length) {
          break
        }
  
        this._allTasksCompleteOrSomeFailedEvent.reset()
        
        // NOTE: no await
        this._runTask(task)
      }
  
      // wait for completion
      if (this.startedCount > 0) {
        await this._allTasksCompleteOrSomeFailedEvent.wait()
      }
  
      this._throwUncaughtErrors()
    } finally {
      this._completeEvent.signal()
    }
  }

  /**
   * Starts executing all queued tasks
   * 
   * This function should not be awaited - it will return immediatly
   */
  startRun() {
    (async () => {
      try {
        await this.run()
      } catch (err) {} // eslint-disable-line no-empty
    })()
  }

  /**
   * Closes the ThreadPool for further task queueing, the ThreadPool's completion can be awaited afte'r it's called
   */
  close() {
    this._closed = true
    this._queuedTasks.close() // NOTE: NO AWAIT
    // TODO: await completion
  }

  /**
   * Closes the ThreadPool, runs it's tasks and awaits their completion
   */
  async runAllQueued(): Promise<Array<R>> {
    this.close()
    await this.run()
    return this._res
  }

  /**
   * Awaits the closing and completion of all ThreadPool tasks
   */
  async waitComplete(): Promise<Array<R>> {
    await this._completeEvent.wait()
    this._throwUncaughtErrors()
    return this._res
  }

  /**
   * Closes the ThreadPool and awaits the running and completion of all ThreadPool tasks
   */
  async closeAndWaitComplete(): Promise<Array<R>> {
    this.close()
    return await this.waitComplete()
  }

  /****************** privates ******************/

  _throwUncaughtErrors() {   
    if (this._uncaughtErrors.length > 0) {
      throw new ExtendedError(`Errors were thrown during execution of ThreadPool`, {
        threadPoolName: this._options.name,
        errorCount: this._uncaughtErrors.length,
        errorMessages: this._uncaughtErrors.map(e => e.message),
        uncaughtErrors: this._uncaughtErrors,
      })
    }
  }

  async _runTask(task: Task<R>) {
    try {
      this.startedCount++
      await task.func()
    } catch (err) {
      try {
        this._errorHandler(err)
      } catch (err2) {
        this._uncaughtErrors.push(err2)
        await this._allTasksCompleteOrSomeFailedEvent.signal()
      }
    } finally {
      this.endedCount++          
      // fire 'progress' event
      this.emit('progress', { endedCount: this.endedCount })
      this._threadsSemaphore.exit()
      if (this._threadsSemaphore.takenCount === 0) {
        this._allTasksCompleteOrSomeFailedEvent.signal()
      }
    }
  }
}