Linked AbortController

LinkedAbortController.mjs

// LinkedAbortController, to create a chain of AbortController instances,
// inspired by .NET CancellationTokenSource.CreateLinkedTokenSource
//
// Run: node --allow-natives-syntax LinkedAbortController.mjs
//
// Gist: https://gist.github.com/noseratio/f037cd08041a72e31e3899747cb85462
// by @noseratio, MIT license

// Also see:
// https://v8.dev/features/weak-references
// https://github.com/nodejs/node/discussions/36423#discussioncomment-149062
// https://github.com/nodejs/node/discussions/36467#discussioncomment-188393
//

// unlimited number of abort listeners can be requested
const setMaxListeners = await async function() {
  if (globalThis.process?.versions?.node) {
    return (await import("events")).setMaxListeners;
  }   
  else {
    return () => undefined;
  }
}();

/**
 * Class representing a linked LinkedAbortController.
 * It will fire "abort" event for its own listeners when 
 * the parent's AbortSignal "abort" event is fired. 
 * @extends AbortController
 */
export class LinkedAbortController extends AbortController  {
  static #staticStrongRefs = new Set();
  #cleanup;

  /**
   * Create a LinkedAbortController.
   * @param {AbortSignal} parentAbortSignal - The parent AbortSignal object
   * @param {boolean} unlimitedListeners - Allow unlimited number of abort listeners.
   */  
  constructor(parentAbortSignal, unlimitedListeners = false) {
    super();

    if (unlimitedListeners) {
      setMaxListeners(Infinity, this.signal);
    }

    if (!parentAbortSignal) {
      return;
    }

    if (parentAbortSignal.aborted) {
      this.abort();
      return;
    }

    // state holds a strong ref to this object's finalizer,
    // to make sure the finalizer callback gets eventually called,
    // when this instance of LinkedAbortController is being GC'ed
    const state = {
      cleanup: null,
      finalizer: null
    };

    // we must avoid strong ref links between this instance of 
    // and LinkedAbortController and its parent AbortSignal
    const selfWeakRef = new WeakRef(this);

    const cleanup = () => state.cleanup?.();

    const onParentAbort = () => {
      selfWeakRef.deref()?.abort();
      cleanup();
    };

    state.cleanup = () => {
      state.cleanup = null;
      parentAbortSignal.removeEventListener("abort", onParentAbort);
      state.finalizer.unregister(state);
      LinkedAbortController.#staticStrongRefs.delete(state);
    }

    // register a finalizer callback for this instance of LinkedAbortController
    state.finalizer = new FinalizationRegistry(({cleanup}) => cleanup?.());
    state.finalizer.register(this, state, state);
    
    // hold a strong external ref to state
    LinkedAbortController.#staticStrongRefs.add(state);

    // listen for parent's abort event
    parentAbortSignal.addEventListener("abort", onParentAbort);
    // cleanup upon our own abort event
    this.signal.addEventListener("abort", cleanup);

    this.#cleanup = cleanup;
  }

  close() {
    this.#cleanup();
  }
}

//
// temporary tests here
//

const assert = (await import("assert")).strict;

const delay = ms => new Promise(r => setTimeout(r), ms); 

const collectGarbage = async () => {
  // only use this for testing!
  await delay(100);
  eval("%CollectGarbage('all')");
  await delay(100);
}

async function true_if_linked_signal_abort_was_fired({keep, forceGC}) {
  const parentAbortController = new AbortController();

  let linkedAbortWasFired = false;

  let linkedAbortController = new LinkedAbortController(parentAbortController.signal);
  linkedAbortController.signal.addEventListener("abort", () => linkedAbortWasFired = true);
  const weakRef = new WeakRef(linkedAbortController);

  if (!keep) {
    linkedAbortController = null;
  }

  if (forceGC) {
    await collectGarbage();
    if (!keep) {
      // linkedAbortController must be GC'ed by now
      assert.ok(weakRef.deref() === undefined);
    }
  }

  parentAbortController.abort();
  linkedAbortController?.close();

  return linkedAbortWasFired;
}

async function main() {
  assert.ok(!await true_if_linked_signal_abort_was_fired({keep: false, forceGC: true}));
  assert.ok(await true_if_linked_signal_abort_was_fired({keep: true, forceGC: true}));
  assert.ok(await true_if_linked_signal_abort_was_fired({keep: false, forceGC: false}));
  assert.ok(await true_if_linked_signal_abort_was_fired({keep: true, forceGC: false}));

  console.log('All tests passed.')
}

await main();