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/async/function/RetryControl.java

@@ -80,15 +80,22 @@ public P getPolicy() {
      * @return {@link RetryAttemptInfo} iff another attempt must be executed.
      * @throws RuntimeException If another attempt must not be executed.
      * The exception thrown represents the failed result of the retryable activity.
+     * @throws Error See above for {@link RuntimeException}.
      */
-    RetryAttemptInfo advanceOrThrow(final Throwable attemptFailedResult) throws RuntimeException {
+    RetryAttemptInfo advanceOrThrow(final Throwable attemptFailedResult) throws RuntimeException, Error {
         assertNotNull(attemptFailedResult);
         try {
             if (disabled) {
                 throw attemptFailedResult;
             }
-            // this `RetryControl` must not be mutated before calling `onAttemptFailure`
-            Decision decision = onAttemptFailure(policy, this, prospectiveFailedResult(), attemptFailedResult);
+            Throwable prospectiveFailedResult = prospectiveFailedResult();
+            if (attemptFailedResult instanceof Error) {
+                onAttemptFatalFailure(policy, prospectiveFailedResult, (Error) attemptFailedResult);
+                policy.onAttemptFatalFailure();
+                throw attemptFailedResult;
+            }
+            // this `RetryControl` must not be mutated before calling `onAttemptFatalFailure`/`onAttemptFailure`
+            Decision decision = onAttemptFailure(policy, this, prospectiveFailedResult, attemptFailedResult);
             mostRecentDecision = decision;
             if (loopControl.isLastIteration() || !decision.getImmediateNextAttemptInfo().isPresent()) {
                 throw decision.getProspectiveFailedResult();
@@ -103,24 +110,48 @@ RetryAttemptInfo advanceOrThrow(final Throwable attemptFailedResult) throws Runt
         }
     }
 
+    private static <P extends RetryPolicy> void onAttemptFatalFailure(
+            final P policy,
+            @Nullable final Throwable prospectiveFailedResult,
+            final Error attemptFailedResult) throws RuntimeException, Error {
+        doAndHandleException(
+                () -> {
+                    policy.onAttemptFatalFailure();
+                    return null;
+                },
+                prospectiveFailedResult,
+                attemptFailedResult);
+    }
+
     private static <P extends RetryPolicy> Decision onAttemptFailure(
             final P policy,
             final RetryContext retryContext,
             @Nullable final Throwable prospectiveFailedResult,
-            final Throwable attemptFailedResult) throws RuntimeException {
-        Decision decision;
+            final Throwable attemptFailedResult) throws RuntimeException, Error {
+        return assertNotNull(doAndHandleException(
+                () -> policy.onAttemptFailure(retryContext, attemptFailedResult),
+                prospectiveFailedResult,
+                attemptFailedResult));
+    }
+
+    @Nullable
+    private static <T> T doAndHandleException(
+            final Supplier<T> action,
+            @Nullable final Throwable prospectiveFailedResult,
+            final Throwable attemptFailedResult) throws RuntimeException, Error {
+        T result;
         try {
-            decision = assertNotNull(policy.onAttemptFailure(retryContext, attemptFailedResult));
-        } catch (Throwable onAttemptFailureException) {
-            if (prospectiveFailedResult != null && prospectiveFailedResult != onAttemptFailureException) {
-                onAttemptFailureException.addSuppressed(prospectiveFailedResult);
+            result = action.get();
+        } catch (Throwable actionException) {
+            if (prospectiveFailedResult != null && prospectiveFailedResult != actionException) {
+                actionException.addSuppressed(prospectiveFailedResult);
             }
-            if (attemptFailedResult != onAttemptFailureException) {
-                onAttemptFailureException.addSuppressed(attemptFailedResult);
+            if (attemptFailedResult != actionException) {
+                actionException.addSuppressed(attemptFailedResult);
             }
-            throw onAttemptFailureException;
+            throw actionException;
         }
-        return decision;
+        return result;
     }
 
     @Override
@@ -162,7 +193,7 @@ private Throwable prospectiveFailedResult() {
      *     <li>this method broke the retry loop.</li>
      * </ul>
      */
-    public void breakAndThrowIfRetryAnd(final Supplier<Boolean> predicate) throws RuntimeException {
+    public void breakAndThrowIfRetryAnd(final Supplier<Boolean> predicate) throws RuntimeException, Error {
         assertFalse(loopControl.isLastIteration());
         if (isFirstAttempt()) {
             return;

driver-core/src/main/com/mongodb/internal/async/function/RetryPolicy.java

@@ -28,25 +28,41 @@
 /**
  * Customizes retrying and may allow for control beyond what {@link RetryControl} itself provides, depending on the implementation.
  * <p>
- * An implementation may be stateful and does not have to be thread-safe.
+ * An implementation may be stateful.
+ * An implementation must be thread-safe iff a {@link RetryPolicy} instance is shared by multiple concurrent executions of a retryable activity.
  * <p>
  * This class is not part of the public API and may be removed or changed at any time.
  */
 @NotThreadSafe
 public interface RetryPolicy {
     /**
      * This method is called exactly once per failed attempt,
-     * even if that is the {@linkplain RetryControl#breakAndThrowIfRetryAnd(Supplier) last attempt},
-     * provided that retrying is not {@linkplain RetryControl#doWhileDisabled(Supplier) disabled}.
+     * even if that is the last attempt due to {@linkplain RetryControl#breakAndThrowIfRetryAnd(Supplier) breaking} the retry loop,
+     * provided that retrying is not {@linkplain RetryControl#doWhileDisabled(Supplier) disabled},
+     * and the failed result of the attempt is not {@linkplain #onAttemptFatalFailure() fatal}.
+     * <p>
      * If this method completes abruptly, then another attempt is not executed,
      * and the exception thrown by the method is used as the failed result of the retryable activity.
      * <p>
      * This method may have side effects, and may mutate {@link RetryContext#getProspectiveFailedResult()}, {@code attemptFailedResult}.
      *
      * @param attemptFailedResult The failed result of the most recent attempt.
+     * @see #onAttemptFatalFailure()
      */
     Decision onAttemptFailure(RetryContext retryContext, Throwable attemptFailedResult);
 
+    /**
+     * This method is similar to {@link #onAttemptFailure(RetryContext, Throwable)}, with the following difference:
+     * it is called only if the failed result of the most recent attempt is an {@link Error}.
+     * The method intentionally accepts neither {@link RetryContext}, nor information about the {@link Error},
+     * to discourage the implementation from doing much.
+     * <p>
+     * If this method is invoked, then another attempt is not executed regardless of whether the method completes normally or abruptly.
+     * However, if the method completes abruptly, then the exception thrown by the method is used as the failed result of the retryable activity.
+     */
+    default void onAttemptFatalFailure() {
+    }
+
     final class Decision {
         private final Throwable prospectiveFailedResult;
         @Nullable
@@ -71,7 +87,7 @@ public Throwable getProspectiveFailedResult() {
         /**
          * Returns {@link Optional#isEmpty()} to signal that another attempt must not be executed.
          * If {@link RetryAttemptInfo} is {@linkplain Optional#isPresent() present},
-         * another attempt is still not executed if most recent attempt was the {@linkplain RetryControl#breakAndThrowIfRetryAnd(Supplier) last one}.
+         * another attempt is still not executed if the retry loop was {@linkplain RetryControl#breakAndThrowIfRetryAnd(Supplier) broken}.
          */
         public Optional<RetryAttemptInfo> getImmediateNextAttemptInfo() {
             return Optional.ofNullable(immediateNextAttemptInfo);

driver-core/src/main/com/mongodb/internal/async/function/RetryingAsyncCallbackSupplier.java

@@ -23,6 +23,7 @@
 
 import java.time.Duration;
 
+import static com.mongodb.assertions.Assertions.fail;
 import static com.mongodb.internal.async.AsyncRunnable.beginAsync;
 
 /**
@@ -69,7 +70,8 @@ public void get(final SingleResultCallback<R> callback) {
                 onAttemptSuccessCallback.complete(onAttemptSuccessCallback);
             }).onErrorIf(e -> true, (attemptFailedResult, onAttemptFailureCallback) -> {
                 if (attemptFailedResult instanceof Error) {
-                    onAttemptFailureCallback.completeExceptionally(attemptFailedResult);
+                    control.advanceOrThrow(attemptFailedResult);
+                    fail("Must not be reached");
                 } else {
                     RetryAttemptInfo retryAttemptInfo = control.advanceOrThrow(attemptFailedResult);
                     clientExecutor.sleepAsync(retryAttemptInfo.getBackoff(), onAttemptFailureCallback);

driver-core/src/main/com/mongodb/internal/async/function/RetryingSyncSupplier.java

@@ -24,6 +24,7 @@
 import java.time.Duration;
 import java.util.function.Supplier;
 
+import static com.mongodb.assertions.Assertions.fail;
 import static com.mongodb.internal.thread.InterruptionUtil.interruptAndCreateMongoInterruptedException;
 import static java.util.concurrent.TimeUnit.NANOSECONDS;
 
@@ -60,7 +61,8 @@ public R get() {
                 // `attemptSuccessfulResult` may be `null`, so we have to wrap it in `MutableValue` for the while check to notice it
                 asyncFunctionSuccessfulResult.set(new MutableValue<>(attemptSuccessfulResult));
             } catch (Error attemptFailedResult) {
-                throw attemptFailedResult;
+                control.advanceOrThrow(attemptFailedResult);
+                fail("Must not be reached");
             } catch (Throwable attemptFailedResult) {
                 RetryAttemptInfo retryAttemptInfo = control.advanceOrThrow(attemptFailedResult);
                 sleep(retryAttemptInfo.getBackoff());

driver-core/src/main/com/mongodb/internal/connection/ClusterClockAdvancingSessionContext.java

@@ -17,6 +17,7 @@
 package com.mongodb.internal.connection;
 
 import com.mongodb.ReadConcern;
+import com.mongodb.internal.session.BaseClientSessionImpl.OverloadRetryPolicyState;
 import com.mongodb.internal.session.SessionContext;
 import com.mongodb.lang.Nullable;
 import org.bson.BsonDocument;
@@ -136,4 +137,9 @@ public void markSessionDirty() {
     public boolean isSessionMarkedDirty() {
         return wrapped.isSessionMarkedDirty();
     }
+
+    @Override
+    public OverloadRetryPolicyState getOverloadRetryPolicyState() {
+        return wrapped.getOverloadRetryPolicyState();
+    }
 }

driver-core/src/main/com/mongodb/internal/connection/CommandMessage.java

@@ -344,12 +344,12 @@ private List<BsonElement> getExtraElements(final OperationContext operationConte
                 extraElements.add(new BsonElement("lsid", sessionContext.getSessionId()));
             }
         }
-        boolean firstMessageInTransaction = sessionContext.notifyMessageSent();
+        boolean startTransaction = sessionContext.notifyMessageSent();
 
         assertFalse(sessionContext.hasActiveTransaction() && sessionContext.isSnapshot());
         if (sessionContext.hasActiveTransaction()) {
             extraElements.add(new BsonElement("txnNumber", new BsonInt64(sessionContext.getTransactionNumber())));
-            if (firstMessageInTransaction) {
+            if (startTransaction) {
                 extraElements.add(new BsonElement("startTransaction", BsonBoolean.TRUE));
                 addReadConcernDocument(extraElements, sessionContext);
             }

driver-core/src/main/com/mongodb/internal/connection/NoOpSessionContext.java

@@ -17,6 +17,7 @@
 package com.mongodb.internal.connection;
 
 import com.mongodb.ReadConcern;
+import com.mongodb.internal.session.BaseClientSessionImpl.OverloadRetryPolicyState;
 import com.mongodb.internal.session.SessionContext;
 import com.mongodb.lang.Nullable;
 import org.bson.BsonDocument;
@@ -34,6 +35,8 @@ public class NoOpSessionContext implements SessionContext {
      */
     public static final NoOpSessionContext INSTANCE = new NoOpSessionContext();
 
+    private final OverloadRetryPolicyState overloadRetryPolicyState = new OverloadRetryPolicyState();
+
     @Override
     public boolean hasSession() {
         return false;
@@ -133,4 +136,9 @@ public void markSessionDirty() {
     public boolean isSessionMarkedDirty() {
         return false;
     }
+
+    @Override
+    public OverloadRetryPolicyState getOverloadRetryPolicyState() {
+        return overloadRetryPolicyState;
+    }
 }

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

@@ -21,6 +21,7 @@
 import com.mongodb.WriteConcern;
 import com.mongodb.internal.MongoNamespaceHelper;
 import com.mongodb.internal.TimeoutContext;
+import com.mongodb.internal.connection.OperationContext;
 import com.mongodb.internal.operation.CommandOperationHelper.CommandCreator;
 import com.mongodb.lang.Nullable;
 import org.bson.BsonDocument;
@@ -36,8 +37,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) {
@@ -68,7 +69,7 @@ CommandCreator getCommandCreator() {
     }
 
     @Override
-    protected Function<BsonDocument, BsonDocument> getRetryCommandModifier(final TimeoutContext timeoutContext) {
+    protected Function<BsonDocument, BsonDocument> getRetryCommandModifier(final OperationContext operationContext) {
         return cmd -> cmd;
     }
 }

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);
     }