Add enableOverloadRetargeting API

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

@@ -20,6 +20,7 @@
 import com.mongodb.annotations.Beta;
 import com.mongodb.annotations.Reason;
 import com.mongodb.connection.ClusterSettings;
+import com.mongodb.connection.ClusterType;
 import com.mongodb.connection.ConnectionPoolSettings;
 import com.mongodb.connection.ServerMonitoringMode;
 import com.mongodb.connection.ServerSettings;
@@ -276,6 +277,8 @@
  * <li>{@code maxAdaptiveRetries=n}: This is {@linkplain Beta Beta API}.
  * The maximum number of retry attempts when encountering a retryable overload error.
  * See {@link MongoClientSettings.Builder#maxAdaptiveRetries(Integer)} for more information.</li>
+ * <li>{@code enableOverloadRetargeting=true|false}: Whether to enable overload retargeting. Defaults to false.
+ *  See {@link MongoClientSettings.Builder#enableOverloadRetargeting(boolean)} for more information.</li>
  * <li>{@code uuidRepresentation=unspecified|standard|javaLegacy|csharpLegacy|pythonLegacy}.  See
  * {@link MongoClientSettings#getUuidRepresentation()} for documentation of semantics of this parameter.  Defaults to "javaLegacy", but
  * will change to "unspecified" in the next major release.</li>
@@ -313,6 +316,7 @@ public class ConnectionString {
     private Boolean retryWrites;
     private Boolean retryReads;
     private Integer maxAdaptiveRetries;
+    private Boolean enableOverloadRetargeting;
     private ReadConcern readConcern;
 
     private Integer minConnectionPoolSize;
@@ -564,6 +568,7 @@ public ConnectionString(final String connectionString, @Nullable final DnsClient
         GENERAL_OPTIONS_KEYS.add("retrywrites");
         GENERAL_OPTIONS_KEYS.add("retryreads");
         GENERAL_OPTIONS_KEYS.add("maxadaptiveretries");
+        GENERAL_OPTIONS_KEYS.add("enableoverloadretargeting");
 
         GENERAL_OPTIONS_KEYS.add("appname");
 
@@ -718,6 +723,9 @@ private void translateOptions(final Map<String, List<String>> optionsMap) {
                         throw new IllegalArgumentException("maxAdaptiveRetries must be >= 0");
                     }
                     break;
+                case "enableoverloadretargeting":
+                    enableOverloadRetargeting = parseBoolean(value, "enableoverloadretargeting");
+                    break;
                 case "uuidrepresentation":
                     uuidRepresentation = createUuidRepresentation(value);
                     break;
@@ -1511,6 +1519,20 @@ public Integer getMaxAdaptiveRetries() {
         return maxAdaptiveRetries;
     }
 
+    /**
+     * Gets whether overload retargeting is enabled.
+     * See {@link MongoClientSettings.Builder#enableOverloadRetargeting(boolean)} for more information.
+     *
+     * @return the enableOverloadRetargeting value, or null if not set
+     * @see MongoClientSettings.Builder#enableOverloadRetargeting(boolean)
+     * @since 5.7
+     */
+    @Beta(Reason.CLIENT)
+    @Nullable
+    public Boolean getEnableOverloadRetargeting() {
+        return enableOverloadRetargeting;
+    }
+
     /**
      * Gets the minimum connection pool size specified in the connection string.
      * @return the minimum connection pool size
@@ -1825,6 +1847,7 @@ public boolean equals(final Object o) {
                 && Objects.equals(retryWrites, that.retryWrites)
                 && Objects.equals(retryReads, that.retryReads)
                 && Objects.equals(maxAdaptiveRetries, that.maxAdaptiveRetries)
+                && Objects.equals(enableOverloadRetargeting, that.enableOverloadRetargeting)
                 && Objects.equals(readConcern, that.readConcern)
                 && Objects.equals(minConnectionPoolSize, that.minConnectionPoolSize)
                 && Objects.equals(maxConnectionPoolSize, that.maxConnectionPoolSize)
@@ -1856,7 +1879,7 @@ public boolean equals(final Object o) {
     @Override
     public int hashCode() {
         return Objects.hash(credential, isSrvProtocol, hosts, database, collection, directConnection, readPreference,
-                writeConcern, retryWrites, retryReads, maxAdaptiveRetries, readConcern, minConnectionPoolSize, maxConnectionPoolSize, maxWaitTime,
+                writeConcern, retryWrites, retryReads, maxAdaptiveRetries, enableOverloadRetargeting, readConcern, minConnectionPoolSize, maxConnectionPoolSize, maxWaitTime,
                 maxConnectionIdleTime, maxConnectionLifeTime, maxConnecting, connectTimeout, timeout, socketTimeout, sslEnabled,
                 sslInvalidHostnameAllowed, requiredReplicaSetName, serverSelectionTimeout, localThreshold, heartbeatFrequency,
                 serverMonitoringMode, applicationName, compressorList, uuidRepresentation, srvServiceName, srvMaxHosts, proxyHost,

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

@@ -25,6 +25,7 @@
 import com.mongodb.client.model.geojson.codecs.GeoJsonCodecProvider;
 import com.mongodb.client.model.mql.ExpressionCodecProvider;
 import com.mongodb.connection.ClusterSettings;
+import com.mongodb.connection.ClusterType;
 import com.mongodb.connection.ConnectionPoolSettings;
 import com.mongodb.connection.ServerSettings;
 import com.mongodb.connection.SocketSettings;
@@ -96,6 +97,7 @@ public final class MongoClientSettings {
     private final boolean retryReads;
     @Nullable
     private final Integer maxAdaptiveRetries;
+    private final boolean enableOverloadRetargeting;
     private final ReadConcern readConcern;
     private final MongoCredential credential;
     private final TransportSettings transportSettings;
@@ -219,6 +221,7 @@ public static final class Builder {
         private boolean retryReads = true;
         @Nullable
         private Integer maxAdaptiveRetries;
+        private boolean enableOverloadRetargeting = false;
         private ReadConcern readConcern = ReadConcern.DEFAULT;
         private CodecRegistry codecRegistry = MongoClientSettings.getDefaultCodecRegistry();
         private TransportSettings transportSettings;
@@ -261,6 +264,7 @@ private Builder(final MongoClientSettings settings) {
             retryWrites = settings.getRetryWrites();
             retryReads = settings.getRetryReads();
             maxAdaptiveRetries = settings.getMaxAdaptiveRetries();
+            enableOverloadRetargeting = settings.getEnableOverloadRetargeting();
             readConcern = settings.getReadConcern();
             credential = settings.getCredential();
             uuidRepresentation = settings.getUuidRepresentation();
@@ -323,6 +327,10 @@ public Builder applyConnectionString(final ConnectionString connectionString) {
             if (connectionString.getMaxAdaptiveRetries() != null) {
                 maxAdaptiveRetries = connectionString.getMaxAdaptiveRetries();
             }
+            Boolean enableOverloadRetargetingValue = connectionString.getEnableOverloadRetargeting();
+            if (enableOverloadRetargetingValue != null) {
+                enableOverloadRetargeting = enableOverloadRetargetingValue;
+            }
             if (connectionString.getUuidRepresentation() != null) {
                 uuidRepresentation = connectionString.getUuidRepresentation();
             }
@@ -559,6 +567,31 @@ public Builder maxAdaptiveRetries(@Nullable final Integer maxAdaptiveRetries) {
             return this;
         }
 
+        /**
+         * Sets whether to enable overload retargeting.
+         *
+         * <p>When enabled, the previously selected servers on which attempts failed with an error
+         * {@linkplain MongoException#hasErrorLabel(String) having}
+         * the {@value MongoException#SYSTEM_OVERLOADED_ERROR_LABEL} label may be deprioritized during
+         * server selection on subsequent retry attempts. This applies to reads when
+         * {@linkplain #retryReads(boolean) retryReads} is enabled, and to writes when
+         * {@linkplain #retryWrites(boolean) retryWrites} is enabled.</p>
+         *
+         * <p>This setting does not take effect for {@linkplain ClusterType#SHARDED sharded clusters}.</p>
+         *
+         * <p>Defaults to {@code false}.</p>
+         *
+         * @param enableOverloadRetargeting whether to enable overload retargeting.
+         * @return this
+         * @see #getEnableOverloadRetargeting()
+         * @since 5.7
+         */
+        @Beta(Reason.CLIENT)
+        public Builder enableOverloadRetargeting(final boolean enableOverloadRetargeting) {
+            this.enableOverloadRetargeting = enableOverloadRetargeting;
+            return this;
+        }
+
         /**
          * Sets the read concern.
          *
@@ -933,6 +966,19 @@ public Integer getMaxAdaptiveRetries() {
         return maxAdaptiveRetries;
     }
 
+    /**
+     * Returns whether overload retargeting is enabled.
+     * See {@link Builder#enableOverloadRetargeting(boolean)} for more information.
+     *
+     * @return the enableOverloadRetargeting value
+     * @see Builder#enableOverloadRetargeting(boolean)
+     * @since 5.7
+     */
+    @Beta(Reason.CLIENT)
+    public boolean getEnableOverloadRetargeting() {
+        return enableOverloadRetargeting;
+    }
+
     /**
      * The read concern to use.
      *
@@ -1207,6 +1253,7 @@ public boolean equals(final Object o) {
         return retryWrites == that.retryWrites
                 && retryReads == that.retryReads
                 && Objects.equals(maxAdaptiveRetries, that.maxAdaptiveRetries)
+                && enableOverloadRetargeting == that.enableOverloadRetargeting
                 && heartbeatSocketTimeoutSetExplicitly == that.heartbeatSocketTimeoutSetExplicitly
                 && heartbeatConnectTimeoutSetExplicitly == that.heartbeatConnectTimeoutSetExplicitly
                 && Objects.equals(readPreference, that.readPreference)
@@ -1236,7 +1283,8 @@ public boolean equals(final Object o) {
 
     @Override
     public int hashCode() {
-        return Objects.hash(readPreference, writeConcern, retryWrites, retryReads, maxAdaptiveRetries, readConcern, credential, transportSettings,
+        return Objects.hash(readPreference, writeConcern, retryWrites, retryReads, maxAdaptiveRetries, enableOverloadRetargeting, readConcern,
+                credential, transportSettings,
                 commandListeners, codecRegistry, loggerSettings, clusterSettings, socketSettings,
                 heartbeatSocketSettings, connectionPoolSettings, serverSettings, sslSettings, applicationName, compressorList,
                 uuidRepresentation, serverApi, autoEncryptionSettings, heartbeatSocketTimeoutSetExplicitly,
@@ -1252,6 +1300,7 @@ public String toString() {
                 + ", retryWrites=" + retryWrites
                 + ", retryReads=" + retryReads
                 + ", maxAdaptiveRetries=" + maxAdaptiveRetries
+                + ", enableOverloadRetargeting=" + enableOverloadRetargeting
                 + ", readConcern=" + readConcern
                 + ", credential=" + credential
                 + ", transportSettings=" + transportSettings
@@ -1281,8 +1330,9 @@ private MongoClientSettings(final Builder builder) {
         readPreference = builder.readPreference;
         writeConcern = builder.writeConcern;
         retryWrites = builder.retryWrites;
-        maxAdaptiveRetries = builder.maxAdaptiveRetries;
         retryReads = builder.retryReads;
+        maxAdaptiveRetries = builder.maxAdaptiveRetries;
+        enableOverloadRetargeting = builder.enableOverloadRetargeting;
         readConcern = builder.readConcern;
         credential = builder.credential;
         transportSettings = builder.transportSettings;

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

@@ -77,6 +77,18 @@ public OperationContext(final RequestContext requestContext, final SessionContex
                 null);
     }
 
+    public OperationContext(final RequestContext requestContext, final SessionContext sessionContext, final TimeoutContext timeoutContext,
+            final TracingManager tracingManager,
+            @Nullable final ServerApi serverApi,
+            @Nullable final String operationName,
+            final ServerDeprioritization serverDeprioritization) {
+        this(NEXT_ID.incrementAndGet(), requestContext, sessionContext, timeoutContext, serverDeprioritization,
+                tracingManager,
+                serverApi,
+                operationName,
+                null);
+    }
+
     static OperationContext simpleOperationContext(
             final TimeoutSettings timeoutSettings, @Nullable final ServerApi serverApi) {
         return new OperationContext(
@@ -119,7 +131,8 @@ public OperationContext withOperationName(final String operationName) {
      * 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,
+        return new OperationContext(id, requestContext, sessionContext, timeoutContext,
+                new ServerDeprioritization(serverDeprioritization.enableOverloadRetargeting), tracingManager, serverApi,
                 operationName, tracingSpan);
     }
 
@@ -206,7 +219,8 @@ public OperationContext withConnectionEstablishmentSessionContext() {
     }
 
     public OperationContext withMinRoundTripTime(final ServerDescription serverDescription) {
-        return withTimeoutContext(timeoutContext.withMinRoundTripTime(TimeUnit.NANOSECONDS.toMillis(serverDescription.getMinRoundTripTimeNanos())));
+        return withTimeoutContext(
+                timeoutContext.withMinRoundTripTime(TimeUnit.NANOSECONDS.toMillis(serverDescription.getMinRoundTripTimeNanos())));
     }
 
     public OperationContext withOverride(final TimeoutContextOverride timeoutContextOverrideFunction) {
@@ -219,11 +233,17 @@ public static final class ServerDeprioritization {
         @Nullable
         private ClusterType clusterType;
         private final Set<ServerAddress> deprioritized;
+        private final boolean enableOverloadRetargeting;
 
-        private ServerDeprioritization() {
-            candidate = null;
-            deprioritized = new HashSet<>();
-            clusterType = null;
+        public ServerDeprioritization() {
+            this(false);
+        }
+
+        public ServerDeprioritization(final boolean enableOverloadRetargeting) {
+            this.enableOverloadRetargeting = enableOverloadRetargeting;
+            this.candidate = null;
+            this.deprioritized = new HashSet<>();
+            this.clusterType = null;
         }
 
         /**
@@ -253,7 +273,8 @@ public void onAttemptFailure(final Throwable failure) {
             // As per spec: sharded clusters deprioritize on any error, other topologies only on overload
             boolean isSystemOverloadedError = failure instanceof MongoException
                     && ((MongoException) failure).hasErrorLabel(SYSTEM_OVERLOADED_ERROR_LABEL);
-            if (clusterType == ClusterType.SHARDED || isSystemOverloadedError) {
+
+            if (clusterType == ClusterType.SHARDED || (isSystemOverloadedError && enableOverloadRetargeting)) {
                 deprioritized.add(candidate);
             }
         }
@@ -303,6 +324,7 @@ public List<ServerDescription> select(final ClusterDescription clusterDescriptio
         }
     }
 
-    public interface TimeoutContextOverride extends Function<TimeoutContext, TimeoutContext> {}
+    public interface TimeoutContextOverride extends Function<TimeoutContext, TimeoutContext> {
+    }
 }
 

driver-core/src/test/unit/com/mongodb/AbstractConnectionStringTest.java

@@ -112,7 +112,7 @@ protected void testValidOptions() {
 
             if (option.getKey().equals("authmechanism")) {
                 String expected = option.getValue().asString().getValue();
-                if (expected.equals("MONGODB-CR"))  {
+                if (expected.equals("MONGODB-CR")) {
                     assertNotNull(connectionString.getCredential());
                     assertNull(connectionString.getCredential().getAuthenticationMechanism());
                 } else {
@@ -125,6 +125,9 @@ protected void testValidOptions() {
             } else if (option.getKey().equalsIgnoreCase("maxadaptiveretries")) {
                 int expected = option.getValue().asInt32().getValue();
                 assertEquals(expected, connectionString.getMaxAdaptiveRetries().intValue());
+            } else if (option.getKey().equalsIgnoreCase("enableoverloadretargeting")) {
+                boolean expected = option.getValue().asBoolean().getValue();
+                assertEquals(expected, connectionString.getEnableOverloadRetargeting().booleanValue());
             } else if (option.getKey().equalsIgnoreCase("replicaset")) {
                 String expected = option.getValue().asString().getValue();
                 assertEquals(expected, connectionString.getRequiredReplicaSetName());

driver-core/src/test/unit/com/mongodb/ConnectionStringUnitTest.java

@@ -46,7 +46,8 @@ void defaults() {
     @ParameterizedTest
     @ValueSource(strings = {
             "serverMonitoringMode=stream",
-            "maxAdaptiveRetries=42"
+            "maxAdaptiveRetries=42",
+            "enableOverloadRetargeting=true"
     })
     void equalAndHashCode(final String connectionStringOptions) {
         ConnectionString default1 = new ConnectionString(DEFAULT_OPTIONS);
@@ -129,4 +130,14 @@ void maxAdaptiveRetries() {
                         () -> new ConnectionString(DEFAULT_OPTIONS + "maxAdaptiveRetries=invalid"))
         );
     }
+
+    @Test
+    void enableOverloadRetargeting() {
+        assertAll(
+                () -> assertNull(new ConnectionString("mongodb://localhost/").getEnableOverloadRetargeting()),
+                () -> assertEquals(false, new ConnectionString(DEFAULT_OPTIONS + "enableOverloadRetargeting=false").getEnableOverloadRetargeting()),
+                () -> assertEquals(true, new ConnectionString(DEFAULT_OPTIONS + "enableOverloadRetargeting=true").getEnableOverloadRetargeting()),
+                () -> assertNull(new ConnectionString(DEFAULT_OPTIONS + "enableOverloadRetargeting=foos").getEnableOverloadRetargeting())
+        );
+    }
 }

driver-core/src/test/unit/com/mongodb/MongoClientSettingsSpecification.groovy

@@ -576,6 +576,7 @@ class MongoClientSettingsSpecification extends Specification {
         def actual = MongoClientSettings.Builder.declaredFields.grep {  !it.synthetic } *.name.sort()
         def expected = ['applicationName', 'autoEncryptionSettings', 'clusterSettingsBuilder', 'codecRegistry', 'commandListeners',
                         'compressorList', 'connectionPoolSettingsBuilder', 'contextProvider', 'credential', 'dnsClient',
+                        'enableOverloadRetargeting',
                         'heartbeatConnectTimeoutMS', 'heartbeatSocketTimeoutMS', 'inetAddressResolver', 'loggerSettingsBuilder',
                         'maxAdaptiveRetries', 'observabilitySettings',
                         'readConcern', 'readPreference', 'retryReads',
@@ -595,6 +596,7 @@ class MongoClientSettingsSpecification extends Specification {
                         'applyToConnectionPoolSettings', 'applyToLoggerSettings', 'applyToServerSettings', 'applyToSocketSettings',
                         'applyToSslSettings', 'autoEncryptionSettings', 'build', 'codecRegistry', 'commandListenerList',
                         'compressorList', 'contextProvider', 'credential', 'dnsClient',
+                        'enableOverloadRetargeting',
                         'heartbeatConnectTimeoutMS',
                         'heartbeatSocketTimeoutMS', 'inetAddressResolver', 'maxAdaptiveRetries', 'observabilitySettings', 'readConcern',
                         'readPreference',