feat(NODE-7491): finalize client backpressure implementation for phase 1 rollout (#4920)

Author: tadjik1

src/connection_string.ts

@@ -643,7 +643,7 @@ interface OptionDescriptor {
 }
 
 export const OPTIONS = {
-  adaptiveRetries: {
+  enableOverloadRetargeting: {
     default: false,
     type: 'boolean'
   },
@@ -885,6 +885,10 @@ export const OPTIONS = {
     default: 15,
     type: 'uint'
   },
+  maxAdaptiveRetries: {
+    default: 2,
+    type: 'uint'
+  },
   maxConnecting: {
     default: 2,
     transform({ name, values: [value] }): number {

src/index.ts

@@ -621,7 +621,6 @@ export type {
   TimeoutContext,
   TimeoutContextOptions
 } from './timeout';
-export type { MAX_RETRIES, TokenBucket } from './token_bucket';
 export type { Transaction, TransactionOptions, TxnState } from './transactions';
 export type {
   BufferPool,

src/mongo_client.ts

@@ -229,8 +229,18 @@ export interface MongoClientOptions extends BSONSerializeOptions, SupportedNodeC
   retryReads?: boolean;
   /** Enable retryable writes. */
   retryWrites?: boolean;
-  /** Whether to enable adaptive retry rate limiting using a token bucket. Defaults to false. */
-  adaptiveRetries?: boolean;
+  /**
+   * The maximum number of retries during server overload. Set to 0 to disable overload retries. Defaults to 2.
+   * @see https://www.mongodb.com/docs/atlas/overload-errors
+   * */
+  maxAdaptiveRetries?: number;
+  /**
+   * Whether or not to enable overload retargeting. Defaults to false.
+   * @see https://www.mongodb.com/docs/atlas/overload-errors
+   * More information about the overload policy in drivers:
+   * @see https://github.com/mongodb/specifications/blob/master/source/client-backpressure/client-backpressure.md#overload-retry-policy
+   * */
+  enableOverloadRetargeting?: boolean;
   /** Allow a driver to force a Single topology type with a connection string containing one host */
   directConnection?: boolean;
   /** Instruct the driver it is connecting to a load balancer fronting a mongos like service */
@@ -1043,7 +1053,8 @@ export interface MongoOptions
   extends Required<
       Pick<
         MongoClientOptions,
-        | 'adaptiveRetries'
+        | 'maxAdaptiveRetries'
+        | 'enableOverloadRetargeting'
         | 'autoEncryption'
         | 'connectTimeoutMS'
         | 'directConnection'

src/operations/execute_operation.ts

@@ -29,13 +29,6 @@ import {
 import type { Topology } from '../sdam/topology';
 import type { ClientSession } from '../sessions';
 import { TimeoutContext } from '../timeout';
-import {
-  BASE_BACKOFF_MS,
-  MAX_BACKOFF_MS,
-  MAX_RETRIES,
-  RETRY_COST,
-  RETRY_TOKEN_RETURN_RATE
-} from '../token_bucket';
 import { abortable, maxWireVersion, supportsRetryableWrites } from '../utils';
 import { AggregateOperation } from './aggregate';
 import { AbstractOperation, Aspect } from './operation';
@@ -176,6 +169,10 @@ type RetryOptions = {
   topology: Topology;
   timeoutContext: TimeoutContext;
 };
+/** @internal The base backoff duration in milliseconds */
+const BASE_BACKOFF_MS = 100;
+/** @internal The maximum backoff duration in milliseconds */
+const MAX_BACKOFF_MS = 10_000;
 
 /**
  * Executes an operation and retries as appropriate
@@ -267,13 +264,6 @@ async function executeOperationWithRetries<
     try {
       try {
         const result = await server.command(operation, timeoutContext);
-        if (topology.s.options.adaptiveRetries) {
-          topology.tokenBucket.deposit(
-            attempt > 0
-              ? RETRY_TOKEN_RETURN_RATE + RETRY_COST // on successful retry
-              : RETRY_TOKEN_RETURN_RATE // otherwise
-          );
-        }
         return operation.handleOk(result);
       } catch (error) {
         return operation.handleError(error);
@@ -282,15 +272,6 @@ async function executeOperationWithRetries<
       // Should never happen but if it does - propagate the error.
       if (!(operationError instanceof MongoError)) throw operationError;
 
-      if (
-        topology.s.options.adaptiveRetries &&
-        attempt > 0 &&
-        !operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)
-      ) {
-        // if a retry attempt fails with a non-overload error, deposit 1 token.
-        topology.tokenBucket.deposit(RETRY_COST);
-      }
-
       // Preserve the original error once a write has been performed.
       // Only update to the latest error if no writes were performed.
       if (error == null) {
@@ -317,7 +298,8 @@ async function executeOperationWithRetries<
       }
 
       if (operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)) {
-        maxAttempts = Math.min(MAX_RETRIES + 1, operation.maxAttempts ?? MAX_RETRIES + 1);
+        const maxOverloadAttempts = topology.s.options.maxAdaptiveRetries + 1;
+        maxAttempts = Math.min(maxOverloadAttempts, operation.maxAttempts ?? maxOverloadAttempts);
       }
 
       if (attempt + 1 >= maxAttempts) {
@@ -352,16 +334,13 @@ async function executeOperationWithRetries<
           throw error;
         }
 
-        if (topology.s.options.adaptiveRetries && !topology.tokenBucket.consume(RETRY_COST)) {
-          throw error;
-        }
-
         await setTimeout(backoffMS);
       }
 
       if (
         topology.description.type === TopologyType.Sharded ||
-        operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError)
+        (operationError.hasErrorLabel(MongoErrorLabel.SystemOverloadedError) &&
+          topology.s.options.enableOverloadRetargeting)
       ) {
         deprioritizedServers.add(server.description);
       }

src/sdam/topology.ts

@@ -35,7 +35,6 @@ 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 { INITIAL_TOKEN_BUCKET_SIZE, TokenBucket } from '../token_bucket';
 import type { Transaction } from '../transactions';
 import {
   addAbortListener,
@@ -146,7 +145,8 @@ export interface TopologyOptions extends BSONSerializeOptions, ServerOptions {
   hosts: HostAddress[];
   retryWrites: boolean;
   retryReads: boolean;
-  adaptiveRetries: boolean;
+  maxAdaptiveRetries: number;
+  enableOverloadRetargeting: boolean;
   /** How long to block for server selection before throwing an error */
   serverSelectionTimeoutMS: number;
   /** The name of the replica set to connect to */
@@ -214,8 +214,6 @@ export class Topology extends TypedEventEmitter<TopologyEvents> {
   hello?: Document;
   _type?: string;
 
-  tokenBucket = new TokenBucket(INITIAL_TOKEN_BUCKET_SIZE);
-
   client!: MongoClient;
 
   private connectionLock?: Promise<Topology>;

src/sessions.ts

@@ -31,7 +31,6 @@ import { ReadConcernLevel } from './read_concern';
 import { ReadPreference } from './read_preference';
 import { _advanceClusterTime, type ClusterTime, TopologyType } from './sdam/common';
 import { TimeoutContext } from './timeout';
-import { MAX_RETRIES } from './token_bucket';
 import {
   isTransactionCommand,
   Transaction,
@@ -499,7 +498,7 @@ export class ClientSession
       readPreference: ReadPreference.primary,
       bypassPinningCheck: true
     });
-    operation.maxAttempts = MAX_RETRIES + 1;
+    operation.maxAttempts = this.clientOptions.maxAdaptiveRetries + 1;
 
     const timeoutContext =
       this.timeoutContext ??
@@ -518,7 +517,7 @@ export class ClientSession
     } catch (firstCommitError) {
       this.commitAttempted = true;
 
-      const remainingAttempts = MAX_RETRIES + 1 - operation.attemptsMade;
+      const remainingAttempts = this.clientOptions.maxAdaptiveRetries + 1 - operation.attemptsMade;
       if (remainingAttempts <= 0) {
         throw firstCommitError;
       }