AsyncWorkTracker.ts

import { isPromise } from '../jsutils/isPromise.js';
import type { PromiseOrValue } from '../jsutils/PromiseOrValue.js';

/** @internal */
export class AsyncWorkTracker {
  pendingAsyncWork: Set<Promise<void>>;

  constructor() {
    this.pendingAsyncWork = new Set<Promise<void>>();
  }

  add(promise: Promise<unknown>): void {
    const pendingAsyncWork = this.pendingAsyncWork;
    const promiseToSettle = promise.then(
      () => {
        pendingAsyncWork.delete(promiseToSettle);
      },
      () => {
        pendingAsyncWork.delete(promiseToSettle);
      },
    );
    pendingAsyncWork.add(promiseToSettle);
  }

  addValues(values: ReadonlyArray<PromiseOrValue<unknown>>): void {
    for (const value of values) {
      if (isPromise(value)) {
        this.add(value);
      }
    }
  }

  async wait(): Promise<void> {
    await Promise.resolve();
    return waitForPendingSet(this.pendingAsyncWork);
  }

  promiseAllTrackOnReject<T>(
    values: ReadonlyArray<PromiseOrValue<T>>,
  ): Promise<Array<T>> {
    const promise = Promise.all(values);
    promise.then(undefined, () => {
      this.addValues(values);
    });
    return promise;
  }

  promiseTrackPending<T, TResult>(
    values: ReadonlyArray<PromiseOrValue<T>>,
    combinator: (promises: ReadonlyArray<Promise<T>>) => Promise<TResult>,
  ): Promise<TResult> {
    const promises = values.map((value) => Promise.resolve(value));
    const settled = promises.map(() => false);

    for (let index = 0; index < promises.length; index++) {
      const promise = promises[index];
      promise.then(
        () => {
          settled[index] = true;
        },
        () => {
          settled[index] = true;
        },
      );
    }

    const trackPending = () => {
      for (let index = 0; index < promises.length; index++) {
        if (!settled[index]) {
          this.add(promises[index]);
        }
      }
    };

    const promise = combinator(promises);
    promise.then(trackPending, trackPending);
    return promise;
  }
}

async function waitForPendingSet(
  pendingPromises: ReadonlySet<Promise<void>>,
): Promise<void> {
  while (pendingPromises.size > 0) {
    // eslint-disable-next-line no-await-in-loop
    await Promise.allSettled(Array.from(pendingPromises));
  }
}