queue-manager-async.ts

export type QueueManagerOptions = {
  parallelism?: number;
};

export type QueueManagerDeadlineOptions = QueueManagerOptions & {
  deadline: number;
};

export type WaitForTurnOptions = {
  signal: AbortSignal;
};

export type Task<T = unknown> = (options: WaitForTurnOptions) => Promise<T> | T;

function withResolvers<T>(): {
  promise: Promise<T>;
  resolve: (value: T | PromiseLike<T>) => void;
  reject: (reason?: unknown) => void;
} {
  let resolve!: (value: T | PromiseLike<T>) => void;
  let reject!: (reason?: unknown) => void;
  const promise = new Promise<T>((res, rej) => {
    resolve = res;
    reject = rej;
  });
  return { promise, resolve, reject };
}


class Entry<T = unknown> {
  fn: Task<T>;
  abortController: AbortController;
  constructor(fn: Task<T>, abortController: AbortController = new AbortController()) {
    this.fn = fn;
    this.abortController = abortController;
  }
}

const abortError: Error & { isAbortError?: boolean } = new Error('abort');
abortError.isAbortError = true;

export class QueueManager extends EventTarget {
  parallelism: number;
  runningEntries: Entry<any>[];
  queue: Entry<any>[];

  constructor({
    parallelism = 1,
  }: QueueManagerOptions = {}) {
    super();

    this.parallelism = parallelism;

    this.runningEntries = [];
    this.queue = [];
  }

  isIdle(): boolean {
    return this.runningEntries.length === 0;
  }

  next(n: number = 1): boolean {
    let i: number;
    for (i = 0; i < n; i++) {
      const entry = this.runningEntries.shift();
      if (entry) {
        entry.abortController.abort(abortError);
        continue;
      }
    }
    // return whether we were able to abort all the entries
    return i === n;
  }

  async waitForTurn<T>(fn: Task<T>): Promise<T> {
    const entry = new Entry<T>(fn);
    return await this.#waitForTurnEntry(entry);
  }

  waitUntilIdle(): Promise<void> {
    if (this.isIdle()) {
      return Promise.resolve();
    }
    const { promise, resolve } = withResolvers<void>();
    const onIdleChange = (e: Event) => {
      const { idle } = (e as MessageEvent<{ idle: boolean }>).data;
      if (idle) {
        this.removeEventListener('idlechange', onIdleChange);
        resolve();
      }
    };
    this.addEventListener('idlechange', onIdleChange);
    return promise;
  }

  async #waitForTurnEntry<T>(entry: Entry<T>): Promise<T> {
    if (this.runningEntries.length < this.parallelism) {
      this.runningEntries.push(entry as unknown as Entry<any>);

      if (this.runningEntries.length === 1) {
        this.dispatchEvent(new MessageEvent('idlechange', {
          data: {
            idle: false,
          },
        }));
      }

      let result!: T;
      let error: unknown;
      try {
        const { fn, abortController } = entry;
        const { signal } = abortController;
        result = await fn({
          signal,
        });
      } catch (err) {
        error = err;
      }

      const index = this.runningEntries.indexOf(entry as unknown as Entry<any>);
      if (index !== -1) {
        this.runningEntries.splice(index, 1);
      }
      if (this.queue.length > 0) {
        const nextEntry = this.queue.shift();
        if (nextEntry) {
          // Fire and forget; no need to await
          this.#waitForTurnEntry(nextEntry);
        }
      } else {
        if (this.runningEntries.length === 0) {
          this.dispatchEvent(new MessageEvent('idlechange', {
            data: {
              idle: true,
            },
          }));
        }
      }

      if (error === undefined || error === abortError) {
        return result;
      } else {
        throw error;
      }
    } else {
      const { fn, abortController } = entry;
      const { promise, resolve, reject } = withResolvers<T>();
      const fn2: Task<T> = async (options) => {
        try {
          const r = await fn(options);
          resolve(r);
          return r;
        } catch (err) {
          reject(err);
          throw err;
        }
      };
      const entry2 = new Entry<T>(fn2, abortController);
      this.queue.push(entry2 as unknown as Entry<any>);
      const result = await promise;
      return result;
    }
  }
}

export class QueueManagerAbortError extends Error {}

export class QueueManagerDeadline extends QueueManager {
  deadline: number;
  private pendingAfterDeadlineRejects: Array<(reason?: unknown) => void> = [];
  private idleAbortListenerAttached = false;
  constructor(opts: QueueManagerDeadlineOptions) {
    super(opts);
    this.deadline = opts.deadline;
  }

  private ensureIdleAbortListener() {
    if (this.idleAbortListenerAttached) return;
    this.idleAbortListenerAttached = true;
    const onIdleChange = (e: Event) => {
      const { idle } = (e as MessageEvent<{ idle: boolean }>).data;
      if (idle) {
        this.removeEventListener('idlechange', onIdleChange);
        this.idleAbortListenerAttached = false;
        for (const r of this.pendingAfterDeadlineRejects) {
          r(new QueueManagerAbortError('QueueManagerDeadline: deadline passed'));
        }
        this.pendingAfterDeadlineRejects.length = 0;
      }
    };
    this.addEventListener('idlechange', onIdleChange);
  }

  async waitForTurn<T>(fn: Task<T>): Promise<T> {
    return await super.waitForTurn(async (options) => {
      if (Date.now() <= this.deadline) {
        return await fn(options);
      } else {
        this.ensureIdleAbortListener();

        const { promise, reject } = withResolvers<T>();
        this.pendingAfterDeadlineRejects.push(reject);
        return await promise;
      }
    });
  }
}

export class MultiQueueManager {
  opts: QueueManagerOptions;
  queueManagers: Map<string, QueueManager>;

  constructor(opts: QueueManagerOptions = {}) {
    this.opts = opts;

    this.queueManagers = new Map();
  }

  async waitForTurn<T>(key: string, fn: Task<T>): Promise<T> {
    let queueManager = this.queueManagers.get(key);
    if (!queueManager) {
      queueManager = new QueueManager(this.opts);
      this.queueManagers.set(key, queueManager);
      queueManager.addEventListener('idlechange', (e: Event) => {
        const { idle } = (e as MessageEvent<{ idle: boolean }>).data;
        if (idle) {
          this.queueManagers.delete(key);
        }
      });
    }
    return await queueManager.waitForTurn(fn);
  }
}