POC

Author: baileympearson

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

@@ -586,6 +586,15 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
         this.throwIfAborted();
       }
     } catch (error) {
+      // Note for Sergey: when retrying a command with `startTransaction`, the spec says drivers must only mark the transaction
+      // as in progress _after_ the server responds, regardless of the command's outcome.
+      // Server errors are thrown from the try-block above _after_ updateSessionFromResponse is called. So, server errors already correctly update the
+      // session's state.
+      // Network errors, however, are thrown from the try-block from `sendWire()`, and as such bypass the call to `updateSessionFromResponse` above. So, we have to
+      // handle updating sessions for non-server errors here.
+      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

src/operations/execute_operation.ts

@@ -1,3 +1,5 @@
+import { setTimeout } from '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 { 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,
@@ -183,8 +193,12 @@ type RetryOptions = {
  * @typeParam TResult - The type of the operation's result, calculated from T
  *
  * @param operation - The operation to execute
- * */
-async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFromOperation<T>>(
+ */
+// Note for Sergey: The rename here could be reverted - I just thought this was more descriptive.
+async function executeOperationWithRetries<
+  T extends AbstractOperation,
+  TResult = ResultTypeFromOperation<T>
+>(
   operation: T,
   { topology, timeoutContext, session, readPreference }: RetryOptions
 ): Promise<TResult> {
@@ -233,33 +247,118 @@ 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();
+  const backoffDelayProvider = new ExponentialBackoffProvider(
+    10_000, // MAX_BACKOFF
+    100, // base backoff
+    2 // backoff rate
+  );
+
+  let maxAttempts =
+    typeof operation.maxAttempts === 'number'
+      ? operation.maxAttempts
+      : willRetry
+        ? timeoutContext.csotEnabled()
+          ? Infinity
+          : 2
+        : 1;
+
+  const shouldRetry =
+    (operation.hasAspect(Aspect.READ_OPERATION) && topology.s.options.retryReads) ||
+    ((operation.hasAspect(Aspect.WRITE_OPERATION) || operation instanceof RunCommandOperation) &&
+      topology.s.options.retryWrites);
+
+  let error: MongoError | null = null;
 
-  for (let tries = 0; tries < maxTries; tries++) {
-    if (previousOperationError) {
-      if (hasWriteAspect && previousOperationError.code === MMAPv1_RETRY_WRITES_ERROR_CODE) {
+  for (let attempt = 0; attempt < maxAttempts; attempt++) {
+    operation.server = server;
+
+    try {
+      try {
+        const result = await server.command(operation, timeoutContext);
+        topology.tokenBucket.deposit(
+          attempt > 0
+            ? // 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);
+      }
+
+      // Note for Sergey: This ended up being a larger refactor than I anticipated. But it made more sense to me to put the error handling
+      // that previously lived in the try-block into the catch-block.
+      // The primary motivator for this change was to make error tracking simpler. The post failure but pre-retry logic needs access to both the
+      // most recently encountered operation error _and_ the error we're storing to potentially throw, if we need to (because in some circumstances
+      // the error raised from here is _not_ the most recently encountered operation error). If we chose to keep the previously existing structure,
+      // this would force us to keep track of two errors outside of the for loop. By moving the pre-retry logic into the catch block, it gains
+      // access to `operationError`, and still only requires keeping track of one error outside the loop.
+    } catch (operationError) {
+      // Should never happen but if it does - propagate the error.
+      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 (error == null) {
+        error = operationError;
+      } else {
+        if (!operationError.hasErrorLabel(MongoErrorLabel.NoWritesPerformed)) {
+          error = operationError;
+        }
+      }
+
+      // Reset timeouts
+      timeoutContext.clear();
+
+      if (hasWriteAspect && operationError.code === MMAPv1_RETRY_WRITES_ERROR_CODE) {
         throw new MongoServerError({
           message: MMAPv1_RETRY_WRITES_ERROR_MESSAGE,
           errmsg: MMAPv1_RETRY_WRITES_ERROR_MESSAGE,
-          originalError: previousOperationError
+          originalError: operationError
         });
       }
 
-      if (operation.hasAspect(Aspect.COMMAND_BATCHING) && !operation.canRetryWrite) {
-        throw previousOperationError;
+      if (!canRetry(operation, operationError)) {
+        throw error;
       }
 
-      if (hasWriteAspect && !isRetryableWriteError(previousOperationError))
-        throw previousOperationError;
+      maxAttempts =
+        shouldRetry && operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)
+          ? 6
+          : maxAttempts;
+
+      if (attempt + 1 >= maxAttempts) {
+        throw error;
+      }
 
-      if (hasReadAspect && !isRetryableReadError(previousOperationError)) {
-        throw previousOperationError;
+      if (operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)) {
+        if (!topology.tokenBucket.consume(RETRY_COST)) {
+          throw error;
+        }
+
+        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: error
+            }
+          );
+        }
+
+        await setTimeout(delayMS);
       }
 
       if (
-        previousOperationError instanceof MongoNetworkError &&
+        operationError instanceof MongoNetworkError &&
         operation.hasAspect(Aspect.CURSOR_CREATING) &&
         session != null &&
         session.isPinned &&
@@ -268,52 +367,62 @@ async function tryOperation<T extends AbstractOperation, TResult = ResultTypeFro
         session.unpin({ force: true, forceClear: true });
       }
 
+      deprioritizedServers.add(server.description);
+
       server = await topology.selectServer(selector, {
         session,
         operationName: operation.commandName,
         deprioritizedServers,
         signal: operation.options.signal
       });
 
-      if (hasWriteAspect && !supportsRetryableWrites(server)) {
+      if (
+        hasWriteAspect &&
+        !supportsRetryableWrites(server) &&
+        !operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)
+      ) {
         throw new MongoUnexpectedServerResponseError(
           'Selected server does not support retryable writes'
         );
       }
-    }
 
-    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)) {
+      // Batched operations must reset the batch before retry,
+      // otherwise building a command will build the _next_ batch, not the current batch.
+      if (operation.hasAspect(Aspect.COMMAND_BATCHING)) {
         operation.resetBatch();
       }
-
-      try {
-        const result = await server.command(operation, timeoutContext);
-        return operation.handleOk(result);
-      } catch (error) {
-        return operation.handleError(error);
-      }
-    } catch (operationError) {
-      if (!(operationError instanceof MongoError)) throw operationError;
-      if (
-        previousOperationError != null &&
-        operationError.hasErrorLabel(MongoErrorLabel.NoWritesPerformed)
-      ) {
-        throw previousOperationError;
-      }
-      deprioritizedServers.add(server.description);
-      previousOperationError = operationError;
-
-      // Reset timeouts
-      timeoutContext.clear();
     }
   }
 
   throw (
-    previousOperationError ??
-    new MongoRuntimeError('Tried to propagate retryability error, but no error was found.')
+    error ??
+    new MongoRuntimeError(
+      'Should never happen: operation execution loop terminated but no error was recorded.'
+    )
   );
+
+  function canRetry(operation: AbstractOperation, error: MongoError) {
+    // always retryable
+    if (
+      error.hasErrorLabel(MongoErrorLabel.SystemOverloadedError) &&
+      error.hasErrorLabel(MongoErrorLabel.RetryableError)
+    ) {
+      return true;
+    }
+
+    // run command is only retryable if we get retryable overload errors
+    if (operation instanceof RunCommandOperation) {
+      return false;
+    }
+
+    // batch operations are only retryable if the batch is retryable
+    if (operation.hasAspect(Aspect.COMMAND_BATCHING)) {
+      return operation.canRetryWrite;
+    }
+
+    return (
+      (hasWriteAspect && willRetryWrite && isRetryableWriteError(error)) ||
+      (hasReadAspect && willRetryRead && isRetryableReadError(error))
+    );
+  }
 }

src/operations/operation.ts

@@ -66,6 +66,14 @@ export abstract class AbstractOperation<TResult = any> {
   /** Specifies the time an operation will run until it throws a timeout error. */
   timeoutMS?: number;
 
+  // Note for Sergey: sort of a hack, one that I planned to revisit before officially opening the PR for review.
+  // For every operation _except_ transactions, all backpressure retries occur inside one `executeOperation` call. `commitTransaction`, however,
+  // currently hard-codes two retry attempts, each with a `RunCommandOperation`. Neither retry attempt is retryable by itself.
+  // Because we use two separate RunCommand operations, it's possible to retry up to 2 * MAX_BACKPRESSURE_RETRIES times, if the server is overloaded.
+  // I hacked around it here by storing the max attempts for an operation on the operation, and sharing the value between both operations.
+  // I don't have a solid idea of a better way to do this right now though.
+  maxAttempts?: number;
+
   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,16 @@ export type TopologyEvents = {
  * @internal
  */
 export class Topology extends TypedEventEmitter<TopologyEvents> {
-  /** @internal */
   s: TopologyPrivate;
-  /** @internal */
   waitQueue: List<ServerSelectionRequest>;
-  /** @internal */
   hello?: Document;
-  /** @internal */
   _type?: string;
 
+  // Note for Sergey: expect this to change, and instead be stored on the server class.
+  tokenBucket = new TokenBucket(1000);
+
   client!: MongoClient;
 
-  /** @internal */
   private connectionLock?: Promise<Topology>;
 
   /** @event */