zero-backpressure-weighted-promise-semaphore/src/zero-backpressure-weighted-promise-semaphore.ts

A modern weighted promise semaphore for Node.js projects, ideal for managing workloads with varying processing requirements. It allows limiting the total weight of concurrently executing jobs, ensuring efficient resource utilization. Offering backpressure

/**
 * Copyright 2024 Ori Cohen https://github.com/ori88c
 * https://github.com/ori88c/zero-backpressure-weighted-promise-semaphore
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

export type SemaphoreJob<T> = () => Promise<T>;
type PromiseResolveType = (value: void | PromiseLike<void>) => void;

const DEFAULT_NUMBERS_OF_SLOTS = 16;
const NO_JOB_IS_PENDING_FOR_WEIGHT = -1;

/**
 * The `ZeroBackpressureWeightedSemaphore` class implements a Promise Semaphore for Node.js projects,
 * enabling users to limit the concurrency of *weighted* jobs.
 * Each job is associated with a natural-number weight (1, 2, 3, ...). The semaphore guarantees that the
 * total weight of concurrently executing jobs never exceeds a user-defined limit.
 * The use of natural numbers for weights is mandated to prevent floating-point precision issues inherent
 * in JavaScript.
 *
 * The weighted jobs functionality is ideal for scenarios where jobs have varying processing requirements,
 * and the backend application needs to manage load on a particular resource. For example, dispatching
 * ML jobs of differing complexities to the same EC2 machine. Although the Node.js application itself
 * does not perform the CPU-intensive ML calculations, users can still benefit from load-limiting on
 * the EC2 machine.
 *
 * This implementation does not queue pending jobs, thereby eliminating backpressure. As a result, users
 * have better control over memory footprint, which enhances performance by reducing garbage-collector
 * overhead.
 *
 * The design addresses the two primary semaphore use cases in Node.js:
 * 1. **Single Job Execution**: A sub-procedure for which the caller must wait before proceeding with
 *    its work. In this case, the job's completion time is crucial to know.
 * 2. **Multiple Jobs Execution**: In this case, the start time of a given job is crucial. Since a
 *    pending job cannot start its execution until the semaphore allows, there is no reason to add
 *    additional jobs that cannot start either.
 *    Once all the jobs are completed, some post-processing logic may be required. The API provides a
 *    designated method to wait until there are no currently-executing jobs.
 *
 * ### Modern API Design
 * Traditional semaphore APIs require explicit acquire and release steps, adding overhead and
 * responsibility on the user.
 * In contrast, `ZeroBackpressureWeightedPromiseSemaphore` manages job execution, abstracting away these
 * details and reducing user responsibility. The acquire and release steps are handled implicitly by the
 * execution methods, reminiscent of the RAII idiom in C++.
 * Method names are chosen to clearly convey their functionality.
 *
 * ### Graceful Termination
 * All the job execution promises are tracked by the semaphore instance, ensuring no dangling promises.
 * This enables graceful termination via the `waitForAllExecutingJobsToComplete` method, which is
 * particularly useful for the multiple jobs execution use-case. This can help perform necessary
 * post-processing logic, and ensure a clear state between unit-tests.
 * If your component has a termination method (`stop`, `terminate`, or similar), keep that in mind.
 *
 * ### Error Handling for Background Jobs
 * Background jobs triggered by `startExecution` may throw errors. Unlike the `waitForCompletion` case,
 * the caller has no reference to the corresponding job promise which executes in the background.
 * Therefore, errors from background jobs are captured by the semaphore and can be extracted using
 * the `extractUncaughtErrors` method. The number of accumulated uncaught errors can be obtained via
 * the `amountOfUncaughtErrors` getter method. This can be useful, for example, if the user wants to
 * handle uncaught errors only after a certain threshold is reached.
 *
 * ### Complexity
 * - **Initialization**: O(estimatedMaxNumberOfConcurrentJobs) for both time and space.
 * - **startExecution, waitForCompletion**: O(1) for both time and space, excluding the job execution itself.
 * - All the getter methods have O(1) complexity for both time and space.
 */
export class ZeroBackpressureWeightedSemaphore<T = void, UncaughtErrorType = Error> {
  private readonly _totalAllowedWeight: number;
  private readonly _availableSlotsStack: Array<number>;
  private readonly _slots: Array<Promise<T> | undefined>;

  private _availableWeight: number;
  private _amountOfCurrentlyExecutingJobs: number = 0;

  // Stores uncaught errors from background jobs triggered by `startExecution`.
  private _uncaughtErrors: UncaughtErrorType[] = [];

  // Weight allotment lock properties:
  // A pending _waitForSufficientWeight promise indicates that a job cannot start
  // due to insufficient available weight. The corresponding resolve callback is used
  // to notify a state change to the _allotWeight promise that acquired the lock.
  // This approach reminds a condition_variable in C++.
  private _pendingWeightAllotment: number = NO_JOB_IS_PENDING_FOR_WEIGHT;
  private _waitForSufficientWeight?: Promise<void>;
  private _notifyPendingAllotment?: PromiseResolveType; // Resolving the above, if in a pending-state.

  /**
   * @param totalAllowedWeight The maximum allowed sum of weights (inclusive) for jobs executed concurrently.
   * @param estimatedMaxNumberOfConcurrentJobs Estimated maximum number of concurrently executing jobs.
   *                                           A higher estimate reduces the likelihood of additional slot
   *                                           allocations during runtime. Please observe that the upper bound
   *                                           is `totalAllowedWeight`, as the minimum weight is 1.
   */
  constructor(
    totalAllowedWeight: number,
    estimatedMaxNumberOfConcurrentJobs = DEFAULT_NUMBERS_OF_SLOTS,
  ) {
    if (!isNaturalNumber(totalAllowedWeight)) {
      // prettier-ignore
      throw new Error(
        `ZeroBackpressureWeightedPromiseSemaphore expects a natural number of ` +
        `totalAllowedWeight, received ${totalAllowedWeight}`,
      );
    }

    if (!isNaturalNumber(estimatedMaxNumberOfConcurrentJobs)) {
      // prettier-ignore
      throw new Error(
        `ZeroBackpressureWeightedPromiseSemaphore expects a natural number of ` +
        `estimatedMaxNumberOfConcurrentJobs, received ` +
        `${estimatedMaxNumberOfConcurrentJobs}`,
      );
    }

    this._totalAllowedWeight = totalAllowedWeight;
    this._availableWeight = totalAllowedWeight;

    // If `estimatedMaxNumberOfConcurrentJobs` is a good estimation, no additional slots
    // will be created during runtime.
    const initialNumberOfSlots = Math.min(estimatedMaxNumberOfConcurrentJobs, totalAllowedWeight);
    this._availableSlotsStack = new Array<number>(initialNumberOfSlots).fill(0);
    for (let i = 1; i < initialNumberOfSlots; ++i) {
      this._availableSlotsStack[i] = i;
    }

    this._slots = new Array(initialNumberOfSlots).fill(undefined);
  }

  /**
   * @returns The maximum allowed sum of weights (inclusive) for jobs executed concurrently.
   */
  public get totalAllowedWeight(): number {
    return this._totalAllowedWeight;
  }

  /**
   * @returns The currently available, non-allotted amount of weight.
   */
  public get availableWeight(): number {
    return this._availableWeight;
  }

  /**
   * @returns The number of jobs currently being executed by the semaphore.
   */
  public get amountOfCurrentlyExecutingJobs(): number {
    return this._amountOfCurrentlyExecutingJobs;
  }

  /**
   * Indicates the number of uncaught errors from background jobs triggered by `startExecution`,
   * that are currently stored by the instance.
   * These errors have not yet been extracted using `extractUncaughtErrors`.
   *
   * Knowing the number of uncaught errors allows users to decide whether to process them immediately
   * or wait for further accumulation.
   *
   * @returns The number of uncaught errors from background jobs.
   */
  public get amountOfUncaughtErrors(): number {
    return this._uncaughtErrors.length;
  }

  /**
   * Resolves once the given job has *started* its execution, indicating that the semaphore has
   * allotted sufficient weight for the job.
   * Users can leverage this to prevent backpressure of pending jobs:
   * If the semaphore is too busy to start a given job `X`, there is no reason to create another
   * job `Y` until `X` has started.
   *
   * This method is particularly useful for executing multiple or background jobs, where no return
   * value is expected. It promotes a just-in-time approach, on which each job is pending execution
   * only when no other job is, thereby eliminating backpressure and reducing memory footprint.
   *
   * ### Graceful Termination
   * Method `waitForAllExecutingJobsToComplete` complements the typical use-cases of `startExecution`.
   * It can be used to perform post-processing, after all the currently-executing jobs have completed.
   *
   * ### Error Handling
   * If the job throws an error, it is captured by the semaphore and can be accessed via the
   * `extractUncaughtError` method. Users are encouraged to specify a custom `UncaughtErrorType`
   * generic parameter to the class if jobs may throw errors.
   *
   * @param backgroundJob The job to be executed once the semaphore is available.
   * @param weight A natural number representing the weight associated with the job.
   * @throws Error if the weight is not a natural number (1, 2, 3, ...).
   * @returns A promise that resolves when the job starts execution.
   */
  public async startExecution(backgroundJob: SemaphoreJob<T>, weight: number): Promise<void> {
    this._validateWeight(weight, 'startExecution');

    await this._allotWeight(weight);
    const availableSlot = this._getAvailableSlot();
    this._slots[availableSlot] = this._handleJobExecution(
      backgroundJob,
      availableSlot,
      weight,
      true,
    );
    return;
  }

  /**
   * Executes the given job in a controlled manner, once the semaphore has allotted sufficient
   * weight for the job. It resolves or rejects when the job finishes execution, returning the
   * job's value or propagating any error it may throw.
   *
   * This method is useful when the flow depends on a job's execution to proceed, such as needing
   * its return value or handling any errors it may throw.
   *
   * ### Example Use Case
   * Suppose you have a route handler that needs to perform a specific code block with limited
   * concurrency (e.g., database access) due to external constraints, such as throttling limits.
   * This method allows you to execute the job with controlled concurrency. Once the job resolves
   * or rejects, you can continue the route handler's flow based on the result.
   *
   * @param job The job to be executed once the semaphore is available.
   * @param weight A natural number representing the weight associated with the job.
   * @throws Error if the weight is not a natural number (1, 2, 3, ...).
   *         Alternatively, an error thrown by the job itself.
   * @returns A promise that resolves with the job's return value or rejects with its error.
   */
  public async waitForCompletion(job: SemaphoreJob<T>, weight: number): Promise<T> {
    this._validateWeight(weight, 'waitForCompletion');

    await this._allotWeight(weight);
    const availableSlot = this._getAvailableSlot();
    return (this._slots[availableSlot] = this._handleJobExecution(
      job,
      availableSlot,
      weight,
      false,
    ));
  }

  /**
   * This method allows the caller to wait until all *currently* executing jobs have finished,
   * meaning once all running promises have either resolved or rejected.
   *
   * This is particularly useful in scenarios where you need to ensure that all jobs are completed
   * before proceeding, such as during shutdown processes or between unit tests.
   *
   * Note that the returned promise only awaits jobs that were executed at the time this method
   * was called. Specifically, it awaits all jobs initiated by this instance that had not completed
   * at the time of invocation.
   * In certain scenarios - such as during application shutdown - you might also want to wait for
   * any potentially pending jobs. This can be achieved using the following pattern:
   * ```
   * while (semaphore.amountOfCurrentlyExecutingJobs > 0) {
   *   await semaphore.waitForAllExecutingJobsToComplete();
   * }
   * ```
   *
   * @returns A promise that resolves when all currently executing jobs are completed.
   */
  public async waitForAllExecutingJobsToComplete(): Promise<void> {
    const busySlots = this._slots.filter((job) => job !== undefined);
    if (busySlots.length > 0) {
      await Promise.allSettled(busySlots);
    }
  }

  /**
   * Returns an array of uncaught errors, captured by the semaphore while executing background
   * jobs added by `startExecution`. The term `extract` implies that the semaphore instance will
   * no longer hold these error references once extracted, unlike `get`.
   * In other words, ownership of these uncaught errors shifts to the caller, while the semaphore
   * clears its list of uncaught errors.
   *
   * Even if the user does not intend to perform error-handling with these uncaught errors, it is
   * important to periodically call this method when using `startExecution` to prevent the
   * accumulation of errors in memory.
   * However, there are a few exceptional cases where the user can safely avoid extracting uncaught
   * errors:
   * - The number of jobs is relatively small and the process is short-lived.
   * - The jobs never throw errors, thus no uncaught errors are possible.
   *
   * @returns An array of uncaught errors from background jobs triggered by `startExecution`.
   */
  public extractUncaughtErrors(): UncaughtErrorType[] {
    const errors = this._uncaughtErrors;
    this._uncaughtErrors = [];
    return errors;
  }

  private _validateWeight(weight: number, callingMethodName: string): void {
    if (weight > this._totalAllowedWeight) {
      // prettier-ignore
      throw new Error(
        `ZeroBackpressureWeightedPromiseSemaphore.${callingMethodName} received a weight of ` +
        `${weight} which exceeds the total allowed weight of ${this._totalAllowedWeight}`,
      );
    }

    if (!isNaturalNumber(weight)) {
      // prettier-ignore
      throw new Error(
        `ZeroBackpressureWeightedPromiseSemaphore.${callingMethodName} received a non-natural ` +
        `weight of ${weight}. Weights must be natural numbers: 1, 2, 3, ...`,
      );
    }
  }

  private async _allotWeight(weight: number): Promise<void> {
    while (this._waitForSufficientWeight) {
      // A previous allotment-request is being served first (FIFO order).
      await this._waitForSufficientWeight;
    }
    // At this stage, no other allotment-request holds the allotment-lock.
    // The current allotment-request will be served before any other pending request.

    if (this._availableWeight >= weight) {
      this._availableWeight -= weight;
      return;
    }

    // Currently available weight is insufficient, we need to wait for an ongoing
    // job to complete.

    // Acquire the allotment-lock:
    // No allotments will occur until this weight-request is served.
    this._pendingWeightAllotment = weight;
    this._waitForSufficientWeight = new Promise((res) => (this._notifyPendingAllotment = res));
    await this._waitForSufficientWeight;

    this._availableWeight -= weight;

    // Release the allotment-lock.
    this._pendingWeightAllotment = NO_JOB_IS_PENDING_FOR_WEIGHT;
    this._waitForSufficientWeight = undefined;
    this._notifyPendingAllotment = undefined;
  }

  private _getAvailableSlot(): number {
    // 1st priority: The pool of unused slots.
    if (this._availableSlotsStack.length > 0) {
      return this._availableSlotsStack.pop();
    }

    // 2nd priority: Create a new slot, if all are currently taken.
    const newSlot = this._slots.length;
    this._slots.push(undefined);
    return newSlot;
  }

  /**
   * Manages the execution of a given job in a controlled manner. It ensures that the
   * job is executed within the constraints of the semaphore and handles updating the
   * internal state once the job has completed.
   *
   * ### Behavior
   * - Waits for the job to either return a value or throw an error.
   * - Updates the internal state to make the allotted slot available again once the job is finished.
   * - Release the weight-allotment lock if the requested amount is available.
   *
   * @param job The job to be executed in the given slot.
   * @param allottedSlot The slot number in which the job should be executed.
   * @param weight A natural number representing the weight associated with the job.
   * @param isBackgroundJob A flag indicating whether the caller expects a return value to proceed
   *                        with its work. If `true`, no return value is expected, and any error
   *                        thrown by the job should not be propagated to the event loop.
   * @returns A promise that resolves with the job's return value or rejects with its error.
   *          Rejection occurs only if triggered by `waitForCompletion`.
   */
  public async _handleJobExecution(
    job: SemaphoreJob<T>,
    allottedSlot: number,
    weight: number,
    isBackgroundJob: boolean,
  ): Promise<T> {
    try {
      ++this._amountOfCurrentlyExecutingJobs;
      const jobResult = await job();
      return jobResult;
    } catch (err) {
      if (!isBackgroundJob) {
        // Triggered by `waitForCompletion`:
        // User is awaiting for either fulfillement or rejection.
        throw err;
      }

      // Triggered by `startExecution`:
      // A background job, the caller does not await for its return value to proceed.
      this._uncaughtErrors.push(err);
    } finally {
      --this._amountOfCurrentlyExecutingJobs;
      this._slots[allottedSlot] = undefined;
      this._availableSlotsStack.push(allottedSlot);
      this._availableWeight += weight;

      // Release the allotment lock if it is currently held, and the available
      // weight suffices to fulfill the pending allotment.
      if (this._notifyPendingAllotment && this._availableWeight >= this._pendingWeightAllotment) {
        this._notifyPendingAllotment();
      }
    }
  }
}

function isNaturalNumber(num: number): boolean {
  const floored = Math.floor(num);
  return floored >= 1 && floored === num;
}