context.ts

import { createHook } from 'async_hooks'

export interface HookInfo<T = any> {
  id: number
  triggerId: number
  activated: boolean
  parent?: HookInfo<T>
  contextData?: T
}

const ROOT_ID = 1

export class Context<T = any> {
  private readonly hookMap = new Map<number, HookInfo<T>>()
  private currentId: number = ROOT_ID

  constructor() {
    // init root
    this.hookMap.set(ROOT_ID, { id: ROOT_ID, triggerId: 0, activated: true })
    // setupHook
    this.setupHook()
  }

  getContextData() {
    const hi = this.hookMap.get(this.currentId)
    return hi && hi.contextData
  }

  setContextData(data: T) {
    const hi = this.hookMap.get(this.currentId)
    if (hi) {
      hi.contextData = data
    }
  }

  private setupHook() {
    // createHook
    createHook({
      init: (asyncId, _, triggerId) => {
        let parent = this.hookMap.get(triggerId)
        if (!parent) {
          triggerId = ROOT_ID
          parent = this.hookMap.get(ROOT_ID)
        }

        this.hookMap.set(asyncId, {
          id: asyncId,
          activated: false,
          parent,
          triggerId
        })
      },
      before: asyncId => {
        this.currentId = asyncId
        const hookInfo = this.hookMap.get(asyncId)
        if (hookInfo) {
          if (!hookInfo.activated) {
            // extend parent info
            const parent = this.findActivatedNode(hookInfo)
            if (parent) {
              hookInfo.parent = parent
              hookInfo.triggerId = parent.id // save relation
              hookInfo.contextData = parent.contextData // extend context data
            }
          }
          hookInfo.activated = true
        }
      },
      after: asyncId => {
        if (asyncId === this.currentId) {
          this.currentId = ROOT_ID
        }
      },
      destroy: asyncId => {
        this.hookMap.has(asyncId) && this.hookMap.delete(asyncId)
      }
    }).enable()
  }

  private findActivatedNode(hi: HookInfo<T>): HookInfo<T> {
    if (!hi) {
      return this.hookMap.get(ROOT_ID)
    }
    if (hi.activated) {
      return hi
    }
    return this.findActivatedNode(hi.parent)
  }
}

const delay = (n: number) => new Promise(resolve => setTimeout(resolve, n))

const ctx = new Context<any>()

const child = async (i: string) => {
  console.log('child start >>>>> ', i, ctx.getContextData())
  await delay(500)
  console.log('child end >>>>> ', i, ctx.getContextData())
}

const childSet = async (i: string) => {
  ctx.setContextData('test-childSet-' + i)
  console.log('childSet start >>>>> ', i, ctx.getContextData())
  await delay(1000)
  console.log('childSet end >>>>> ', i, ctx.getContextData())
}

const run = async (i: string) => {
  ctx.setContextData('test-' + i)
  await child(i)
  await childSet(i)
}

Promise.all([
  Promise.resolve().then(() => run('1')),
  Promise.resolve().then(() => run('2'))
])