WIP feat(DRIVERS-3239): add exponential backoff in operation retry loop

src/cmap/connect.ts

@@ -224,6 +224,7 @@ export interface HandshakeDocument extends Document {
   compression: string[];
   saslSupportedMechs?: string;
   loadBalanced?: boolean;
+  backpressure: true;
 }
 
 /**
@@ -241,6 +242,7 @@ export async function prepareHandshakeDocument(
 
   const handshakeDoc: HandshakeDocument = {
     [serverApi?.version || options.loadBalanced === true ? 'hello' : LEGACY_HELLO_COMMAND]: 1,
+    backpressure: true,
     helloOk: true,
     client: clientMetadata,
     compression: compressors

src/cmap/connection.ts

@@ -582,6 +582,9 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
         this.throwIfAborted();
       }
     } catch (error) {
+      if (options.session != null && !(error instanceof MongoServerError)) {
+        updateSessionFromResponse(options.session, MongoDBResponse.empty);
+      }
       if (this.shouldEmitAndLogCommand) {
         this.emitAndLogCommand(
           this.monitorCommands,

src/index.ts

@@ -87,6 +87,7 @@ export {
   MongoWriteConcernError,
   WriteConcernErrorResult
 } from './error';
+export { TokenBucket } from './token_bucket';
 export {
   AbstractCursor,
   // Actual driver classes exported
@@ -539,7 +540,12 @@ export type {
 export type { InsertManyResult, InsertOneOptions, InsertOneResult } from './operations/insert';
 export type { CollectionInfo, ListCollectionsOptions } from './operations/list_collections';
 export type { ListDatabasesOptions, ListDatabasesResult } from './operations/list_databases';
-export type { AbstractOperation, Hint, OperationOptions } from './operations/operation';
+export type {
+  AbstractOperation,
+  Hint,
+  OperationOptions,
+  RetryContext
+} from './operations/operation';
 export type { ProfilingLevelOptions } from './operations/profiling_level';
 export type { RemoveUserOptions } from './operations/remove_user';
 export type { RenameOptions } from './operations/rename';

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,
@@ -10,6 +12,7 @@ import {
   MongoInvalidArgumentError,
   MongoNetworkError,
   MongoNotConnectedError,
+  MongoOperationTimeoutError,
   MongoRuntimeError,
   MongoServerError,
   MongoTransactionError,
@@ -26,9 +29,16 @@ import {
 import type { Topology } from '../sdam/topology';
 import type { ClientSession } from '../sessions';
 import { TimeoutContext } from '../timeout';
-import { abortable, maxWireVersion, supportsRetryableWrites } from '../utils';
+import { RETRY_COST, TOKEN_REFRESH_RATE } from '../token_bucket';
+import {
+  abortable,
+  ExponentialBackoffProvider,
+  maxWireVersion,
+  supportsRetryableWrites
+} from '../utils';
 import { AggregateOperation } from './aggregate';
-import { AbstractOperation, Aspect } from './operation';
+import { AbstractOperation, Aspect, RetryContext } from './operation';
+import { RunCommandOperation } from './run_command';
 
 const MMAPv1_RETRY_WRITES_ERROR_CODE = MONGODB_ERROR_CODES.IllegalOperation;
 const MMAPv1_RETRY_WRITES_ERROR_MESSAGE =
@@ -50,7 +60,7 @@ type ResultTypeFromOperation<TOperation extends AbstractOperation> = ReturnType<
  * The expectation is that this function:
  * - Connects the MongoClient if it has not already been connected, see {@link autoConnect}
  * - Creates a session if none is provided and cleans up the session it creates
- * - Tries an operation and retries under certain conditions, see {@link tryOperation}
+ * - Tries an operation and retries under certain conditions, see {@link executeOperationWithRetries}
  *
  * @typeParam T - The operation's type
  * @typeParam TResult - The type of the operation's result, calculated from T
@@ -120,7 +130,7 @@ export async function executeOperation<
   });
 
   try {
-    return await tryOperation(operation, {
+    return await executeOperationWithRetries(operation, {
       topology,
       timeoutContext,
       session,
@@ -184,7 +194,10 @@ type RetryOptions = {
  *
  * @param operation - The operation to execute
  * */
-async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFromOperation<T>>(
+async function executeOperationWithRetries<
+  T extends AbstractOperation,
+  TResult = ResultTypeFromOperation<T>
+>(
   operation: T,
   { topology, timeoutContext, session, readPreference }: RetryOptions
 ): Promise<TResult> {
@@ -233,11 +246,21 @@ async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFro
     session.incrementTransactionNumber();
   }
 
-  const maxTries = willRetry ? (timeoutContext.csotEnabled() ? Infinity : 2) : 1;
   let previousOperationError: MongoError | undefined;
   const deprioritizedServers = new DeprioritizedServers();
 
-  for (let tries = 0; tries < maxTries; tries++) {
+  const backoffDelayProvider = new ExponentialBackoffProvider(
+    10_000, // MAX_BACKOFF
+    100, // base backoff
+    2 // backoff rate
+  );
+
+  const retryContext = new RetryContext(
+    willRetry,
+    willRetry ? (timeoutContext.csotEnabled() ? Infinity : 2) : 1
+  );
+
+  for (; retryContext.shouldRetry(); retryContext.recordFailure(previousOperationError)) {
     if (previousOperationError) {
       if (hasWriteAspect && previousOperationError.code === MMAPv1_RETRY_WRITES_ERROR_CODE) {
         throw new MongoServerError({
@@ -247,15 +270,46 @@ async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFro
         });
       }
 
-      if (operation.hasAspect(Aspect.COMMAND_BATCHING) && !operation.canRetryWrite) {
+      const canRetryBackpressureError =
+        (operation.hasAspect(Aspect.WRITE_OPERATION) && topology.s.options.retryWrites) ||
+        (operation.hasAspect(Aspect.READ_OPERATION) && topology.s.options.retryReads) ||
+        (operation instanceof RunCommandOperation && topology.s.options.retryReads);
+      // TODO: think about whether or not willRetry checks are necessary here.
+      const isRetryable =
+        // bulk write commands are retryable if all operations in the batch are retryable
+        (willRetryWrite &&
+          operation.hasAspect(Aspect.COMMAND_BATCHING) &&
+          operation.canRetryWrite) ||
+        // if we have a retryable read or write operation, we can retry
+        (hasWriteAspect && willRetryWrite && isRetryableWriteError(previousOperationError)) ||
+        (hasReadAspect && willRetryRead && isRetryableReadError(previousOperationError)) ||
+        // if we have a retryable, system overloaded error, we can retry
+        (canRetryBackpressureError &&
+          previousOperationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError) &&
+          previousOperationError.hasErrorLabel(MongoErrorLabel.RetryableError));
+
+      if (!isRetryable) {
         throw previousOperationError;
       }
 
-      if (hasWriteAspect && !isRetryableWriteError(previousOperationError))
-        throw previousOperationError;
+      if (previousOperationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)) {
+        const delayMS = backoffDelayProvider.getNextBackoffDuration();
 
-      if (hasReadAspect && !isRetryableReadError(previousOperationError)) {
-        throw previousOperationError;
+        // if the delay would exhaust the CSOT timeout, short-circuit.
+        if (timeoutContext.csotEnabled() && delayMS > timeoutContext.remainingTimeMS) {
+          throw new MongoOperationTimeoutError(
+            `MongoDB SystemOverload exponential backoff would exceed timeoutMS deadline: remaining CSOT deadline=${timeoutContext.remainingTimeMS}, backoff delayMS=${delayMS}`,
+            {
+              cause: previousOperationError
+            }
+          );
+        }
+
+        if (!topology.tokenBucket.consume(RETRY_COST)) {
+          throw previousOperationError;
+        }
+
+        await setTimeout(delayMS);
       }
 
       if (
@@ -285,19 +339,32 @@ 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 (retryContext.isRetry && operation.hasAspect(Aspect.COMMAND_BATCHING)) {
         operation.resetBatch();
       }
 
       try {
         const result = await server.command(operation, timeoutContext);
+        topology.tokenBucket.deposit(
+          retryContext.isRetry
+            ? // on successful retry, deposit the retry cost + the refresh rate.
+              TOKEN_REFRESH_RATE + RETRY_COST
+            : // otherwise, just deposit the refresh rate.
+              TOKEN_REFRESH_RATE
+        );
         return operation.handleOk(result);
       } catch (error) {
         return operation.handleError(error);
       }
     } catch (operationError) {
       if (!(operationError instanceof MongoError)) throw operationError;
+
+      if (!operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)) {
+        // if an operation fails with an error that does not contain the SystemOverloadError, deposit 1 token.
+        topology.tokenBucket.deposit(RETRY_COST);
+      }
+
       if (
         previousOperationError != null &&
         operationError.hasErrorLabel(MongoErrorLabel.NoWritesPerformed)
@@ -312,8 +379,5 @@ async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFro
     }
   }
 
-  throw (
-    previousOperationError ??
-    new MongoRuntimeError('Tried to propagate retryability error, but no error was found.')
-  );
+  throw previousOperationError ?? new MongoRuntimeError('ahh');
 }

src/operations/operation.ts

@@ -1,6 +1,7 @@
 import { type Connection, type MongoError } from '..';
 import { type BSONSerializeOptions, type Document, resolveBSONOptions } from '../bson';
 import { type MongoDBResponse } from '../cmap/wire_protocol/responses';
+import { MongoErrorLabel } from '../error';
 import { type Abortable } from '../mongo_types';
 import { ReadPreference, type ReadPreferenceLike } from '../read_preference';
 import type { Server, ServerCommandOptions } from '../sdam/server';
@@ -45,6 +46,42 @@ export interface OperationOptions extends BSONSerializeOptions {
   timeoutMS?: number;
 }
 
+/**
+ * @internal
+ */
+export class RetryContext {
+  private maxAttempts: number;
+  private willRetry: boolean;
+  private attempts: number = 0;
+
+  constructor(willRetry: boolean, maxAttempts: number) {
+    this.maxAttempts = maxAttempts;
+    this.willRetry = willRetry;
+  }
+
+  get isRetry() {
+    return this.attempts > 0;
+  }
+
+  shouldRetry() {
+    return this.attempts < this.maxAttempts;
+  }
+
+  recordFailure(error: MongoError) {
+    this.attempts++;
+
+    const isRetryableOverloadError =
+      error.hasErrorLabel(MongoErrorLabel.RetryableError) &&
+      error.hasErrorLabel(MongoErrorLabel.SystemOverloadedError);
+
+    if (!(this.willRetry || isRetryableOverloadError)) return;
+
+    this.maxAttempts = error.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)
+      ? 6
+      : this.maxAttempts;
+  }
+}
+
 /**
  * This class acts as a parent class for any operation and is responsible for setting this.options,
  * as well as setting and getting a session.
@@ -66,6 +103,8 @@ export abstract class AbstractOperation<TResult = any> {
   /** Specifies the time an operation will run until it throws a timeout error. */
   timeoutMS?: number;
 
+  retryContext?: RetryContext;
+
   private _session: ClientSession | undefined;
 
   static aspects?: Set<symbol>;

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,
@@ -207,18 +208,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 */