Direct retries to another mongos if one is available (#1367)

JAVA-4254

---------

Co-authored-by: Maxim Katcharov <[email protected]>

Author: stIncMale

driver-core/src/main/com/mongodb/connection/ClusterSettings.java

@@ -468,16 +468,18 @@ public String getRequiredReplicaSetName() {
      *
      * <p>The server selector augments the normal server selection rules applied by the driver when determining
      * which server to send an operation to.  At the point that it's called by the driver, the
-     * {@link com.mongodb.connection.ClusterDescription} which is passed to it contains a list of
-     * {@link com.mongodb.connection.ServerDescription} instances which satisfy either the configured {@link com.mongodb.ReadPreference}
-     * for any read operation or ones that can take writes (e.g. a standalone, mongos, or replica set primary).
+     * {@link ClusterDescription} which is passed to it {@linkplain ClusterDescription#getServerDescriptions() contains} a list of
+     * {@link ServerDescription} instances which satisfy either the configured {@link com.mongodb.ReadPreference}
+     * for any read operation or ones that can take writes (e.g. a standalone, mongos, or replica set primary),
+     * barring those corresponding to servers that the driver considers unavailable or potentially problematic.
      * </p>
      * <p>The server selector can then filter the {@code ServerDescription} list using whatever criteria that is required by the
      * application.</p>
-     * <p>After this selector executes, two additional selectors are applied by the driver:</p>
+     * <p>After this selector executes, three additional selectors are applied by the driver:</p>
      * <ul>
      * <li>select from within the latency window</li>
-     * <li>select a random server from those remaining</li>
+     * <li>select at most two random servers from those remaining</li>
+     * <li>select the one with fewer outstanding concurrent operations</li>
      * </ul>
      * <p>To skip the latency window selector, an application can:</p>
      * <ul>
@@ -486,6 +488,7 @@ public String getRequiredReplicaSetName() {
      * </ul>
      *
      * @return the server selector, which may be null
+     * @see Builder#serverSelector(ServerSelector)
      */
     @Nullable
     public ServerSelector getServerSelector() {

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

@@ -78,32 +78,33 @@ public RetryState() {
      * which is usually synchronous code.
      *
      * @param attemptException The exception produced by the most recent attempt.
-     * It is passed to the {@code retryPredicate} and to the {@code exceptionTransformer}.
-     * @param exceptionTransformer A function that chooses which exception to preserve as a prospective failed result of the associated
-     * retryable activity and may also transform or mutate the exceptions.
-     * The choice is between
+     * It is passed to the {@code retryPredicate} and to the {@code onAttemptFailureOperator}.
+     * @param onAttemptFailureOperator The action that is called once per failed attempt before (in the happens-before order) the
+     * {@code retryPredicate}, regardless of whether the {@code retryPredicate} is called.
+     * This action is allowed to have side effects.
+     * <p>
+     * It also has to choose which exception to preserve as a prospective failed result of the associated retryable activity.
+     * The {@code onAttemptFailureOperator} may mutate its arguments, choose from the arguments, or return a different exception,
+     * but it must return a {@code @}{@link NonNull} value.
+     * The choice is between</p>
      * <ul>
      *     <li>the previously chosen exception or {@code null} if none has been chosen
-     *     (the first argument of the {@code exceptionTransformer})</li>
-     *     <li>and the exception from the most recent attempt (the second argument of the {@code exceptionTransformer}).</li>
+     *     (the first argument of the {@code onAttemptFailureOperator})</li>
+     *     <li>and the exception from the most recent attempt (the second argument of the {@code onAttemptFailureOperator}).</li>
      * </ul>
-     * The {@code exceptionTransformer} may either choose from its arguments, or return a different exception, a.k.a. transform,
-     * but it must return a {@code @}{@link NonNull} value.
-     * The {@code exceptionTransformer} is called once before (in the happens-before order) the {@code retryPredicate},
-     * regardless of whether the {@code retryPredicate} is called. The result of the {@code exceptionTransformer} does not affect
-     * what exception is passed to the {@code retryPredicate}.
+     * The result of the {@code onAttemptFailureOperator} does not affect the exception passed to the {@code retryPredicate}.
      * @param retryPredicate {@code true} iff another attempt needs to be made. The {@code retryPredicate} is called not more than once
      * per attempt and only if all the following is true:
      * <ul>
-     *     <li>{@code exceptionTransformer} completed normally;</li>
+     *     <li>{@code onAttemptFailureOperator} completed normally;</li>
      *     <li>the most recent attempt is not the {@linkplain #isLastAttempt() last} one.</li>
      * </ul>
      * The {@code retryPredicate} accepts this {@link RetryState} and the exception from the most recent attempt,
      * and may mutate the exception. The {@linkplain RetryState} advances to represent the state of a new attempt
      * after (in the happens-before order) testing the {@code retryPredicate}, and only if the predicate completes normally.
      * @throws RuntimeException Iff any of the following is true:
      * <ul>
-     *     <li>the {@code exceptionTransformer} completed abruptly;</li>
+     *     <li>the {@code onAttemptFailureOperator} completed abruptly;</li>
      *     <li>the most recent attempt is the {@linkplain #isLastAttempt() last} one;</li>
      *     <li>the {@code retryPredicate} completed abruptly;</li>
      *     <li>the {@code retryPredicate} is {@code false}.</li>
@@ -112,10 +113,10 @@ public RetryState() {
      * i.e., the caller must not do any more attempts.
      * @see #advanceOrThrow(Throwable, BinaryOperator, BiPredicate)
      */
-    void advanceOrThrow(final RuntimeException attemptException, final BinaryOperator<Throwable> exceptionTransformer,
+    void advanceOrThrow(final RuntimeException attemptException, final BinaryOperator<Throwable> onAttemptFailureOperator,
             final BiPredicate<RetryState, Throwable> retryPredicate) throws RuntimeException {
         try {
-            doAdvanceOrThrow(attemptException, exceptionTransformer, retryPredicate, true);
+            doAdvanceOrThrow(attemptException, onAttemptFailureOperator, retryPredicate, true);
         } catch (RuntimeException | Error unchecked) {
             throw unchecked;
         } catch (Throwable checked) {
@@ -129,18 +130,19 @@ void advanceOrThrow(final RuntimeException attemptException, final BinaryOperato
      *
      * @see #advanceOrThrow(RuntimeException, BinaryOperator, BiPredicate)
      */
-    void advanceOrThrow(final Throwable attemptException, final BinaryOperator<Throwable> exceptionTransformer,
+    void advanceOrThrow(final Throwable attemptException, final BinaryOperator<Throwable> onAttemptFailureOperator,
             final BiPredicate<RetryState, Throwable> retryPredicate) throws Throwable {
-        doAdvanceOrThrow(attemptException, exceptionTransformer, retryPredicate, false);
+        doAdvanceOrThrow(attemptException, onAttemptFailureOperator, retryPredicate, false);
     }
 
     /**
      * @param onlyRuntimeExceptions {@code true} iff the method must expect {@link #exception} and {@code attemptException} to be
      * {@link RuntimeException}s and must not explicitly handle other {@link Throwable} types, of which only {@link Error} is possible
      * as {@link RetryState} does not have any source of {@link Exception}s.
+     * @param onAttemptFailureOperator See {@link #advanceOrThrow(RuntimeException, BinaryOperator, BiPredicate)}.
      */
     private void doAdvanceOrThrow(final Throwable attemptException,
-            final BinaryOperator<Throwable> exceptionTransformer,
+            final BinaryOperator<Throwable> onAttemptFailureOperator,
             final BiPredicate<RetryState, Throwable> retryPredicate,
             final boolean onlyRuntimeExceptions) throws Throwable {
         assertTrue(attempt() < attempts);
@@ -149,7 +151,7 @@ private void doAdvanceOrThrow(final Throwable attemptException,
             assertTrue(isRuntime(attemptException));
         }
         assertTrue(!isFirstAttempt() || exception == null);
-        Throwable newlyChosenException = transformException(exception, attemptException, onlyRuntimeExceptions, exceptionTransformer);
+        Throwable newlyChosenException = callOnAttemptFailureOperator(exception, attemptException, onlyRuntimeExceptions, onAttemptFailureOperator);
         if (isLastAttempt()) {
             exception = newlyChosenException;
             throw exception;
@@ -167,27 +169,31 @@ private void doAdvanceOrThrow(final Throwable attemptException,
 
     /**
      * @param onlyRuntimeExceptions See {@link #doAdvanceOrThrow(Throwable, BinaryOperator, BiPredicate, boolean)}.
+     * @param onAttemptFailureOperator See {@link #advanceOrThrow(RuntimeException, BinaryOperator, BiPredicate)}.
      */
-    private static Throwable transformException(@Nullable final Throwable previouslyChosenException, final Throwable attemptException,
-            final boolean onlyRuntimeExceptions, final BinaryOperator<Throwable> exceptionTransformer) {
+    private static Throwable callOnAttemptFailureOperator(
+            @Nullable final Throwable previouslyChosenException,
+            final Throwable attemptException,
+            final boolean onlyRuntimeExceptions,
+            final BinaryOperator<Throwable> onAttemptFailureOperator) {
         if (onlyRuntimeExceptions && previouslyChosenException != null) {
             assertTrue(isRuntime(previouslyChosenException));
         }
         Throwable result;
         try {
-            result = assertNotNull(exceptionTransformer.apply(previouslyChosenException, attemptException));
+            result = assertNotNull(onAttemptFailureOperator.apply(previouslyChosenException, attemptException));
             if (onlyRuntimeExceptions) {
                 assertTrue(isRuntime(result));
             }
-        } catch (Throwable exceptionTransformerException) {
-            if (onlyRuntimeExceptions && !isRuntime(exceptionTransformerException)) {
-                throw exceptionTransformerException;
+        } catch (Throwable onAttemptFailureOperatorException) {
+            if (onlyRuntimeExceptions && !isRuntime(onAttemptFailureOperatorException)) {
+                throw onAttemptFailureOperatorException;
             }
             if (previouslyChosenException != null) {
-                exceptionTransformerException.addSuppressed(previouslyChosenException);
+                onAttemptFailureOperatorException.addSuppressed(previouslyChosenException);
             }
-            exceptionTransformerException.addSuppressed(attemptException);
-            throw exceptionTransformerException;
+            onAttemptFailureOperatorException.addSuppressed(attemptException);
+            throw onAttemptFailureOperatorException;
         }
         return result;
     }

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

@@ -41,31 +41,34 @@
 public final class RetryingAsyncCallbackSupplier<R> implements AsyncCallbackSupplier<R> {
     private final RetryState state;
     private final BiPredicate<RetryState, Throwable> retryPredicate;
-    private final BinaryOperator<Throwable> failedResultTransformer;
+    private final BinaryOperator<Throwable> onAttemptFailureOperator;
     private final AsyncCallbackSupplier<R> asyncFunction;
 
     /**
      * @param state The {@link RetryState} to be deemed as initial for the purpose of the new {@link RetryingAsyncCallbackSupplier}.
-     * @param failedResultTransformer A function that chooses which failed result of the {@code asyncFunction} to preserve as a prospective
-     * failed result of this {@link RetryingAsyncCallbackSupplier} and may also transform or mutate the exceptions.
-     * The choice is between
+     * @param onAttemptFailureOperator The action that is called once per failed attempt before (in the happens-before order) the
+     * {@code retryPredicate}, regardless of whether the {@code retryPredicate} is called.
+     * This action is allowed to have side effects.
+     * <p>
+     * It also has to choose which exception to preserve as a prospective failed result of this {@link RetryingAsyncCallbackSupplier}.
+     * The {@code onAttemptFailureOperator} may mutate its arguments, choose from the arguments, or return a different exception,
+     * but it must return a {@code @}{@link NonNull} value.
+     * The choice is between</p>
      * <ul>
      *     <li>the previously chosen failed result or {@code null} if none has been chosen
-     *     (the first argument of the {@code failedResultTransformer})</li>
-     *     <li>and the failed result from the most recent attempt (the second argument of the {@code failedResultTransformer}).</li>
+     *     (the first argument of the {@code onAttemptFailureOperator})</li>
+     *     <li>and the failed result from the most recent attempt (the second argument of the {@code onAttemptFailureOperator}).</li>
      * </ul>
-     * The {@code failedResultTransformer} may either choose from its arguments, or return a different exception, a.k.a. transform,
-     * but it must return a {@code @}{@link NonNull} value.
-     * If it completes abruptly, then the {@code asyncFunction} cannot be retried and the exception thrown by
-     * the {@code failedResultTransformer} is used as a failed result of this {@link RetryingAsyncCallbackSupplier}.
-     * The {@code failedResultTransformer} is called before (in the happens-before order) the {@code retryPredicate}.
-     * The result of the {@code failedResultTransformer} does not affect what exception is passed to the {@code retryPredicate}.
+     * The result of the {@code onAttemptFailureOperator} does not affect the exception passed to the {@code retryPredicate}.
+     * <p>
+     * If {@code onAttemptFailureOperator} completes abruptly, then the {@code asyncFunction} cannot be retried and the exception thrown by
+     * the {@code onAttemptFailureOperator} is used as a failed result of this {@link RetryingAsyncCallbackSupplier}.</p>
      * @param retryPredicate {@code true} iff another attempt needs to be made. If it completes abruptly,
      * then the {@code asyncFunction} cannot be retried and the exception thrown by the {@code retryPredicate}
      * is used as a failed result of this {@link RetryingAsyncCallbackSupplier}. The {@code retryPredicate} is called not more than once
      * per attempt and only if all the following is true:
      * <ul>
-     *     <li>{@code failedResultTransformer} completed normally;</li>
+     *     <li>{@code onAttemptFailureOperator} completed normally;</li>
      *     <li>the most recent attempt is not the {@linkplain RetryState#isLastAttempt() last} one.</li>
      * </ul>
      * The {@code retryPredicate} accepts this {@link RetryState} and the exception from the most recent attempt,
@@ -75,12 +78,12 @@ public final class RetryingAsyncCallbackSupplier<R> implements AsyncCallbackSupp
      */
     public RetryingAsyncCallbackSupplier(
             final RetryState state,
-            final BinaryOperator<Throwable> failedResultTransformer,
+            final BinaryOperator<Throwable> onAttemptFailureOperator,
             final BiPredicate<RetryState, Throwable> retryPredicate,
             final AsyncCallbackSupplier<R> asyncFunction) {
         this.state = state;
         this.retryPredicate = retryPredicate;
-        this.failedResultTransformer = failedResultTransformer;
+        this.onAttemptFailureOperator = onAttemptFailureOperator;
         this.asyncFunction = asyncFunction;
     }
 
@@ -113,7 +116,7 @@ private class RetryingCallback implements SingleResultCallback<R> {
         public void onResult(@Nullable final R result, @Nullable final Throwable t) {
             if (t != null) {
                 try {
-                    state.advanceOrThrow(t, failedResultTransformer, retryPredicate);
+                    state.advanceOrThrow(t, onAttemptFailureOperator, retryPredicate);
                 } catch (Throwable failedResult) {
                     wrapped.onResult(null, failedResult);
                     return;

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

@@ -37,26 +37,26 @@
 public final class RetryingSyncSupplier<R> implements Supplier<R> {
     private final RetryState state;
     private final BiPredicate<RetryState, Throwable> retryPredicate;
-    private final BinaryOperator<Throwable> failedResultTransformer;
+    private final BinaryOperator<Throwable> onAttemptFailureOperator;
     private final Supplier<R> syncFunction;
 
     /**
      * See {@link RetryingAsyncCallbackSupplier#RetryingAsyncCallbackSupplier(RetryState, BinaryOperator, BiPredicate, AsyncCallbackSupplier)}
      * for the documentation of the parameters.
      *
-     * @param failedResultTransformer Even though the {@code failedResultTransformer} accepts {@link Throwable},
+     * @param onAttemptFailureOperator Even though the {@code onAttemptFailureOperator} accepts {@link Throwable},
      * only {@link RuntimeException}s are passed to it.
      * @param retryPredicate Even though the {@code retryPredicate} accepts {@link Throwable},
      * only {@link RuntimeException}s are passed to it.
      */
     public RetryingSyncSupplier(
             final RetryState state,
-            final BinaryOperator<Throwable> failedResultTransformer,
+            final BinaryOperator<Throwable> onAttemptFailureOperator,
             final BiPredicate<RetryState, Throwable> retryPredicate,
             final Supplier<R> syncFunction) {
         this.state = state;
         this.retryPredicate = retryPredicate;
-        this.failedResultTransformer = failedResultTransformer;
+        this.onAttemptFailureOperator = onAttemptFailureOperator;
         this.syncFunction = syncFunction;
     }
 
@@ -66,10 +66,10 @@ public R get() {
             try {
                 return syncFunction.get();
             } catch (RuntimeException attemptException) {
-                state.advanceOrThrow(attemptException, failedResultTransformer, retryPredicate);
+                state.advanceOrThrow(attemptException, onAttemptFailureOperator, retryPredicate);
             } catch (Exception attemptException) {
                 // wrap potential sneaky / Kotlin exceptions
-                state.advanceOrThrow(new RuntimeException(attemptException), failedResultTransformer, retryPredicate);
+                state.advanceOrThrow(new RuntimeException(attemptException), onAttemptFailureOperator, retryPredicate);
             }
         }
     }

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

@@ -31,9 +31,11 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -123,14 +125,13 @@ public void close() {
     }
 
     @Override
-    public ClusterableServer getServer(final ServerAddress serverAddress) {
+    public ServersSnapshot getServersSnapshot() {
         isTrue("is open", !isClosed());
-
-        ServerTuple serverTuple = addressToServerTupleMap.get(serverAddress);
-        if (serverTuple == null) {
-            return null;
-        }
-        return serverTuple.server;
+        Map<ServerAddress, ServerTuple> nonAtomicSnapshot = new HashMap<>(addressToServerTupleMap);
+        return serverAddress -> {
+            ServerTuple serverTuple = nonAtomicSnapshot.get(serverAddress);
+            return serverTuple == null ? null : serverTuple.server;
+        };
     }
 
     void onChange(final Collection<ServerAddress> newHosts) {