Direct retries to another mongos if one is available

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

@@ -33,6 +33,7 @@
 import com.mongodb.event.ClusterOpeningEvent;
 import com.mongodb.internal.VisibleForTesting;
 import com.mongodb.internal.async.SingleResultCallback;
+import com.mongodb.internal.connection.OperationContext.ServerDeprioritization;
 import com.mongodb.internal.diagnostics.logging.Logger;
 import com.mongodb.internal.diagnostics.logging.Loggers;
 import com.mongodb.internal.logging.LogMessage;
@@ -122,8 +123,9 @@ public ServerTuple selectServer(final ServerSelector serverSelector, final Opera
             CountDownLatch currentPhase = phase.get();
             ClusterDescription curDescription = description;
             logServerSelectionStarted(clusterId, operationContext, serverSelector, curDescription);
-            ServerSelector compositeServerSelector = getCompositeServerSelector(serverSelector);
-            ServerTuple serverTuple = selectServer(compositeServerSelector, curDescription);
+            ServerDeprioritization serverDeprioritization = operationContext.getServerDeprioritization();
+            ServerSelector completeServerSelector = getCompleteServerSelector(serverSelector, serverDeprioritization);
+            ServerTuple serverTuple = selectServer(completeServerSelector, curDescription);
 
             boolean selectionWaitingLogged = false;
 
@@ -137,8 +139,10 @@ public ServerTuple selectServer(final ServerSelector serverSelector, final Opera
                 }
 
                 if (serverTuple != null) {
+                    ServerAddress serverAddress = serverTuple.getServerDescription().getAddress();
                     logServerSelectionSucceeded(
-                            clusterId, operationContext, serverTuple.getServerDescription().getAddress(), serverSelector, curDescription);
+                            clusterId, operationContext, serverAddress, serverSelector, curDescription);
+                    serverDeprioritization.updateCandidate(serverAddress, curDescription.getType());
                     return serverTuple;
                 }
 
@@ -163,7 +167,7 @@ public ServerTuple selectServer(final ServerSelector serverSelector, final Opera
 
                 currentPhase = phase.get();
                 curDescription = description;
-                serverTuple = selectServer(compositeServerSelector, curDescription);
+                serverTuple = selectServer(completeServerSelector, curDescription);
             }
 
         } catch (InterruptedException e) {
@@ -180,8 +184,9 @@ public void selectServerAsync(final ServerSelector serverSelector, final Operati
         ClusterDescription currentDescription = description;
 
         logServerSelectionStarted(clusterId, operationContext, serverSelector, currentDescription);
-        ServerSelectionRequest request = new ServerSelectionRequest(operationContext, serverSelector, getCompositeServerSelector(serverSelector),
-                                                                    getMaxWaitTimeNanos(), callback);
+        ServerSelectionRequest request = new ServerSelectionRequest(operationContext, serverSelector,
+                getCompleteServerSelector(serverSelector, operationContext.getServerDeprioritization()),
+                getMaxWaitTimeNanos(), callback);
 
         if (!handleServerSelectionRequest(request, currentPhase, currentDescription)) {
             notifyWaitQueueHandler(request);
@@ -276,10 +281,12 @@ private boolean handleServerSelectionRequest(final ServerSelectionRequest reques
                     return true;
                 }
 
-                ServerTuple serverTuple = selectServer(request.compositeSelector, description);
+                ServerTuple serverTuple = selectServer(request.completeSelector, description);
                 if (serverTuple != null) {
-                    logServerSelectionSucceeded(clusterId, request.operationContext, serverTuple.getServerDescription().getAddress(),
+                    ServerAddress serverAddress = serverTuple.getServerDescription().getAddress();
+                    logServerSelectionSucceeded(clusterId, request.operationContext, serverAddress,
                             request.originalSelector, description);
+                    request.operationContext.getServerDeprioritization().updateCandidate(serverAddress, description.getType());
                     request.onResult(serverTuple, null);
                     return true;
                 }
@@ -343,14 +350,13 @@ private static List<ServerTuple> atMostNRandom(final ArrayList<ServerDescription
         return result;
     }
 
-    private ServerSelector getCompositeServerSelector(final ServerSelector serverSelector) {
+    private ServerSelector getCompleteServerSelector(final ServerSelector serverSelector, final ServerDeprioritization serverDeprioritization) {
         ServerSelector latencyMinimizingServerSelector =
                 new LatencyMinimizingServerSelector(settings.getLocalThreshold(MILLISECONDS), MILLISECONDS);
-        if (settings.getServerSelector() == null) {
-            return new CompositeServerSelector(asList(serverSelector, latencyMinimizingServerSelector));
-        } else {
-            return new CompositeServerSelector(asList(serverSelector, settings.getServerSelector(), latencyMinimizingServerSelector));
-        }
+        CompositeServerSelector compositeSelector = settings.getServerSelector() == null
+                ? new CompositeServerSelector(asList(serverSelector, latencyMinimizingServerSelector))
+                : new CompositeServerSelector(asList(serverSelector, settings.getServerSelector(), latencyMinimizingServerSelector));
+        return serverDeprioritization.apply(compositeSelector);
     }
 
     protected ClusterableServer createServer(final ServerAddress serverAddress) {
@@ -399,21 +405,21 @@ private MongoException createAndLogTimeoutException(
     private static final class ServerSelectionRequest {
         private final OperationContext operationContext;
         private final ServerSelector originalSelector;
-        private final ServerSelector compositeSelector;
+        private final ServerSelector completeSelector;
         @Nullable
         private final Long maxWaitTimeNanos;
         private final SingleResultCallback<ServerTuple> callback;
         private final long startTimeNanos = System.nanoTime();
         private CountDownLatch phase;
 
         ServerSelectionRequest(final OperationContext operationContext,
-                               final ServerSelector serverSelector, final ServerSelector compositeSelector,
+                               final ServerSelector serverSelector, final ServerSelector completeSelector,
                                @Nullable
                                final Long maxWaitTimeNanos,
                                final SingleResultCallback<ServerTuple> callback) {
             this.operationContext = operationContext;
             this.originalSelector = serverSelector;
-            this.compositeSelector = compositeSelector;
+            this.completeSelector = completeSelector;
             this.maxWaitTimeNanos = maxWaitTimeNanos;
             this.callback = callback;
         }

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

@@ -15,6 +15,17 @@
  */
 package com.mongodb.internal.connection;
 
+import com.mongodb.MongoConnectionPoolClearedException;
+import com.mongodb.ServerAddress;
+import com.mongodb.connection.ClusterDescription;
+import com.mongodb.connection.ClusterType;
+import com.mongodb.connection.ServerDescription;
+import com.mongodb.lang.Nullable;
+import com.mongodb.selector.ServerSelector;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
@@ -23,12 +34,88 @@
 public class OperationContext {
     private static final AtomicLong NEXT_ID = new AtomicLong(0);
     private final long id;
+    private final ServerDeprioritization serverDeprioritization;
 
     public OperationContext() {
         id = NEXT_ID.incrementAndGet();
+        serverDeprioritization = new ServerDeprioritization();
     }
 
     public long getId() {
         return id;
     }
+
+    /**
+     * @return The same {@link ServerDeprioritization} if called on the same {@link OperationContext}.
+     */
+    public ServerDeprioritization getServerDeprioritization() {
+        return serverDeprioritization;
+    }
+
+    public static final class ServerDeprioritization {
+        @Nullable
+        private ServerAddress candidate;
+        private final Set<ServerAddress> deprioritized;
+
+        private ServerDeprioritization() {
+            candidate = null;
+            deprioritized = new HashSet<>();
+        }
+
+        ServerSelector apply(final ServerSelector selector) {
+            return new DeprioritizingSelector(selector);
+        }
+
+        void updateCandidate(final ServerAddress serverAddress, final ClusterType clusterType) {
+            candidate = isEnabled(clusterType) ? serverAddress : null;
+        }
+
+        public void onAttemptFailure(final Throwable failure) {
+            if (candidate == null || failure instanceof MongoConnectionPoolClearedException) {
+                candidate = null;
+                return;
+            }
+            deprioritized.add(candidate);
+        }
+
+        private static boolean isEnabled(final ClusterType clusterType) {
+            return clusterType == ClusterType.SHARDED;
+        }
+
+        /**
+         * {@link ServerSelector} requires thread safety, but that is only because a user may specify
+         * {@link com.mongodb.connection.ClusterSettings.Builder#serverSelector(ServerSelector)},
+         * which indeed may be used concurrently. {@link DeprioritizingSelector} does not need to be thread-safe.
+         */
+        private final class DeprioritizingSelector implements ServerSelector {
+            private final ServerSelector wrapped;
+
+            private DeprioritizingSelector(final ServerSelector wrapped) {
+                this.wrapped = wrapped;
+            }
+
+            @Override
+            public List<ServerDescription> select(final ClusterDescription clusterDescription) {
+                if (isEnabled(clusterDescription.getType())) {
+                    List<ServerDescription> filteredServerDescriptions = ClusterDescriptionHelper.getServersByPredicate(
+                            clusterDescription, serverDescription -> !deprioritized.contains(serverDescription.getAddress()));
+                    ClusterDescription filteredClusterDescription = new ClusterDescription(
+                            clusterDescription.getConnectionMode(),
+                            clusterDescription.getType(),
+                            clusterDescription.getSrvResolutionException(),
+                            filteredServerDescriptions,
+                            clusterDescription.getClusterSettings(),
+                            clusterDescription.getServerSettings());
+                    List<ServerDescription> result = wrapped.select(filteredClusterDescription);
+                    if (result.isEmpty()) {
+                        // fall back to selecting from all servers ignoring the deprioritized ones
+                        result = wrapped.select(clusterDescription);
+                    }
+                    return result;
+                } else {
+                    return wrapped.select(clusterDescription);
+                }
+            }
+        }
+    }
 }

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

@@ -44,6 +44,7 @@
 
 import java.util.Collections;
 import java.util.List;
+import java.util.function.BinaryOperator;
 
 import static com.mongodb.assertions.Assertions.assertNotNull;
 import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
@@ -285,7 +286,11 @@ static <D, T> void createReadCommandAndExecuteAsync(
 
     static <R> AsyncCallbackSupplier<R> decorateReadWithRetriesAsync(final RetryState retryState, final OperationContext operationContext,
             final AsyncCallbackSupplier<R> asyncReadFunction) {
-        return new RetryingAsyncCallbackSupplier<>(retryState, CommandOperationHelper::chooseRetryableReadException,
+        BinaryOperator<Throwable> onAttemptFailure =
+                (@Nullable Throwable previouslyChosenException, Throwable mostRecentAttemptException) ->
+                        CommandOperationHelper.onRetryableReadAttemptFailure(
+                                operationContext, previouslyChosenException, mostRecentAttemptException);
+        return new RetryingAsyncCallbackSupplier<>(retryState, onAttemptFailure,
                 CommandOperationHelper::shouldAttemptToRetryRead, callback -> {
             logRetryExecute(retryState, operationContext);
             asyncReadFunction.get(callback);
@@ -294,7 +299,11 @@ static <R> AsyncCallbackSupplier<R> decorateReadWithRetriesAsync(final RetryStat
 
     static <R> AsyncCallbackSupplier<R> decorateWriteWithRetriesAsync(final RetryState retryState, final OperationContext operationContext,
             final AsyncCallbackSupplier<R> asyncWriteFunction) {
-        return new RetryingAsyncCallbackSupplier<>(retryState, CommandOperationHelper::chooseRetryableWriteException,
+        BinaryOperator<Throwable> onAttemptFailure =
+                (@Nullable Throwable previouslyChosenException, Throwable mostRecentAttemptException) ->
+                        CommandOperationHelper.onRetryableWriteAttemptFailure(
+                                operationContext, previouslyChosenException, mostRecentAttemptException);
+        return new RetryingAsyncCallbackSupplier<>(retryState, onAttemptFailure,
                 CommandOperationHelper::shouldAttemptToRetryWrite, callback -> {
             logRetryExecute(retryState, operationContext);
             asyncWriteFunction.get(callback);

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

@@ -51,8 +51,15 @@ interface CommandCreator {
         BsonDocument create(ServerDescription serverDescription, ConnectionDescription connectionDescription);
     }
 
+    static Throwable onRetryableReadAttemptFailure(
+            final OperationContext operationContext,
+            @Nullable final Throwable previouslyChosenException,
+            final Throwable mostRecentAttemptException) {
+        operationContext.getServerDeprioritization().onAttemptFailure(mostRecentAttemptException);
+        return chooseRetryableReadException(previouslyChosenException, mostRecentAttemptException);
+    }
 
-    static Throwable chooseRetryableReadException(
+    private static Throwable chooseRetryableReadException(
             @Nullable final Throwable previouslyChosenException, final Throwable mostRecentAttemptException) {
         assertFalse(mostRecentAttemptException instanceof ResourceSupplierInternalException);
         if (previouslyChosenException == null
@@ -64,7 +71,15 @@ static Throwable chooseRetryableReadException(
         }
     }
 
-    static Throwable chooseRetryableWriteException(
+    static Throwable onRetryableWriteAttemptFailure(
+            final OperationContext operationContext,
+            @Nullable final Throwable previouslyChosenException,
+            final Throwable mostRecentAttemptException) {
+        operationContext.getServerDeprioritization().onAttemptFailure(mostRecentAttemptException);
+        return chooseRetryableWriteException(previouslyChosenException, mostRecentAttemptException);
+    }
+
+    private static Throwable chooseRetryableWriteException(
             @Nullable final Throwable previouslyChosenException, final Throwable mostRecentAttemptException) {
         if (previouslyChosenException == null) {
             if (mostRecentAttemptException instanceof ResourceSupplierInternalException) {

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

@@ -51,6 +51,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.BinaryOperator;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
@@ -140,7 +141,11 @@ public Boolean getRetryWrites() {
 
     private <R> Supplier<R> decorateWriteWithRetries(final RetryState retryState, final OperationContext operationContext,
             final Supplier<R> writeFunction) {
-        return new RetryingSyncSupplier<>(retryState, CommandOperationHelper::chooseRetryableWriteException,
+        BinaryOperator<Throwable> onAttemptFailure =
+                (@Nullable Throwable previouslyChosenException, Throwable mostRecentAttemptException) ->
+                        CommandOperationHelper.onRetryableWriteAttemptFailure(
+                                operationContext, previouslyChosenException, mostRecentAttemptException);
+        return new RetryingSyncSupplier<>(retryState, onAttemptFailure,
                 this::shouldAttemptToRetryWrite, () -> {
             logRetryExecute(retryState, operationContext);
             return writeFunction.get();
@@ -149,7 +154,11 @@ private <R> Supplier<R> decorateWriteWithRetries(final RetryState retryState, fi
 
     private <R> AsyncCallbackSupplier<R> decorateWriteWithRetries(final RetryState retryState, final OperationContext operationContext,
             final AsyncCallbackSupplier<R> writeFunction) {
-        return new RetryingAsyncCallbackSupplier<>(retryState, CommandOperationHelper::chooseRetryableWriteException,
+        BinaryOperator<Throwable> onAttemptFailure =
+                (@Nullable Throwable previouslyChosenException, Throwable mostRecentAttemptException) ->
+                        CommandOperationHelper.onRetryableWriteAttemptFailure(
+                                operationContext, previouslyChosenException, mostRecentAttemptException);
+        return new RetryingAsyncCallbackSupplier<>(retryState, onAttemptFailure,
                 this::shouldAttemptToRetryWrite, callback -> {
             logRetryExecute(retryState, operationContext);
             writeFunction.get(callback);

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

@@ -41,6 +41,7 @@
 import org.bson.codecs.Decoder;
 
 import java.util.function.BiFunction;
+import java.util.function.BinaryOperator;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
@@ -274,7 +275,11 @@ static <D, T> T createReadCommandAndExecute(
 
     static <R> Supplier<R> decorateWriteWithRetries(final RetryState retryState,
             final OperationContext operationContext, final Supplier<R> writeFunction) {
-        return new RetryingSyncSupplier<>(retryState, CommandOperationHelper::chooseRetryableWriteException,
+        BinaryOperator<Throwable> onAttemptFailure =
+                (@Nullable Throwable previouslyChosenException, Throwable mostRecentAttemptException) ->
+                        CommandOperationHelper.onRetryableWriteAttemptFailure(
+                                operationContext, previouslyChosenException, mostRecentAttemptException);
+        return new RetryingSyncSupplier<>(retryState, onAttemptFailure,
                 CommandOperationHelper::shouldAttemptToRetryWrite, () -> {
             logRetryExecute(retryState, operationContext);
             return writeFunction.get();
@@ -283,7 +288,11 @@ static <R> Supplier<R> decorateWriteWithRetries(final RetryState retryState,
 
     static <R> Supplier<R> decorateReadWithRetries(final RetryState retryState, final OperationContext operationContext,
             final Supplier<R> readFunction) {
-        return new RetryingSyncSupplier<>(retryState, CommandOperationHelper::chooseRetryableReadException,
+        BinaryOperator<Throwable> onAttemptFailure =
+                (@Nullable Throwable previouslyChosenException, Throwable mostRecentAttemptException) ->
+                        CommandOperationHelper.onRetryableReadAttemptFailure(
+                                operationContext, previouslyChosenException, mostRecentAttemptException);
+        return new RetryingSyncSupplier<>(retryState, onAttemptFailure,
                 CommandOperationHelper::shouldAttemptToRetryRead, () -> {
             logRetryExecute(retryState, operationContext);
             return readFunction.get();