add backoff and jitter logic into operation retry loop

Author: baileympearson

src/error.ts

@@ -99,7 +99,9 @@ export const MongoErrorLabel = Object.freeze({
   ResetPool: 'ResetPool',
   PoolRequestedRetry: 'PoolRequestedRetry',
   InterruptInUseConnections: 'InterruptInUseConnections',
-  NoWritesPerformed: 'NoWritesPerformed'
+  NoWritesPerformed: 'NoWritesPerformed',
+  SystemOverloadError: 'SystemOverloadError',
+  RetryableError: 'RetryableError'
 } as const);
 
 /** @public */

src/operations/execute_operation.ts

@@ -1,3 +1,5 @@
+import { setTimeout } from 'node:timers/promises';
+
 import { MIN_SUPPORTED_SNAPSHOT_READS_WIRE_VERSION } from '../cmap/wire_protocol/constants';
 import {
   isRetryableReadError,
@@ -26,6 +28,7 @@ import {
 import type { Topology } from '../sdam/topology';
 import type { ClientSession } from '../sessions';
 import { TimeoutContext } from '../timeout';
+import { RETRY_COST, TOKEN_REFRESH_RATE } from '../token_bucket';
 import { abortable, maxWireVersion, supportsRetryableWrites } from '../utils';
 import { AggregateOperation } from './aggregate';
 import { AbstractOperation, Aspect } from './operation';
@@ -232,11 +235,12 @@ async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFro
     session.incrementTransactionNumber();
   }
 
-  const maxTries = willRetry ? (timeoutContext.csotEnabled() ? Infinity : 2) : 1;
+  const MAX_ATTEMPTS = 5;
+  const maxTries = willRetry ? MAX_ATTEMPTS : 1;
   let previousOperationError: MongoError | undefined;
   let previousServer: ServerDescription | undefined;
 
-  for (let tries = 0; tries < maxTries; tries++) {
+  for (let attempt = 0; attempt < maxTries; attempt++) {
     if (previousOperationError) {
       if (hasWriteAspect && previousOperationError.code === MMAPv1_RETRY_WRITES_ERROR_CODE) {
         throw new MongoServerError({
@@ -246,14 +250,40 @@ async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFro
         });
       }
 
-      if (operation.hasAspect(Aspect.COMMAND_BATCHING) && !operation.canRetryWrite) {
+      const isRetryableError =
+        // any command with the RetryableError label is retryable
+        previousOperationError.hasErrorLabel(MongoErrorLabel.RetryableError) ||
+        // COMMAND_BATCHING commands are retryable depending on the contents of the batches
+        (operation.hasAspect(Aspect.COMMAND_BATCHING) && !operation.canRetryWrite) ||
+        (hasWriteAspect && isRetryableWriteError(previousOperationError)) ||
+        (hasReadAspect && !isRetryableReadError(previousOperationError));
+
+      if (!isRetryableError) {
         throw previousOperationError;
       }
 
-      if (hasWriteAspect && !isRetryableWriteError(previousOperationError))
-        throw previousOperationError;
+      const isSystemOverloadError = previousOperationError.hasErrorLabel(
+        MongoErrorLabel.SystemOverloadError
+      );
+      if (isSystemOverloadError) {
+        const delay =
+          Math.random() *
+          Math.min(
+            10_000, // MAX_BACKOFF
+            100 * 2 ** attempt
+          );
+
+        // short circuit if CSOT would expire while sleeping
+        if (timeoutContext.csotEnabled() && timeoutContext.remainingTimeMS < delay) {
+          throw previousOperationError;
+        }
+
+        await setTimeout(delay, undefined, {
+          // TODO: handle abort signal here?
+        });
+      }
 
-      if (hasReadAspect && !isRetryableReadError(previousOperationError)) {
+      if (!topology.tokenBucket.consume(RETRY_COST)) {
         throw previousOperationError;
       }
 
@@ -284,19 +314,35 @@ async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFro
     operation.server = server;
 
     try {
-      // If tries > 0 and we are command batching we need to reset the batch.
-      if (tries > 0 && operation.hasAspect(Aspect.COMMAND_BATCHING)) {
+      // If attempt > 0 and we are command batching we need to reset the batch.
+      if (attempt > 0 && operation.hasAspect(Aspect.COMMAND_BATCHING)) {
         operation.resetBatch();
       }
 
       try {
         const result = await server.command(operation, timeoutContext);
+
+        // On success, deposit the refresh rate into the bucket.  If the attempt
+        // was a retry, also deposit the retry cost to allow for another retry.
+        topology.tokenBucket.deposit(TOKEN_REFRESH_RATE);
+        if (attempt > 0) {
+          topology.tokenBucket.deposit(RETRY_COST);
+        }
+
         return operation.handleOk(result);
       } catch (error) {
         return operation.handleError(error);
       }
     } catch (operationError) {
       if (!(operationError instanceof MongoError)) throw operationError;
+
+      // If a retry fails with a non-SystemOverloadError, this indicates that the server
+      // either was not overloaded OR was overloaded but has recovered enough to let this
+      // request through.  Either way, deposit the retry cost to allow for another retry.
+      if (attempt > 0 && !operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadError)) {
+        topology.tokenBucket.deposit(RETRY_COST);
+      }
+
       if (
         previousOperationError != null &&
         operationError.hasErrorLabel(MongoErrorLabel.NoWritesPerformed)

src/sdam/topology.ts

@@ -35,6 +35,7 @@ import { type Abortable, TypedEventEmitter } from '../mongo_types';
 import { ReadPreference, type ReadPreferenceLike } from '../read_preference';
 import type { ClientSession } from '../sessions';
 import { Timeout, TimeoutContext, TimeoutError } from '../timeout';
+import { TokenBucket } from '../token_bucket';
 import type { Transaction } from '../transactions';
 import {
   addAbortListener,
@@ -201,18 +202,15 @@ export type TopologyEvents = {
  * @internal
  */
 export class Topology extends TypedEventEmitter<TopologyEvents> {
-  /** @internal */
   s: TopologyPrivate;
-  /** @internal */
   waitQueue: List<ServerSelectionRequest>;
-  /** @internal */
   hello?: Document;
-  /** @internal */
   _type?: string;
 
+  tokenBucket = new TokenBucket(1000);
+
   client!: MongoClient;
 
-  /** @internal */
   private connectionLock?: Promise<Topology>;
 
   /** @event */

src/token_bucket.ts

@@ -1,4 +1,4 @@
-class TokenBucket {
+export class TokenBucket {
   private budget: number;
   constructor(allowance: number) {
     this.budget = allowance;
@@ -15,22 +15,6 @@ class TokenBucket {
   }
 }
 
-const TOKEN_REFRESH_RATE = 0.1;
-const INITIAL_SIZE = 1000;
-const RETRY_COST = 1;
-
-export class RetryTokenBucket {
-  private bucket = new TokenBucket(INITIAL_SIZE);
-
-  onCommandSuccess(retry: boolean) {
-    this.bucket.deposit(TOKEN_REFRESH_RATE);
-
-    if (retry) {
-      this.bucket.deposit(RETRY_COST);
-    }
-  }
-
-  canRetry() {
-    return this.bucket.consume(RETRY_COST);
-  }
-}
+export const TOKEN_REFRESH_RATE = 0.1;
+export const INITIAL_SIZE = 1000;
+export const RETRY_COST = 1;