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

src/change_stream.ts

@@ -18,6 +18,7 @@ import { MongoClient } from './mongo_client';
 import { type InferIdType, TypedEventEmitter } from './mongo_types';
 import type { AggregateOptions } from './operations/aggregate';
 import type { OperationParent } from './operations/command';
+import { DeprioritizedServers } from './sdam/server_selection';
 import type { ServerSessionId } from './sessions';
 import { CSOTTimeoutContext, type TimeoutContext } from './timeout';
 import { type AnyOptions, getTopology, type MongoDBNamespace, squashError } from './utils';
@@ -1073,7 +1074,8 @@ export class ChangeStream<
     try {
       await topology.selectServer(this.cursor.readPreference, {
         operationName: 'reconnect topology in change stream',
-        timeoutContext: this.timeoutContext
+        timeoutContext: this.timeoutContext,
+        deprioritizedServers: new DeprioritizedServers()
       });
       this.cursor = this._createChangeStreamCursor(this.cursor.resumeOptions);
     } catch {

src/cmap/connect.ts

@@ -209,6 +209,7 @@ export interface HandshakeDocument extends Document {
   compression: string[];
   saslSupportedMechs?: string;
   loadBalanced?: boolean;
+  backpressure: true;
 }
 
 /**
@@ -226,6 +227,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/error.ts

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

src/index.ts

@@ -87,6 +87,7 @@ export {
   MongoWriteConcernError,
   WriteConcernErrorResult
 } from './error';
+export { TokenBucket } from './token_bucket';
 export {
   AbstractCursor,
   // Actual driver classes exported
@@ -588,7 +589,7 @@ export type {
   TagSet,
   TopologyVersion
 } from './sdam/server_description';
-export type { ServerSelector } from './sdam/server_selection';
+export type { DeprioritizedServers, ServerSelector } from './sdam/server_selection';
 export type { SrvPoller, SrvPollerEvents, SrvPollerOptions } from './sdam/srv_polling';
 export type {
   ConnectOptions,

src/mongo_client.ts

@@ -48,7 +48,7 @@ import type { ReadConcern, ReadConcernLevel, ReadConcernLike } from './read_conc
 import { ReadPreference, type ReadPreferenceMode } from './read_preference';
 import type { ServerMonitoringMode } from './sdam/monitor';
 import type { TagSet } from './sdam/server_description';
-import { readPreferenceServerSelector } from './sdam/server_selection';
+import { DeprioritizedServers, readPreferenceServerSelector } from './sdam/server_selection';
 import type { SrvPoller } from './sdam/srv_polling';
 import { Topology, type TopologyEvents } from './sdam/topology';
 import { ClientSession, type ClientSessionOptions, ServerSessionPool } from './sessions';
@@ -789,7 +789,7 @@ export class MongoClient extends TypedEventEmitter<MongoClientEvents> implements
       // to avoid the server selection timeout.
       const selector = readPreferenceServerSelector(ReadPreference.primaryPreferred);
       const serverDescriptions = Array.from(topologyDescription.servers.values());
-      const servers = selector(topologyDescription, serverDescriptions);
+      const servers = selector(topologyDescription, serverDescriptions, new DeprioritizedServers());
       if (servers.length !== 0) {
         const endSessions = Array.from(client.s.sessionPool.sessions, ({ id }) => id);
         if (endSessions.length !== 0) {

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,23 +12,30 @@ import {
   MongoInvalidArgumentError,
   MongoNetworkError,
   MongoNotConnectedError,
+  MongoOperationTimeoutError,
   MongoRuntimeError,
   MongoServerError,
   MongoTransactionError,
   MongoUnexpectedServerResponseError
 } from '../error';
 import type { MongoClient } from '../mongo_client';
 import { ReadPreference } from '../read_preference';
-import type { ServerDescription } from '../sdam/server_description';
 import {
+  DeprioritizedServers,
   sameServerSelector,
   secondaryWritableServerSelector,
   type ServerSelector
 } from '../sdam/server_selection';
 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';
 
@@ -50,7 +59,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 +129,7 @@ export async function executeOperation<
   });
 
   try {
-    return await tryOperation(operation, {
+    return await executeOperationWithRetries(operation, {
       topology,
       timeoutContext,
       session,
@@ -184,7 +193,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> {
@@ -207,7 +219,8 @@ async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFro
     session,
     operationName: operation.commandName,
     timeoutContext,
-    signal: operation.options.signal
+    signal: operation.options.signal,
+    deprioritizedServers: new DeprioritizedServers()
   });
 
   const hasReadAspect = operation.hasAspect(Aspect.READ_OPERATION);
@@ -232,11 +245,23 @@ async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFro
     session.incrementTransactionNumber();
   }
 
-  const maxTries = willRetry ? (timeoutContext.csotEnabled() ? Infinity : 2) : 1;
   let previousOperationError: MongoError | undefined;
-  let previousServer: ServerDescription | 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
+  );
+
+  for (
+    let attempt = 0, maxAttempts = willRetry ? (timeoutContext.csotEnabled() ? Infinity : 2) : 1;
+    attempt <= maxAttempts;
+    attempt++,
+      maxAttempts = previousOperationError?.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)
+        ? 5
+        : maxAttempts
+  ) {
     if (previousOperationError) {
       if (hasWriteAspect && previousOperationError.code === MMAPv1_RETRY_WRITES_ERROR_CODE) {
         throw new MongoServerError({
@@ -246,15 +271,39 @@ async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFro
         });
       }
 
-      if (operation.hasAspect(Aspect.COMMAND_BATCHING) && !operation.canRetryWrite) {
+      const isRetryable =
+        // bulk write commands are retryable if all operations in the batch are retryable
+        (operation.hasAspect(Aspect.COMMAND_BATCHING) && operation.canRetryWrite) ||
+        // if we have a retryable read or write operation, we can retry
+        (hasWriteAspect && isRetryableWriteError(previousOperationError)) ||
+        (hasReadAspect && isRetryableReadError(previousOperationError)) ||
+        // if we have a retryable, system overloaded error, we can retry
+        (previousOperationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError) &&
+          previousOperationError.hasErrorLabel(MongoErrorLabel.RetryableError));
+
+      if (!isRetryable) {
         throw previousOperationError;
       }
 
-      if (hasWriteAspect && !isRetryableWriteError(previousOperationError))
-        throw previousOperationError;
-
-      if (hasReadAspect && !isRetryableReadError(previousOperationError)) {
-        throw previousOperationError;
+      if (previousOperationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)) {
+        const delayMS = backoffDelayProvider.getNextBackoffDuration();
+
+        // if the delay would exhaust the CSOT timeout, short-circuit.
+        if (timeoutContext.csotEnabled() && delayMS > timeoutContext.remainingTimeMS) {
+          // TODO: is this the right error to throw?
+          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 (
@@ -270,7 +319,7 @@ async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFro
       server = await topology.selectServer(selector, {
         session,
         operationName: operation.commandName,
-        previousServer,
+        deprioritizedServers,
         signal: operation.options.signal
       });
 
@@ -284,35 +333,47 @@ 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)) {
+      const isRetry = attempt > 0;
+
+      // If attempt > 0 and we are command batching we need to reset the batch.
+      if (isRetry && operation.hasAspect(Aspect.COMMAND_BATCHING)) {
         operation.resetBatch();
       }
 
       try {
         const result = await server.command(operation, timeoutContext);
+        topology.tokenBucket.deposit(
+          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)
       ) {
         throw previousOperationError;
       }
-      previousServer = server.description;
+      deprioritizedServers.add(server.description);
       previousOperationError = operationError;
 
       // Reset timeouts
       timeoutContext.clear();
     }
   }
 
-  throw (
-    previousOperationError ??
-    new MongoRuntimeError('Tried to propagate retryability error, but no error was found.')
-  );
+  throw previousOperationError ?? new MongoRuntimeError('ahh');
 }