promise-local-storage.js

const async_hooks = require('async_hooks');

let store = new Map();

let promiseStorageHook = async_hooks.createHook({
  init(id, type, triggerId, resource) {
    if (store.get(triggerId)) {
      if (type === 'PROMISE') {
        store.set(id, store.get(triggerId));
      }
    }
  },
  destroy(id) {
    store.delete(id);
  }
});

export function getPromiseLocalStorage() {
  let box = store.get(async_hooks.executionAsyncId());
  return box ? box.data : {};
}

let runningDepth = 0;
function enable() {
  if (runningDepth === 0) {
    promiseStorageHook.enable();
  }

  runningDepth++;
}

function disable() {
  runningDepth--;

  if (runningDepth === 0) {
    promiseStorageHook.disable();
  }
}

export function withoutPromiseLocalStorage(func) {
  disable();

  let result = func();
  if (result instanceof Promise) {
    result.then(enable);
  } else {
    enable();
  }

  return result;
}

export function withPromiseLocalStorage(data, func) {
  enable();

  let currentData = getPromiseLocalStorage();

  const asyncResource = new async_hooks.AsyncResource('AsyncStorage');
  return asyncResource.runInAsyncScope(() => {
    store.set(async_hooks.executionAsyncId(), {
      data: { ...currentData, ...data }
    });
    let result = func();

    if (result instanceof Promise) {
      return result.finally(() => {
        store.get(
          async_hooks.executionAsyncId(),
          currentData
        ).data = currentData;
        disable();
      });
    } else {
      // We don't need to change any of the pointers back to the
      // original data because we're not returning a promise, so
      // there's no way the parent scope can "attach" itself to this
      // scope and accidentall inherit it
      disable();
    }

    return result;
  });
}

// Usage:

import {
  withPromiseLocalStorage,
  getPromiseLocalStorage
} from './promise-local-storage';

function wait(n) {
  return new Promise(resolve => setTimeout(resolve, n));
}

withPromiseLocalStorage({ name: 'james' }, async () => {
  await wait(100);
  console.log(getPromiseLocalStorage()); // { name: 'james' }
  await wait(200);

  await withPromiseLocalStorage({ name: 'sarah' }, async () => {
    await wait(200);
    console.log(getPromiseLocalStorage()); // { name: 'sarah' }
  });

  await withPromiseLocalStorage({ id: 100 }, async () => {
    await wait(200);
    console.log(getPromiseLocalStorage()); // { name: 'james', id: 100 }
  });

  console.log(getPromiseLocalStorage()); // { name: 'james' }
});