add retry context

Author: baileympearson

src/operations/execute_operation.ts

@@ -37,7 +37,7 @@ import {
   supportsRetryableWrites
 } from '../utils';
 import { AggregateOperation } from './aggregate';
-import { AbstractOperation, Aspect } from './operation';
+import { AbstractOperation, Aspect, RetryContext } from './operation';
 
 const MMAPv1_RETRY_WRITES_ERROR_CODE = MONGODB_ERROR_CODES.IllegalOperation;
 const MMAPv1_RETRY_WRITES_ERROR_MESSAGE =
@@ -254,18 +254,9 @@ async function executeOperationWithRetries<
     2 // backoff rate
   );
 
-  let maxAttempts =
-    (operation.maxAttempts ?? willRetry) ? (timeoutContext.csotEnabled() ? Infinity : 2) : 1;
-
-  for (
-    let attempt = 0;
-    attempt < maxAttempts;
-    attempt++,
-      maxAttempts =
-        willRetry && previousOperationError?.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)
-          ? 6
-          : maxAttempts
-  ) {
+  const retryContext = new RetryContext(willRetry, timeoutContext.csotEnabled() ? Infinity : 2);
+
+  for (; retryContext.shouldRetry(); retryContext.recordFailure(previousOperationError)) {
     if (previousOperationError) {
       if (hasWriteAspect && previousOperationError.code === MMAPv1_RETRY_WRITES_ERROR_CODE) {
         throw new MongoServerError({
@@ -294,7 +285,6 @@ async function executeOperationWithRetries<
 
         // 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}`,
             {
@@ -337,17 +327,15 @@ async function executeOperationWithRetries<
     operation.server = server;
 
     try {
-      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)) {
+      if (retryContext.isRetry && operation.hasAspect(Aspect.COMMAND_BATCHING)) {
         operation.resetBatch();
       }
 
       try {
         const result = await server.command(operation, timeoutContext);
         topology.tokenBucket.deposit(
-          isRetry
+          retryContext.isRetry
             ? // on successful retry, deposit the retry cost + the refresh rate.
               TOKEN_REFRESH_RATE + RETRY_COST
             : // otherwise, just deposit the refresh rate.

src/operations/operation.ts

@@ -1,4 +1,4 @@
-import { type Connection, type MongoError } from '..';
+import { type Connection, type MongoError, MongoErrorLabel } from '..';
 import { type BSONSerializeOptions, type Document, resolveBSONOptions } from '../bson';
 import { type MongoDBResponse } from '../cmap/wire_protocol/responses';
 import { type Abortable } from '../mongo_types';
@@ -45,6 +45,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,7 +102,7 @@ export abstract class AbstractOperation<TResult = any> {
   /** Specifies the time an operation will run until it throws a timeout error. */
   timeoutMS?: number;
 
-  maxAttempts?: number;
+  retryContext?: RetryContext;
 
   private _session: ClientSession | undefined;
 

src/sessions.ts

@@ -25,7 +25,7 @@ import {
 import type { MongoClient, MongoOptions } from './mongo_client';
 import { TypedEventEmitter } from './mongo_types';
 import { executeOperation } from './operations/execute_operation';
-import { RetryAttemptContext } from './operations/operation';
+import { RetryContext } from './operations/operation';
 import { RunCommandOperation } from './operations/run_command';
 import { ReadConcernLevel } from './read_concern';
 import { ReadPreference } from './read_preference';
@@ -494,14 +494,14 @@ export class ClientSession
       command.recoveryToken = this.transaction.recoveryToken;
     }
 
-    const retryContext = new RetryAttemptContext(5);
+    const retryContext = new RetryContext(false, 5);
 
     const operation = new RunCommandOperation(new MongoDBNamespace('admin'), command, {
       session: this,
       readPreference: ReadPreference.primary,
       bypassPinningCheck: true
     });
-    operation.attempts = retryContext;
+    operation.retryContext = retryContext;
 
     const timeoutContext =
       this.timeoutContext ??
@@ -531,7 +531,7 @@ export class ClientSession
             readPreference: ReadPreference.primary,
             bypassPinningCheck: true
           });
-          op.attempts = retryContext;
+          op.retryContext = operation.retryContext;
           await executeOperation(this.client, op, timeoutContext);
           return;
         } catch (retryCommitError) {