Implement the overload retry policy for commands that are read/write retryable

config/checkstyle/checkstyle.xml

@@ -122,7 +122,7 @@
 
         <module name="MethodLength"/>
         <module name="ParameterNumber">
-            <property name="max" value="12"/>
+            <property name="max" value="13"/>
         </module>
 
 

driver-core/src/main/com/mongodb/MongoClientSettings.java

@@ -32,6 +32,7 @@
 import com.mongodb.connection.SslSettings;
 import com.mongodb.connection.TransportSettings;
 import com.mongodb.event.CommandListener;
+import com.mongodb.internal.operation.CommandOperationHelper;
 import com.mongodb.lang.Nullable;
 import com.mongodb.observability.ObservabilitySettings;
 import com.mongodb.spi.dns.DnsClient;
@@ -503,9 +504,11 @@ public Builder retryReads(final boolean retryReads) {
          * the {@value MongoException#SYSTEM_OVERLOADED_ERROR_LABEL} and {@value MongoException#RETRYABLE_ERROR_LABEL} labels.
          * Such errors are referred to as retryable overload errors.
          * <p>
-         * Default is {@code null}, implies the value 2 and the above retry behavior. The implied value and behavior may change in
+         * Default is {@code null}, implies the value {@value CommandOperationHelper#DEFAULT_MAX_ADAPTIVE_RETRIES} and the above retry behavior.
+         * The implied value and behavior may change in
          * the future in a <a href="https://semver.org/spec/v2.0.0.html#summary">minor version</a>.
-         * This means, there is no guarantee that not setting a value is equivalent to setting the value 2.
+         * This means, there is no guarantee that not setting a value is equivalent to setting the value
+         * {@value CommandOperationHelper#DEFAULT_MAX_ADAPTIVE_RETRIES}.
          * The value 0 results in not retrying the attempts failed due to retryable overload errors.
          *
          * <table>
@@ -961,7 +964,6 @@ public boolean getRetryReads() {
      */
     @Beta(Reason.CLIENT)
     @Nullable
-    // TODO-BACKPRESSURE Valentin Use the `maxAdaptiveRetries` setting when retrying.
     public Integer getMaxAdaptiveRetries() {
         return maxAdaptiveRetries;
     }

driver-core/src/main/com/mongodb/internal/operation/AbortTransactionOperation.java

@@ -36,8 +36,8 @@ public class AbortTransactionOperation extends TransactionOperation {
     private static final String COMMAND_NAME = "abortTransaction";
     private BsonDocument recoveryToken;
 
-    public AbortTransactionOperation(final WriteConcern writeConcern) {
-        super(writeConcern);
+    public AbortTransactionOperation(final WriteConcern writeConcern, @Nullable final Integer maxAdaptiveRetriesSetting) {
+        super(writeConcern, maxAdaptiveRetriesSetting);
     }
 
     public AbortTransactionOperation recoveryToken(@Nullable final BsonDocument recoveryToken) {

driver-core/src/main/com/mongodb/internal/operation/AggregateOperation.java

@@ -45,13 +45,14 @@
 public class AggregateOperation<T> implements ReadOperationExplainable<T> {
     private final AggregateOperationImpl<T> wrapped;
 
-    public AggregateOperation(final MongoNamespace namespace, final List<BsonDocument> pipeline, final Decoder<T> decoder) {
-        this(namespace, pipeline, decoder, AggregationLevel.COLLECTION);
+    public AggregateOperation(final MongoNamespace namespace, final List<BsonDocument> pipeline, final Decoder<T> decoder,
+            @Nullable final Integer maxAdaptiveRetriesSetting) {
+        this(namespace, pipeline, decoder, AggregationLevel.COLLECTION, maxAdaptiveRetriesSetting);
     }
 
     public AggregateOperation(final MongoNamespace namespace, final List<BsonDocument> pipeline, final Decoder<T> decoder,
-            final AggregationLevel aggregationLevel) {
-        this.wrapped = new AggregateOperationImpl<>(namespace, pipeline, decoder, aggregationLevel);
+            final AggregationLevel aggregationLevel, @Nullable final Integer maxAdaptiveRetriesSetting) {
+        this.wrapped = new AggregateOperationImpl<>(namespace, pipeline, decoder, aggregationLevel, maxAdaptiveRetriesSetting);
     }
 
     public List<BsonDocument> getPipeline() {

driver-core/src/main/com/mongodb/internal/operation/AggregateOperationImpl.java

@@ -65,6 +65,8 @@ class AggregateOperationImpl<T> implements ReadOperationCursor<T> {
     private final PipelineCreator pipelineCreator;
 
     private boolean retryReads;
+    @Nullable
+    private final Integer maxAdaptiveRetriesSetting;
     private Boolean allowDiskUse;
     private Integer batchSize;
     private Collation collation;
@@ -75,21 +77,24 @@ class AggregateOperationImpl<T> implements ReadOperationCursor<T> {
     private CursorType cursorType;
 
     AggregateOperationImpl(final MongoNamespace namespace,
-            final List<BsonDocument> pipeline, final Decoder<T> decoder, final AggregationLevel aggregationLevel) {
+            final List<BsonDocument> pipeline, final Decoder<T> decoder, final AggregationLevel aggregationLevel,
+            @Nullable final Integer maxAdaptiveRetriesSetting) {
         this(namespace, pipeline, decoder,
                 defaultAggregateTarget(notNull("aggregationLevel", aggregationLevel),
                         notNull("namespace", namespace).getCollectionName()),
-                defaultPipelineCreator(pipeline));
+                defaultPipelineCreator(pipeline), maxAdaptiveRetriesSetting);
     }
 
     AggregateOperationImpl(final MongoNamespace namespace,
             final List<BsonDocument> pipeline, final Decoder<T> decoder, final AggregateTarget aggregateTarget,
-            final PipelineCreator pipelineCreator) {
+            final PipelineCreator pipelineCreator,
+            @Nullable final Integer maxAdaptiveRetriesSetting) {
         this.namespace = notNull("namespace", namespace);
         this.pipeline = notNull("pipeline", pipeline);
         this.decoder = notNull("decoder", decoder);
         this.aggregateTarget = notNull("aggregateTarget", aggregateTarget);
         this.pipelineCreator = notNull("pipelineCreator", pipelineCreator);
+        this.maxAdaptiveRetriesSetting = maxAdaptiveRetriesSetting;
     }
 
     List<BsonDocument> getPipeline() {
@@ -196,15 +201,15 @@ public MongoNamespace getNamespace() {
     public BatchCursor<T> execute(final ReadBinding binding, final OperationContext operationContext) {
         return executeRetryableRead(binding, applyTimeoutModeToOperationContext(timeoutMode, operationContext), namespace.getDatabaseName(),
                 getCommandCreator(), CommandResultDocumentCodec.create(decoder, FIELD_NAMES_WITH_RESULT),
-                transformer(), retryReads);
+                transformer(), retryReads, maxAdaptiveRetriesSetting);
     }
 
     @Override
     public void executeAsync(final AsyncReadBinding binding, final OperationContext operationContext, final SingleResultCallback<AsyncBatchCursor<T>> callback) {
         SingleResultCallback<AsyncBatchCursor<T>> errHandlingCallback = errorHandlingCallback(callback, LOGGER);
         executeRetryableReadAsync(binding, applyTimeoutModeToOperationContext(timeoutMode, operationContext),  namespace.getDatabaseName(),
                 getCommandCreator(), CommandResultDocumentCodec.create(decoder, FIELD_NAMES_WITH_RESULT),
-                asyncTransformer(), retryReads,
+                asyncTransformer(), retryReads, maxAdaptiveRetriesSetting,
                 errHandlingCallback);
     }
 

driver-core/src/main/com/mongodb/internal/operation/AggregateToCollectionOperation.java

@@ -178,7 +178,8 @@ public Void execute(final ReadBinding binding, final OperationContext operationC
                 getCommandCreator(),
                 new BsonDocumentCodec(),
                 transformer(),
-                false);
+                false,
+                null);
     }
 
     @Override
@@ -194,6 +195,7 @@ public void executeAsync(final AsyncReadBinding binding, final OperationContext
                 new BsonDocumentCodec(),
                 asyncTransformer(),
                 false,
+                null,
                 callback);
     }
 

driver-core/src/main/com/mongodb/internal/operation/AsyncOperationHelper.java

@@ -176,9 +176,10 @@ static <D, T> void executeRetryableReadAsync(
             final Decoder<D> decoder,
             final CommandReadTransformerAsync<D, T> transformer,
             final boolean retryReadsSetting,
+            @Nullable final Integer maxAdaptiveRetriesSetting,
             final SingleResultCallback<T> callback) {
         executeRetryableReadAsync(binding, operationContext, binding::getReadConnectionSource, database, commandCreator,
-                                  decoder, transformer, retryReadsSetting, callback);
+                                  decoder, transformer, retryReadsSetting, maxAdaptiveRetriesSetting, callback);
     }
 
     static <D, T> void executeRetryableReadAsync(
@@ -190,9 +191,10 @@ static <D, T> void executeRetryableReadAsync(
             final Decoder<D> decoder,
             final CommandReadTransformerAsync<D, T> transformer,
             final boolean retryReadsSetting,
+            @Nullable final Integer maxAdaptiveRetriesSetting,
             final SingleResultCallback<T> callback) {
         RetryControl<SpecRetryPolicy> retryControl = createSpecRetryControl(
-                new SpecRetryPolicy.IndividualPolicies(retryReadsSetting).includeRead(operationContext),
+                new SpecRetryPolicy.IndividualPolicies(retryReadsSetting).includeRead(operationContext).includeOverload(maxAdaptiveRetriesSetting),
                 operationContext);
         binding.retain();
         AsyncCallbackSupplier<T> asyncRead = decorateWithRetriesAsync(retryControl, operationContext,
@@ -257,12 +259,13 @@ static <T, R> void executeRetryableWriteAsync(
             final CommandWriteTransformerAsync<T, R> transformer,
             final Function<BsonDocument, BsonDocument> retryCommandModifier,
             final boolean effectiveRetryWritesSetting,
+            @Nullable final Integer maxAdaptiveRetriesSetting,
             final SingleResultCallback<R> callback) {
         beginAsync().<R>thenSupply(c -> {
             binding.retain();
             MutableValue<BsonDocument> command = new MutableValue<>();
             RetryControl<SpecRetryPolicy> retryControl = createSpecRetryControl(
-                    new SpecRetryPolicy.IndividualPolicies(effectiveRetryWritesSetting).includeWrite(),
+                    new SpecRetryPolicy.IndividualPolicies(effectiveRetryWritesSetting).includeWrite().includeOverload(maxAdaptiveRetriesSetting),
                     operationContext);
             AsyncCallbackSupplier<R> retryingWrite = decorateWithRetriesAsync(retryControl, operationContext, supplierCallback -> {
                 beginAsync().<R>thenSupply(withSourceAndConnectionCallback -> {

driver-core/src/main/com/mongodb/internal/operation/BaseFindAndModifyOperation.java

@@ -51,6 +51,8 @@ public abstract class BaseFindAndModifyOperation<T> implements WriteOperation<T>
     private final MongoNamespace namespace;
     private final WriteConcern writeConcern;
     private final boolean retryWrites;
+    @Nullable
+    private final Integer maxAdaptiveRetriesSetting;
     private final Decoder<T> decoder;
 
     private BsonDocument filter;
@@ -62,11 +64,14 @@ public abstract class BaseFindAndModifyOperation<T> implements WriteOperation<T>
     private BsonValue comment;
     private BsonDocument variables;
 
-    protected BaseFindAndModifyOperation(final MongoNamespace namespace, final WriteConcern writeConcern, final boolean retryWrites,
+    protected BaseFindAndModifyOperation(final MongoNamespace namespace, final WriteConcern writeConcern,
+            final boolean retryWrites,
+            @Nullable final Integer maxAdaptiveRetriesSetting,
             final Decoder<T> decoder) {
         this.namespace = notNull("namespace", namespace);
         this.writeConcern = notNull("writeConcern", writeConcern);
         this.retryWrites = retryWrites;
+        this.maxAdaptiveRetriesSetting = maxAdaptiveRetriesSetting;
         this.decoder = notNull("decoder", decoder);
     }
 
@@ -83,15 +88,16 @@ public T execute(final WriteBinding binding, final OperationContext operationCon
                                      getCommandCreator(),
                                      FindAndModifyHelper.transformer(),
                                      cmd -> cmd,
-                                     retryWrites);
+                                     retryWrites,
+                                     maxAdaptiveRetriesSetting);
     }
 
     @Override
     public void executeAsync(final AsyncWriteBinding binding, final OperationContext operationContext, final SingleResultCallback<T> callback) {
         executeRetryableWriteAsync(binding, operationContext, getDatabaseName(), null, getFieldNameValidator(),
                                    CommandResultDocumentCodec.create(getDecoder(), "value"),
                                    getCommandCreator(),
-                FindAndModifyHelper.asyncTransformer(), cmd -> cmd, retryWrites, callback);
+                FindAndModifyHelper.asyncTransformer(), cmd -> cmd, retryWrites, maxAdaptiveRetriesSetting, callback);
     }
 
     @Override

driver-core/src/main/com/mongodb/internal/operation/ChangeStreamOperation.java

@@ -69,14 +69,15 @@ public class ChangeStreamOperation<T> implements ReadOperationCursor<T> {
     @VisibleForTesting(otherwise = VisibleForTesting.AccessModifier.PRIVATE)
     public ChangeStreamOperation(final MongoNamespace namespace, final FullDocument fullDocument,
             final FullDocumentBeforeChange fullDocumentBeforeChange, final List<BsonDocument> pipeline, final Decoder<T> decoder) {
-        this(namespace, fullDocument, fullDocumentBeforeChange, pipeline, decoder, ChangeStreamLevel.COLLECTION);
+        this(namespace, fullDocument, fullDocumentBeforeChange, pipeline, decoder, ChangeStreamLevel.COLLECTION, null);
     }
 
     public ChangeStreamOperation(final MongoNamespace namespace, final FullDocument fullDocument,
             final FullDocumentBeforeChange fullDocumentBeforeChange, final List<BsonDocument> pipeline, final Decoder<T> decoder,
-            final ChangeStreamLevel changeStreamLevel) {
+            final ChangeStreamLevel changeStreamLevel,
+            @Nullable final Integer maxAdaptiveRetriesSetting) {
         this.wrapped = new AggregateOperationImpl<>(namespace, pipeline, RAW_BSON_DOCUMENT_CODEC, getAggregateTarget(),
-                getPipelineCreator()).cursorType(CursorType.TailableAwait);
+                getPipelineCreator(), maxAdaptiveRetriesSetting).cursorType(CursorType.TailableAwait);
         this.fullDocument = notNull("fullDocument", fullDocument);
         this.fullDocumentBeforeChange = notNull("fullDocumentBeforeChange", fullDocumentBeforeChange);
         this.decoder = notNull("decoder", decoder);

driver-core/src/main/com/mongodb/internal/operation/ClientBulkWriteOperation.java

@@ -156,6 +156,8 @@ public final class ClientBulkWriteOperation implements WriteOperation<ClientBulk
     private final ConcreteClientBulkWriteOptions options;
     private final WriteConcern writeConcernSetting;
     private final boolean retryWritesSetting;
+    @Nullable
+    private final Integer maxAdaptiveRetriesSetting;
     private final CodecRegistry codecRegistry;
 
     /**
@@ -166,11 +168,13 @@ public ClientBulkWriteOperation(
             @Nullable final ClientBulkWriteOptions options,
             final WriteConcern writeConcernSetting,
             final boolean retryWritesSetting,
+            @Nullable final Integer maxAdaptiveRetriesSetting,
             final CodecRegistry codecRegistry) {
         this.models = models;
         this.options = options == null ? EMPTY_OPTIONS : (ConcreteClientBulkWriteOptions) options;
         this.writeConcernSetting = writeConcernSetting;
         this.retryWritesSetting = retryWritesSetting;
+        this.maxAdaptiveRetriesSetting = maxAdaptiveRetriesSetting;
         this.codecRegistry = codecRegistry;
     }
 
@@ -280,7 +284,7 @@ private Integer executeBatch(
         assertFalse(unexecutedModels.isEmpty());
         SessionContext sessionContext = operationContext.getSessionContext();
         RetryControl<SpecRetryPolicy> retryControl = createSpecRetryControl(
-                new SpecRetryPolicy.IndividualPolicies(retryWritesSetting).includeWrite(),
+                new SpecRetryPolicy.IndividualPolicies(retryWritesSetting).includeWrite().includeOverload(maxAdaptiveRetriesSetting),
                 operationContext);
         BatchEncoder batchEncoder = new BatchEncoder();
 
@@ -334,7 +338,7 @@ private void executeBatchAsync(
             assertFalse(unexecutedModels.isEmpty());
             SessionContext sessionContext = operationContext.getSessionContext();
             RetryControl<SpecRetryPolicy> retryControl = createSpecRetryControl(
-                    new SpecRetryPolicy.IndividualPolicies(retryWritesSetting).includeWrite(),
+                    new SpecRetryPolicy.IndividualPolicies(retryWritesSetting).includeWrite().includeOverload(maxAdaptiveRetriesSetting),
                     operationContext);
             BatchEncoder batchEncoder = new BatchEncoder();
 

driver-core/src/main/com/mongodb/internal/operation/CommandOperationHelper.java

@@ -17,6 +17,7 @@
 package com.mongodb.internal.operation;
 
 import com.mongodb.MongoClientException;
+import com.mongodb.MongoClientSettings;
 import com.mongodb.MongoCommandException;
 import com.mongodb.MongoConnectionPoolClearedException;
 import com.mongodb.MongoException;
@@ -42,6 +43,11 @@
 
 @SuppressWarnings("overloads")
 public final class CommandOperationHelper {
+    /**
+     * The value used when {@link MongoClientSettings#getMaxAdaptiveRetries()} is {@code null}.
+     */
+    public static final int DEFAULT_MAX_ADAPTIVE_RETRIES = 2;
+
     static WriteConcern validateAndGetEffectiveWriteConcern(final WriteConcern writeConcernSetting, final SessionContext sessionContext)
             throws MongoClientException {
         boolean activeTransaction = sessionContext.hasActiveTransaction();

driver-core/src/main/com/mongodb/internal/operation/CommandReadOperation.java

@@ -68,14 +68,14 @@ public MongoNamespace getNamespace() {
     @Override
     public T execute(final ReadBinding binding, final OperationContext operationContext) {
         return executeRetryableRead(binding, operationContext, databaseName, commandCreator, decoder,
-                transformer(), false);
+                transformer(), false, null);
     }
 
     @Override
     public void executeAsync(final AsyncReadBinding binding, final OperationContext operationContext,
                              final SingleResultCallback<T> callback) {
         executeRetryableReadAsync(binding, operationContext,  databaseName, commandCreator, decoder,
-                asyncTransformer(), false, callback);
+                asyncTransformer(), false, null, callback);
     }
 
     private static <T> CommandReadTransformer<T, T> transformer() {

driver-core/src/main/com/mongodb/internal/operation/CommitTransactionOperation.java

@@ -53,12 +53,8 @@ public class CommitTransactionOperation extends TransactionOperation {
     private final boolean alreadyCommitted;
     private BsonDocument recoveryToken;
 
-    public CommitTransactionOperation(final WriteConcern writeConcern) {
-        this(writeConcern, false);
-    }
-
-    public CommitTransactionOperation(final WriteConcern writeConcern, final boolean alreadyCommitted) {
-        super(writeConcern);
+    public CommitTransactionOperation(final WriteConcern writeConcern, @Nullable final Integer maxAdaptiveRetriesSetting, final boolean alreadyCommitted) {
+        super(writeConcern, maxAdaptiveRetriesSetting);
         this.alreadyCommitted = alreadyCommitted;
     }
 

driver-core/src/main/com/mongodb/internal/operation/CountDocumentsOperation.java

@@ -43,15 +43,18 @@ public class CountDocumentsOperation implements ReadOperationSimple<Long> {
     private static final Decoder<BsonDocument> DECODER = new BsonDocumentCodec();
     private final MongoNamespace namespace;
     private boolean retryReads;
+    @Nullable
+    private final Integer maxAdaptiveRetriesSetting;
     private BsonDocument filter;
     private BsonValue hint;
     private BsonValue comment;
     private long skip;
     private long limit;
     private Collation collation;
 
-    public CountDocumentsOperation(final MongoNamespace namespace) {
+    public CountDocumentsOperation(final MongoNamespace namespace, @Nullable final Integer maxAdaptiveRetriesSetting) {
         this.namespace = notNull("namespace", namespace);
+        this.maxAdaptiveRetriesSetting = maxAdaptiveRetriesSetting;
     }
 
     @Nullable
@@ -157,7 +160,7 @@ public void executeAsync(final AsyncReadBinding binding, final OperationContext
     }
 
     private AggregateOperation<BsonDocument> getAggregateOperation() {
-        return new AggregateOperation<>(namespace, getPipeline(), DECODER)
+        return new AggregateOperation<>(namespace, getPipeline(), DECODER, maxAdaptiveRetriesSetting)
                 .retryReads(retryReads)
                 .collation(collation)
                 .comment(comment)