@jshaftoe/infra/src/Execution.js

//@flow

import autoBind from 'auto-bind'

import System from './System'
import Module from './Module'
import { sleep } from './asyncUtils'
import { SECOND_MS, MINUTE_MS } from './consts'

const { MODULE_NAME, log, warn, error, noteGauge, noteCount, noteTimer, trackOp } = new Module(__filename) // eslint-disable-line no-unused-vars

type ExecutionOptions = {|
  dontTerminateOnCompletion?: ?boolean, // Used for cases where code returns but some callbacks still keep the process alive (like web server)
|}

export default class Execution<T> {

  _options: ExecutionOptions
  _startTimeMS: ?number
  _endTimeMS: ?number
  _terminating: boolean = false
  _onTerminateHandlers: Array<(err: ?Error, data: ?T) => Promise<void>> = []

  /**
  * An async function that is called AFTER a termination has accured and infra resources were released,
  * This handler is called after infra resources were flushed, so it should avoid using them.
  */
  terminationHandler: ?(err: ?Error, data: ?T) => Promise<void>


  constructor(options?: ?ExecutionOptions) {
    this._options = options || {
      dontTerminateOnCompletion: false,
    }

    autoBind(this)
  }


  async run(runFunc: () => Promise<T>): Promise<void> {
    try {
      noteCount(`execution.start`)

      // Catches ctrl+c and PM2 shutdown
      process.on('SIGINT', async () => await this._terminate(new Error('SIGINT')))

      if (System.getConfig().trackVitalsIntervalMS) {
        this._trackVitals()
      }

      this._startTimeMS = Date.now()
      log('Execution started')
      const returnValue = await runFunc()

      if (!this._options.dontTerminateOnCompletion) {
        await this._terminate(null, returnValue)
      }
    } catch (err) {
      await this._terminate(err)
    } 
  }

  /**
  * Allows adding an async function that is called BEFORE a termination has accured.
  * @param {*} terminationHandler
  */
  addOnTerminateHandler(onTerminateHandler: (err: ?Error, data: ?T) => Promise<void>): void {
    this._onTerminateHandlers.push(onTerminateHandler)
  }


  async _trackVitals() {
    while (!this._terminating) {
      try {
        const used = process.memoryUsage()
        for (const key in used) {
          noteGauge(`memory.${key}MB`, Math.round(used[key] / 1024 / 1024 * 100) / 100)
        }
        await sleep(System.getConfig().trackVitalsIntervalMS || 10 * SECOND_MS)
      } catch (err) {
        error('Failed to track vitals', { err })
        await sleep(1 * MINUTE_MS)
      }
    }
  }

  async _terminate(terminationError: ?Error, data: ?T) {
    if (this._terminating)
      return
    this._terminating = true
    
    this._endTimeMS = Date.now()
    // $FlowIgnore
    const durationMS = this._endTimeMS - this._startTimeMS

    if (terminationError == null) {
      await noteCount(`execution.succeed`)
      await noteTimer('execution.succeed_duration', durationMS)
      log('Execution ended successfully', { durationMS })
    } else if (terminationError.message === 'SIGINT') {
      await noteCount(`execution.terminate`)
      await noteTimer('execution.terminate_duration', durationMS)
      warn('Execution terminated by SIGINT signal, terminating...', { durationMS })
    } else {
      await noteCount(`execution.fail`)
      await noteTimer('execution.fail_duration', durationMS)
      error('Execution threw an unhandled exception, terminating...', { durationMS, err: terminationError })
    }
    await noteCount('execution.end')
    await noteTimer('execution.end_duration', durationMS)

    // onTerniate handlers
    for (const onTerminateHandler of this._onTerminateHandlers) {
      try {
        await onTerminateHandler(terminationError, data)
      } catch (err) {
        error('one of the onTerminate handlers threw an error', { err })
      }
    }

    await System.flush()
    // NOTE: logging or tracking behind this point are not guaranteed to be transmitted

    if (!this.terminationHandler) {
      process.exit(terminationError == null ? 0 : 1)
    } else {
      // terminationHandler
      try {
        await this.terminationHandler(terminationError, data)
      } catch (err) {
        error('terminationHandler threw an error', { err })
        await System.flush()
        process.exit(1)
      }
    }
  }

}