track-promise-redirect.js

import ah from "async_hooks";

const print = process._rawDebug;

const knownPromises = new Map();
const promiseRecords = new WeakMap();

let promiseIdStack = [];
const getCurrentPromiseId = () => promiseIdStack.slice(-1)[0];

const unlink = (asyncId) => {
  const record = knownPromises.get(asyncId);
  if (!record) return;
  // print(`unlink ${asyncId} refs=${record.refs}`);
  if (record.refs === 1) {
    knownPromises.delete(asyncId);
    if (record.redirectedTo) {
      unlink(record.redirectedTo);
    }
  } else {
    record.refs -= 1;
  }
};

const PromisePrototypeThen = Promise.prototype.then;

const RawStackTrace = (_, stack) => stack;

const getStack = () => {
  const previousPrepareStackTrace = Error.prepareStackTrace;
  const previousStackTraceLimit = Error.stackTraceLimit;
  Error.prepareStackTrace = RawStackTrace;
  Error.stackTraceLimit = 10;
  const stack = new Error().stack;
  Error.prepareStackTrace = previousPrepareStackTrace;
  Error.stackTraceLimit = previousStackTraceLimit;
  return stack.slice(1);
};

// Sanity check that the resolution actually triggers the expected hooked promise
let activeResolutionId = 0;
const wrapTrackedResolution = (promiseId, resolution) =>
  typeof resolution === "function"
    ? function (value) {
        const previousActive = activeResolutionId;
        if (previousActive) {
          print(`unexpected active resolution ${activeResolutionId}`);
        }
        activeResolutionId = promiseId;
        try {
          return Reflect.apply(resolution, this, [value]);
        } finally {
          activeResolutionId = previousActive;
        }
      }
    : resolution;

const wrapResolution = (resolution) =>
  typeof resolution === "function"
    ? function (value) {
        return Reflect.apply(resolution, this, [value]);
      }
    : resolution;

// If `.then` is alone on the stack, it's because the engine
// is trying to assimilate the promise.
// In that case, the promise trying to assimilate will have its
// `before` and `after` hooks trigger around this invocation
// See [NewPromiseResolveThenableJob](https://tc39.es/ecma262/#sec-newpromiseresolvethenablejob)
// Queued from [steps 13-15 of Promise Resolve Functions](https://tc39.es/ecma262/#sec-promise-resolve-functions)
Promise.prototype.then = function (onFulfilled, onRejected) {
  const currentPromiseId = getCurrentPromiseId();
  const thisPromiseRecord = promiseRecords.get(this) || { asyncId: 0, refs: 0 };
  const thisPromiseId = thisPromiseRecord.asyncId;
  // print(`then thisId=${thisPromiseId}, currentId=${currentPromiseId}`);
  const redirectedPromiseRecord =
    currentPromiseId && knownPromises.get(currentPromiseId);
  if (redirectedPromiseRecord) {
    const stack = getStack();
    // print(
    //   `then thisId=${thisPromiseId}, currentId=${currentPromiseId}, stackLength=${stack.length}`
    // );
    if (stack.length === 1) {
      redirectedPromiseRecord.redirectedTo = thisPromiseId;
      thisPromiseRecord.refs += 1;
      print(`redirect ${currentPromiseId} to ${thisPromiseId}`);
      return Reflect.apply(PromisePrototypeThen, this, [
        wrapTrackedResolution(currentPromiseId, onFulfilled),
        wrapTrackedResolution(currentPromiseId, onRejected),
      ]);
    }
  }
  return Reflect.apply(PromisePrototypeThen, this, [
    wrapResolution(onFulfilled),
    wrapResolution(onRejected),
  ]);
};

ah.createHook({
  init(asyncId, type, triggerAsyncId, resource) {
    // print(
    //   `init: asyncId=${asyncId}, type=${type}, triggerAsyncId=${triggerAsyncId}`
    // );
    if (type === "PROMISE") {
      const record = { asyncId, redirectedTo: 0, refs: 1 };
      knownPromises.set(asyncId, record);
      promiseRecords.set(resource, record);
    }
  },
  before(asyncId) {
    // print(`before: asyncId=${asyncId}, stackLength=${promiseIdStack.length}`);
    if (promiseIdStack.length) {
      print(`unexpected promises on the stack: ${promiseIdStack.join(",")}`);
    }
    if (knownPromises.has(asyncId)) {
      promiseIdStack.push(asyncId);
    }
  },
  after(asyncId) {
    // print(`after: asyncId=${asyncId}`);
    if (knownPromises.has(asyncId)) {
      const expected = promiseIdStack.pop();
      if (expected !== asyncId) {
        print(`unexpected after for ${asyncId}, `);
        if (expected) promiseIdStack.push(expected);
      }
    }
  },
  promiseResolve(asyncId) {
    // print(`promiseResolve: asyncId=${asyncId}`);
    if (activeResolutionId) {
      if (asyncId !== activeResolutionId) {
        print(
          `unexpected resolution ${asyncId}, expected=${activeResolutionId}`
        );
      }
    }
  },
  destroy(asyncId) {
    // print(`destroy: asyncId=${asyncId}`);
    unlink(asyncId);
  },
}).enable();

const makePromiseKit = () => {
  let resolve, reject;
  const promise = new Promise((res, rej) => {
    resolve = res;
    reject = rej;
  });
  return { promise, resolve, reject };
};

const crank = () => new Promise((res) => setTimeout(res, 0));

const outsideP = (() => {
  print("start");
  const pk0 = makePromiseKit();
  print(`pk0.promise=${promiseRecords.get(pk0.promise).asyncId}`);
  const af = async () => {
    print("step1");
    await null;
    print("after await");
    return pk0.promise;
  };
  print("step0");
  const p1 = af();
  print(`p1=${promiseRecords.get(p1).asyncId}`);
  print("step2");
  const p1ThenRet = p1.then((val) => {
    print(`value=${val}`);
    const pk2 = makePromiseKit();
    print(`pk2.promise=${promiseRecords.get(pk2.promise).asyncId}`);

    // Attempt to trip up our custom `.then` stack detection
    // This should end up redirecting p3 to the legitimate promise returned by the bound then
    // If the stack detection was faulty p3 could be mistakenly found redirected to pk0.promise
    const p3 = pk0.promise.then(
      pk0.promise.then.bind(pk0.promise, pk2.resolve, pk2.reject)
    );
    print(`p3=${promiseRecords.get(p3).asyncId}`);
    const ret = pk2.promise.then(async () => {
      const { redirectedTo, asyncId } = promiseRecords.get(p3);
      print(`test asyncId=${asyncId}, redirectedTo=${redirectedTo}`);
    });
    print(`ret=${promiseRecords.get(ret).asyncId}`);
    return ret;
  });
  print(`p1ThenRet=${promiseRecords.get(p1ThenRet).asyncId}`);
  print("step3");
  pk0.resolve("hello");
  print("end");
  return p1ThenRet;
})().then(async () => {
  print("await before gc");
  await crank();
  print("before gc");
  gc();
  await crank();
  for (const { asyncId, refs, redirectedTo } of knownPromises.values()) {
    print(`leftover ${asyncId}, refs=${refs}, redirectedTo=${redirectedTo}`);
  }
});
print(`outsideP=${promiseRecords.get(outsideP).asyncId}`);

zzz-output.txt

start
pk0.promise=8
step0
step1
p1=9
step2
p1ThenRet=12
step3
end
outsideP=13
after await
redirect 9 to 8
value=hello
pk2.promise=20
p3=21
ret=22
redirect 12 to 22
redirect 21 to 24
test asyncId=21, redirectedTo=24
redirect 22 to 29
await before gc
redirect 13 to 32
before gc
leftover 13, refs=1, redirectedTo=32
leftover 32, refs=2, redirectedTo=0
leftover 35, refs=1, redirectedTo=0
leftover 36, refs=1, redirectedTo=0
leftover 37, refs=1, redirectedTo=0
leftover 39, refs=1, redirectedTo=0