Add support for server selection's deprioritized servers to all topologies.

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

@@ -113,9 +113,9 @@ abstract class BaseCluster implements Cluster {
     private volatile ClusterDescription description;
 
     BaseCluster(final ClusterId clusterId,
-                final ClusterSettings settings,
-                final ClusterableServerFactory serverFactory,
-                final ClientMetadata clientMetadata) {
+                          final ClusterSettings settings,
+                          final ClusterableServerFactory serverFactory,
+                          final ClientMetadata clientMetadata) {
         this.clusterId = notNull("clusterId", clusterId);
         this.settings = notNull("settings", settings);
         this.serverFactory = notNull("serverFactory", serverFactory);
@@ -159,7 +159,7 @@ public ServerTuple selectServer(final ServerSelector serverSelector, final Opera
             if (serverTuple != null) {
                 ServerAddress serverAddress = serverTuple.getServerDescription().getAddress();
                 logServerSelectionSucceeded(operationContext, clusterId, serverAddress, serverSelector, currentDescription);
-                serverDeprioritization.updateCandidate(serverAddress);
+                serverDeprioritization.updateCandidate(serverAddress, currentDescription.getType());
                 return serverTuple;
             }
             computedServerSelectionTimeout.onExpired(() ->
@@ -302,7 +302,7 @@ private boolean handleServerSelectionRequest(
                 if (serverTuple != null) {
                     ServerAddress serverAddress = serverTuple.getServerDescription().getAddress();
                     logServerSelectionSucceeded(operationContext, clusterId, serverAddress, request.originalSelector, description);
-                    serverDeprioritization.updateCandidate(serverAddress);
+                    serverDeprioritization.updateCandidate(serverAddress, description.getType());
                     request.onResult(serverTuple, null);
                     return true;
                 }
@@ -361,8 +361,7 @@ private static ServerSelector getCompleteServerSelector(
             final ClusterSettings settings) {
         List<ServerSelector> selectors = Stream.of(
                 getRaceConditionPreFilteringSelector(serversSnapshot),
-                serverSelector,
-                serverDeprioritization.getServerSelector(),
+                serverDeprioritization.apply(serverSelector),
                 settings.getServerSelector(), // may be null
                 new LatencyMinimizingServerSelector(settings.getLocalThreshold(MILLISECONDS), MILLISECONDS),
                 AtMostTwoRandomServerSelector.instance(),

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

@@ -17,6 +17,7 @@
 
 import com.mongodb.Function;
 import com.mongodb.MongoConnectionPoolClearedException;
+import com.mongodb.MongoException;
 import com.mongodb.ReadConcern;
 import com.mongodb.RequestContext;
 import com.mongodb.ServerAddress;
@@ -27,7 +28,6 @@
 import com.mongodb.internal.IgnorableRequestContext;
 import com.mongodb.internal.TimeoutContext;
 import com.mongodb.internal.TimeoutSettings;
-import com.mongodb.internal.VisibleForTesting;
 import com.mongodb.internal.observability.micrometer.Span;
 import com.mongodb.internal.observability.micrometer.TracingManager;
 import com.mongodb.internal.session.SessionContext;
@@ -40,6 +40,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
+import static com.mongodb.MongoException.SYSTEM_OVERLOADED_ERROR_LABEL;
 import static java.util.stream.Collectors.toList;
 
 /**
@@ -76,7 +77,7 @@ public OperationContext(final RequestContext requestContext, final SessionContex
                 null);
     }
 
-    public static OperationContext simpleOperationContext(
+    static OperationContext simpleOperationContext(
             final TimeoutSettings timeoutSettings, @Nullable final ServerApi serverApi) {
         return new OperationContext(
                 IgnorableRequestContext.INSTANCE,
@@ -113,6 +114,15 @@ public OperationContext withOperationName(final String operationName) {
                 operationName, tracingSpan);
     }
 
+    /**
+     * TODO-JAVA-6058: This method enables overriding the ServerDeprioritization state.
+     * It is a temporary solution to handle cases where deprioritization state persists across operations.
+     */
+    public OperationContext withNewServerDeprioritization() {
+        return new OperationContext(id, requestContext, sessionContext, timeoutContext, new ServerDeprioritization(), tracingManager, serverApi,
+                operationName, tracingSpan);
+    }
+
     public long getId() {
         return id;
     }
@@ -152,8 +162,7 @@ public void setTracingSpan(final Span tracingSpan) {
         this.tracingSpan = tracingSpan;
     }
 
-    @VisibleForTesting(otherwise = VisibleForTesting.AccessModifier.PRIVATE)
-    public OperationContext(final long id,
+    private OperationContext(final long id,
             final RequestContext requestContext,
             final SessionContext sessionContext,
             final TimeoutContext timeoutContext,
@@ -174,26 +183,6 @@ public OperationContext(final long id,
         this.tracingSpan = tracingSpan;
     }
 
-    @VisibleForTesting(otherwise = VisibleForTesting.AccessModifier.PRIVATE)
-    public OperationContext(final long id,
-            final RequestContext requestContext,
-            final SessionContext sessionContext,
-            final TimeoutContext timeoutContext,
-            final TracingManager tracingManager,
-            @Nullable final ServerApi serverApi,
-            @Nullable final String operationName) {
-        this.id = id;
-        this.serverDeprioritization = new ServerDeprioritization();
-        this.requestContext = requestContext;
-        this.sessionContext = sessionContext;
-        this.timeoutContext = timeoutContext;
-        this.tracingManager = tracingManager;
-        this.serverApi = serverApi;
-        this.operationName = operationName;
-        this.tracingSpan = null;
-    }
-
-
     /**
      * @return The same {@link ServerDeprioritization} if called on the same {@link OperationContext}.
      */
@@ -227,34 +216,45 @@ public OperationContext withOverride(final TimeoutContextOverride timeoutContext
     public static final class ServerDeprioritization {
         @Nullable
         private ServerAddress candidate;
+        @Nullable
+        private ClusterType clusterType;
         private final Set<ServerAddress> deprioritized;
-        private final DeprioritizingSelector selector;
 
         private ServerDeprioritization() {
             candidate = null;
             deprioritized = new HashSet<>();
-            selector = new DeprioritizingSelector();
+            clusterType = null;
         }
 
         /**
-         * The returned {@link ServerSelector} tries to {@linkplain ServerSelector#select(ClusterDescription) select}
-         * only the {@link ServerDescription}s that do not have deprioritized {@link ServerAddress}es.
-         * If no such {@link ServerDescription} can be selected, then it selects {@link ClusterDescription#getServerDescriptions()}.
+         * The returned {@link ServerSelector} wraps the provided selector and attempts
+         * {@linkplain ServerSelector#select(ClusterDescription) server selection} in two passes:
+         * <ol>
+         *   <li>First pass: selects using the wrapped selector with only non-deprioritized {@link ServerDescription}s.</li>
+         *   <li>Second pass: if the first pass selects no {@link ServerDescription}s,
+         *   selects using the wrapped selector again with all {@link ServerDescription}s, including deprioritized ones.</li>
+         * </ol>
          */
-        ServerSelector getServerSelector() {
-            return selector;
+        ServerSelector apply(final ServerSelector wrappedSelector) {
+            return new DeprioritizingSelector(wrappedSelector);
         }
 
-        void updateCandidate(final ServerAddress serverAddress) {
-            candidate = serverAddress;
+        void updateCandidate(final ServerAddress serverAddress, final ClusterType clusterType) {
+            this.candidate = serverAddress;
+            this.clusterType = clusterType;
         }
 
         public void onAttemptFailure(final Throwable failure) {
             if (candidate == null || failure instanceof MongoConnectionPoolClearedException) {
                 candidate = null;
                 return;
             }
-            deprioritized.add(candidate);
+
+            boolean isSystemOverloadedError = failure instanceof MongoException
+                    && ((MongoException) failure).hasErrorLabel(SYSTEM_OVERLOADED_ERROR_LABEL);
+            if (clusterType == ClusterType.SHARDED || isSystemOverloadedError) {
+                deprioritized.add(candidate);
+            }
         }
 
         /**
@@ -263,24 +263,41 @@ public void onAttemptFailure(final Throwable failure) {
          * which indeed may be used concurrently. {@link DeprioritizingSelector} does not need to be thread-safe.
          */
         private final class DeprioritizingSelector implements ServerSelector {
-            private DeprioritizingSelector() {
+            private final ServerSelector wrappedSelector;
+
+            private DeprioritizingSelector(final ServerSelector wrappedSelector) {
+                this.wrappedSelector = wrappedSelector;
             }
 
             @Override
             public List<ServerDescription> select(final ClusterDescription clusterDescription) {
                 List<ServerDescription> serverDescriptions = clusterDescription.getServerDescriptions();
-                if (!isEnabled(clusterDescription.getType())) {
-                    return serverDescriptions;
+
+                // TODO-JAVA-5908: Evaluate whether using the early-return optimization has a meaningful performance impact on server selection.
+                if (serverDescriptions.size() == 1 || deprioritized.isEmpty()) {
+                    return wrappedSelector.select(clusterDescription);
                 }
+
+                // TODO-JAVA-5908: Evaluate whether using a loop instead of Stream has a meaningful performance impact on server selection.
                 List<ServerDescription> nonDeprioritizedServerDescriptions = serverDescriptions
                         .stream()
                         .filter(serverDescription -> !deprioritized.contains(serverDescription.getAddress()))
                         .collect(toList());
-                return nonDeprioritizedServerDescriptions.isEmpty() ? serverDescriptions : nonDeprioritizedServerDescriptions;
-            }
 
-            private boolean isEnabled(final ClusterType clusterType) {
-                return clusterType == ClusterType.SHARDED;
+                // TODO-JAVA-5908: Evaluate whether using the early-return optimization has a meaningful performance impact on server selection.
+                if (nonDeprioritizedServerDescriptions.isEmpty()) {
+                    return wrappedSelector.select(clusterDescription);
+                }
+
+                List<ServerDescription> selected = wrappedSelector.select(
+                        new ClusterDescription(
+                                clusterDescription.getConnectionMode(),
+                                clusterDescription.getType(),
+                                clusterDescription.getSrvResolutionException(),
+                                nonDeprioritizedServerDescriptions,
+                                clusterDescription.getClusterSettings(),
+                                clusterDescription.getServerSettings()));
+                return selected.isEmpty() ? wrappedSelector.select(clusterDescription) : selected;
             }
         }
     }

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

@@ -73,7 +73,9 @@ final class AsyncChangeStreamBatchCursor<T> implements AsyncAggregateResponseBat
         this.wrapped = new AtomicReference<>(assertNotNull(wrapped));
         this.binding = binding;
         binding.retain();
-        this.initialOperationContext = operationContext.withOverride(TimeoutContext::withMaxTimeAsMaxAwaitTimeOverride);
+        this.initialOperationContext = operationContext
+                .withOverride(TimeoutContext::withMaxTimeAsMaxAwaitTimeOverride)
+                .withNewServerDeprioritization();
         this.resumeToken = resumeToken;
         this.maxWireVersion = maxWireVersion;
         isClosed = new AtomicBoolean();

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

@@ -85,7 +85,9 @@ final class ChangeStreamBatchCursor<T> implements AggregateResponseBatchCursor<T
                             final int maxWireVersion) {
         this.changeStreamOperation = changeStreamOperation;
         this.binding = binding.retain();
-        this.initialOperationContext = operationContext.withOverride(TimeoutContext::withMaxTimeAsMaxAwaitTimeOverride);
+        this.initialOperationContext = operationContext
+                .withOverride(TimeoutContext::withMaxTimeAsMaxAwaitTimeOverride)
+                .withNewServerDeprioritization();
         this.wrapped = wrapped;
         this.resumeToken = resumeToken;
         this.maxWireVersion = maxWireVersion;