deposit retry tokens only on retries, update maxAttempts property conditionally

Author: tadjik1

package-lock.json

@@ -586,12 +586,6 @@ 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);
       }

src/cmap/connection.ts

@@ -12,7 +12,6 @@ import {
   MongoInvalidArgumentError,
   MongoNetworkError,
   MongoNotConnectedError,
-  MongoOperationTimeoutError,
   MongoRuntimeError,
   MongoServerError,
   MongoTransactionError,
@@ -29,13 +28,8 @@ 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,
-  ExponentialBackoffProvider,
-  maxWireVersion,
-  supportsRetryableWrites
-} from '../utils';
+import { RETRY_COST, RETRY_TOKEN_RETURN_RATE } from '../token_bucket';
+import { abortable, maxWireVersion, supportsRetryableWrites } from '../utils';
 import { AggregateOperation } from './aggregate';
 import { AbstractOperation, Aspect } from './operation';
 import { RunCommandOperation } from './run_command';
@@ -194,7 +188,6 @@ type RetryOptions = {
  *
  * @param operation - The operation to execute
  */
-// 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>
@@ -248,11 +241,6 @@ async function executeOperationWithRetries<
   }
 
   const deprioritizedServers = new DeprioritizedServers();
-  const backoffDelayProvider = new ExponentialBackoffProvider(
-    10_000, // MAX_BACKOFF
-    100, // base backoff
-    2 // backoff rate
-  );
 
   let maxAttempts =
     typeof operation.maxAttempts === 'number'
@@ -271,6 +259,7 @@ async function executeOperationWithRetries<
   let error: MongoError | null = null;
 
   for (let attempt = 0; attempt < maxAttempts; attempt++) {
+    operation.attemptsMade = attempt + 1;
     operation.server = server;
 
     try {
@@ -279,28 +268,20 @@ async function executeOperationWithRetries<
         topology.tokenBucket.deposit(
           attempt > 0
             ? // on successful retry, deposit the retry cost + the refresh rate.
-              TOKEN_REFRESH_RATE + RETRY_COST
+              RETRY_TOKEN_RETURN_RATE + RETRY_COST
             : // otherwise, just deposit the refresh rate.
-              TOKEN_REFRESH_RATE
+              RETRY_TOKEN_RETURN_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.
+      if (attempt > 0 && !operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)) {
+        // if a retry attempt fails with a non-overload error, deposit 1 token.
         topology.tokenBucket.deposit(RETRY_COST);
       }
 
@@ -327,10 +308,9 @@ async function executeOperationWithRetries<
         throw error;
       }
 
-      maxAttempts =
-        shouldRetry && operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)
-          ? 6
-          : maxAttempts;
+      if (operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)) {
+        maxAttempts = Math.min(6, operation.maxAttempts ?? 6);
+      }
 
       if (attempt + 1 >= maxAttempts) {
         throw error;
@@ -341,17 +321,11 @@ async function executeOperationWithRetries<
           throw error;
         }
 
-        const delayMS = backoffDelayProvider.getNextBackoffDuration();
+        const delayMS = Math.random() * Math.min(10_000, 100 * 2 ** attempt);
 
         // 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
-            }
-          );
+          throw error;
         }
 
         await setTimeout(delayMS);

src/operations/execute_operation.ts

@@ -66,14 +66,12 @@ 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.
+  /** @internal Used by commitTransaction to share the retry budget across two executeOperatin calls. */
   maxAttempts?: number;
 
+  /** @internal Tracks how many attempts were made in the last executeOperation call. */
+  attemptsMade?: number;
+
   private _session: ClientSession | undefined;
 
   static aspects?: Set<symbol>;

src/operations/operation.ts

@@ -213,7 +213,6 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
   hello?: Document;
   _type?: string;
 
-  // Note for Sergey: expect this to change, and instead be stored on the server class.
   tokenBucket = new TokenBucket(1000);
 
   client!: MongoClient;

src/sdam/topology.ts

@@ -497,7 +497,7 @@ export class ClientSession
       readPreference: ReadPreference.primary,
       bypassPinningCheck: true
     });
-    operation.maxAttempts = 5;
+    operation.maxAttempts = 6;
 
     const timeoutContext =
       this.timeoutContext ??
@@ -515,8 +515,12 @@ export class ClientSession
       return;
     } catch (firstCommitError) {
       this.commitAttempted = true;
-      // Backpressure exhausted all retries on the initial retry loop.
-      if (operation.maxAttempts === 5) throw firstCommitError;
+
+      const remainingAttempts = 6 - (operation.attemptsMade ?? 1);
+      if (remainingAttempts <= 0) {
+        throw firstCommitError;
+      }
+
       if (firstCommitError instanceof MongoError && isRetryableWriteError(firstCommitError)) {
         // SPEC-1185: apply majority write concern when retrying commitTransaction
         WriteConcern.apply(command, { wtimeoutMS: 10000, ...wc, w: 'majority' });
@@ -529,7 +533,7 @@ export class ClientSession
             readPreference: ReadPreference.primary,
             bypassPinningCheck: true
           });
-          op.maxAttempts = operation.maxAttempts;
+          op.maxAttempts = remainingAttempts;
           await executeOperation(this.client, op, timeoutContext);
           return;
         } catch (retryCommitError) {

src/sessions.ts

@@ -26,7 +26,7 @@ export class TokenBucket {
  * @internal
  * The amount to deposit on successful operations, as defined in the backpressure specification.
  */
-export const TOKEN_REFRESH_RATE = 0.1;
+export const RETRY_TOKEN_RETURN_RATE = 0.1;
 /**
  * @internal
  * The initial size of the token bucket, as defined in the backpressure specification.

src/token_bucket.ts

@@ -1419,23 +1419,3 @@ export async function abortable<T>(
     abortListener?.[kDispose]();
   }
 }
-
-// Note for Sergey: overkill, I know - but sometimes you gotta over-engineer a bit for fun and revert before opening for review ¯\_(ツ)_/¯
-// feel free to remove this class. Or leave it. If you keep it, we should use this for the convenient transactions' API as well.
-export class ExponentialBackoffProvider {
-  constructor(
-    public readonly maxBackoff: number,
-    public readonly baseBackoff: number,
-    public readonly backoffIncreaseRate: number,
-    public iteration = 0
-  ) {}
-
-  getNextBackoffDuration(): number {
-    const iteration = this.iteration;
-    this.iteration += 1;
-    return (
-      Math.random() *
-      Math.min(this.maxBackoff, this.baseBackoff * this.backoffIncreaseRate ** iteration)
-    );
-  }
-}