LLM-Iterable-Thenable/src/utils/exponential-backoff.ts

/**
 * Exponential backoff is a technique used to retry a function after a delay.
 *
 * The delay increases exponentially with each attempt, up to a maximum delay.
 *
 * The jitter is a random amount of time added to the delay to prevent thundering herd problems.
 *
 * The growth rate is the factor by which the delay increases with each attempt.
 */
export class ExponentialBackoff {
  /**
   * Create a new ExponentialBackoff instance
   *
   * @param config - The configuration for the exponential backoff
   * @returns The ExponentialBackoff instance
   */
  static from(config?: Partial<ExponentialBackoffOptions>): ExponentialBackoff {
    const backoff = new ExponentialBackoff(config);
    return backoff;
  }

  /**
   * Run the function with exponential backoff
   *
   * @param fn - The function to run
   * @param onError - The callback to call when an error occurs
   * @param options - The configuration for the exponential backoff
   *
   * @throws The last error if the function fails and we have hit the max attempts
   *
   * @returns The result of the function
   */
  static run<T>(
    fn: () => Promise<T>,
    onError = (_error: Error) => {},
    options?: Partial<ExponentialBackoffOptions>,
  ): Promise<T> {
    const backoff = ExponentialBackoff.from(options);
    return backoff.run(fn, onError);
  }

  private readonly options: ExponentialBackoffOptions;

  constructor(options?: Partial<ExponentialBackoffOptions>) {
    this.options = {
      maxDelay: 10000,
      maxAttempts: 10,
      baseDelay: 1000,
      growthRate: 2,
      jitter: 0.1,
      ...options,
    };
  }

  /**
   * Run the function with exponential backoff
   *
   * If the function fails but we have not hit the max attempts, the error will be passed to the onError callback
   * and the function will be retried with an exponential delay
   *
   * If the function fails and we have hit the max attempts, the last error will be thrown
   *
   * @param fn - The function to run
   * @param onError - The callback to call when an error occurs
   *
   * @throws The last error if the function fails and we have hit the max attempts
   *
   * @returns The result of the function
   */
  async run<T>(
    fn: () => Promise<T>,
    onError = (_error: Error) => {},
  ): Promise<T> {
    let lastError: Error = new Error('Exponential backoff: Max retries hit');

    let attempt = 0;

    while (
      attempt < this.options.maxAttempts ||
      this.options.maxAttempts == 0
    ) {
      try {
        return await fn();
      } catch (error) {
        // Store the error in case we fail every attempt
        lastError = error instanceof Error ? error : new Error(`${error}`);
        onError(lastError);

        // Wait before going to the next attempt
        const delay = this.calculateDelay(attempt);
        await new Promise((resolve) => setTimeout(resolve, delay));
      }

      attempt++;
    }

    // We completed the loop without ever succeeding. Throw the last error we got
    throw lastError;
  }

  /**
   * Calculate the delay before we should attempt to retry
   *
   * NOTE: The maximum delay is (maxDelay * (1 + jitter))
   *
   * @param attempt
   * @returns The time in milliseconds before another attempt should be made
   */
  private calculateDelay(attempt: number): number {
    // Get the power of the growth rate
    const power = Math.pow(this.options.growthRate, attempt);

    // Get the delay before jitter or limit
    const rawDelay = this.options.baseDelay * power;

    // Cap the delay to the maximum. Do this before the jitter so jitter does not become larger than delay
    const cappedDelay = Math.min(rawDelay, this.options.maxDelay);

    // Get the jitter direction. This will be between -1 and 1
    const jitterDirection = 2 * Math.random() - 1;

    // Calculate the jitter
    const jitter = jitterDirection * this.options.jitter * cappedDelay;

    // Add the jitter to the delay
    return cappedDelay + jitter;
  }
}

export type ExponentialBackoffOptions = {
  /**
   * The maximum delay between attempts in milliseconds
   */
  maxDelay: number;

  /**
   * The maximum number of attempts. Passing 0 will result in infinite attempts.
   */
  maxAttempts: number;

  /**
   * The base delay between attempts in milliseconds
   */
  baseDelay: number;

  /**
   * The growth rate of the delay
   */
  growthRate: number;

  /**
   * The jitter of the delay as a percentage of growthRate
   */
  jitter: number;
};