JAVA-5949 preserve connection pool on backpressure errors when establishing connections

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

@@ -0,0 +1,112 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.internal.connection;
+
+import com.mongodb.MongoException;
+import com.mongodb.MongoSecurityException;
+import com.mongodb.MongoSocketException;
+
+import javax.net.ssl.SSLHandshakeException;
+import javax.net.ssl.SSLPeerUnverifiedException;
+import javax.net.ssl.SSLProtocolException;
+import java.net.UnknownHostException;
+import java.security.cert.CertPathBuilderException;
+import java.security.cert.CertPathValidatorException;
+import java.security.cert.CertificateException;
+import java.util.Locale;
+
+/**
+ * Attaches {@link MongoException#SYSTEM_OVERLOADED_ERROR_LABEL} and
+ * {@link MongoException#RETRYABLE_ERROR_LABEL} to network errors encountered during connection
+ * establishment or the hello message, per the CMAP specification.
+ *
+ * <p>This is topology-agnostic: it must be invoked from the connection-establishment path so that
+ * both default SDAM and load-balanced modes are covered.
+ */
+final class BackpressureErrorLabeler {
+
+    private BackpressureErrorLabeler() {
+    }
+
+    static void applyLabelsIfEligible(final Throwable t) {
+        if (!(t instanceof MongoException)) {
+            return;
+        }
+        if (t instanceof MongoSecurityException) {
+            return;
+        }
+        if (!(t instanceof MongoSocketException)) {
+            return;
+        }
+        if (isDnsLookupFailure(t)) {
+            return;
+        }
+        if (isTlsConfigurationError(t)) {
+            return;
+        }
+        // TODO-BACKPRESSURE Nabil - SOCKS5 Revisit alongside JAVA-5205 (SOCKS5 in async) so both sync and
+        // async proxy error surfaces can be handled together — likely via a dedicated internal
+        // exception thrown from the proxy code path.
+        MongoException mongoException = (MongoException) t;
+        mongoException.addLabel(MongoException.SYSTEM_OVERLOADED_ERROR_LABEL);
+        mongoException.addLabel(MongoException.RETRYABLE_ERROR_LABEL);
+    }
+
+    static boolean isDnsLookupFailure(final Throwable t) {
+        Throwable cause = t.getCause();
+        while (cause != null) {
+            if (cause instanceof UnknownHostException) {
+                return true;
+            }
+            cause = cause.getCause();
+        }
+        return false;
+    }
+
+    static boolean isTlsConfigurationError(final Throwable t) {
+        if (!(t instanceof MongoSocketException)) {
+            return false;
+        }
+        Throwable cause = t.getCause();
+        while (cause != null) {
+            if (cause instanceof CertificateException
+                    || cause instanceof CertPathBuilderException
+                    || cause instanceof CertPathValidatorException
+                    || cause instanceof SSLPeerUnverifiedException
+                    || cause instanceof SSLProtocolException) {
+                return true;
+            }
+            if (cause instanceof SSLHandshakeException) {
+                String message = cause.getMessage();
+                if (message != null) {
+                    String lowerMessage = message.toLowerCase(Locale.ROOT);
+                    if (lowerMessage.contains("certificate")
+                            || lowerMessage.contains("verify")
+                            || lowerMessage.contains("trust")
+                            || lowerMessage.contains("hostname")
+                            || lowerMessage.contains("protocol")
+                            || lowerMessage.contains("cipher")
+                            || lowerMessage.contains("handshake_failure")) {
+                        return true;
+                    }
+                }
+            }
+            cause = cause.getCause();
+        }
+        return false;
+    }
+}

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

@@ -137,6 +137,9 @@ private void handleException(final SdamIssue sdamIssue, final boolean beforeHand
             serverMonitor.connect();
         } else if (sdamIssue.relatedToNetworkNotTimeout()
                 || (beforeHandshake && (sdamIssue.relatedToNetworkTimeout() || sdamIssue.relatedToAuth()))) {
+            if (sdamIssue.hasBackpressureLabel()) {
+                return;
+            }
             updateDescription(sdamIssue.serverDescription());
             connectionPool.invalidate(sdamIssue.exception().orElse(null));
             serverMonitor.cancelCurrentCheck();

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

@@ -230,9 +230,11 @@ public void open(final OperationContext originalOperationContext) {
         isTrue("Open already called", stream == null);
         stream = streamFactory.create(serverId.getAddress());
         OperationContext operationContext = originalOperationContext;
+        boolean beforeHandshake = true;
         try {
             stream.open(operationContext);
             InternalConnectionInitializationDescription initializationDescription = connectionInitializer.startHandshake(this, operationContext);
+            beforeHandshake = false;
 
             operationContext = operationContext.withOverride(TimeoutContext::withNewlyStartedMaintenanceTimeout);
             initAfterHandshakeStart(initializationDescription);
@@ -241,6 +243,9 @@ public void open(final OperationContext originalOperationContext) {
             initAfterHandshakeFinish(initializationDescription);
         } catch (Throwable t) {
             close();
+            if (beforeHandshake) {
+                BackpressureErrorLabeler.applyLabelsIfEligible(t);
+            }
             if (t instanceof MongoException) {
                 throw (MongoException) t;
             } else {
@@ -263,6 +268,7 @@ public void completed(@Nullable final Void aVoid) {
                             (initialResult, initialException) -> {
                                     if (initialException != null) {
                                         close();
+                                        BackpressureErrorLabeler.applyLabelsIfEligible(initialException);
                                         callback.onResult(null, initialException);
                                     } else {
                                         assertNotNull(initialResult);
@@ -278,11 +284,13 @@ public void completed(@Nullable final Void aVoid) {
                 @Override
                 public void failed(final Throwable t) {
                     close();
+                    BackpressureErrorLabeler.applyLabelsIfEligible(t);
                     callback.onResult(null, t);
                 }
             });
         } catch (Throwable t) {
             close();
+            BackpressureErrorLabeler.applyLabelsIfEligible(t);
             callback.onResult(null, t);
         }
     }

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

@@ -17,6 +17,7 @@
 package com.mongodb.internal.connection;
 
 import com.mongodb.MongoCommandException;
+import com.mongodb.MongoException;
 import com.mongodb.MongoNodeIsRecoveringException;
 import com.mongodb.MongoNotPrimaryException;
 import com.mongodb.MongoSecurityException;
@@ -162,6 +163,11 @@ boolean relatedToWriteConcern() {
             return exception instanceof MongoWriteConcernWithResponseException;
         }
 
+        boolean hasBackpressureLabel() {
+            return exception instanceof MongoException
+                    && ((MongoException) exception).hasErrorLabel(MongoException.SYSTEM_OVERLOADED_ERROR_LABEL);
+        }
+
         private static boolean stale(@Nullable final Throwable t, final ServerDescription currentServerDescription) {
             return TopologyVersionHelper.topologyVersion(t)
                     .map(candidateTopologyVersion -> TopologyVersionHelper.newerOrEqual(

driver-core/src/test/unit/com/mongodb/internal/connection/AbstractServerDiscoveryAndMonitoringTest.java

@@ -112,6 +112,7 @@ protected void applyApplicationError(final BsonDocument applicationError) {
 
         switch (when) {
             case "beforeHandshakeCompletes":
+                BackpressureErrorLabeler.applyLabelsIfEligible(exception);
                 server.sdamServerDescriptionManager().handleExceptionBeforeHandshake(
                         SdamIssue.of(exception, new SdamIssue.Context(server.serverId(), errorGeneration, maxWireVersion)));
                 break;

driver-core/src/test/unit/com/mongodb/internal/connection/DefaultServerSpecification.groovy

@@ -259,6 +259,28 @@ class DefaultServerSpecification extends Specification {
         ]
     }
 
+    def 'DNS lookup failure should invalidate the pool'() {
+        given:
+        def exceptionToThrow = new MongoSocketException('DNS lookup failed', new ServerAddress(),
+                new UnknownHostException('no such host'))
+        BackpressureErrorLabeler.applyLabelsIfEligible(exceptionToThrow)
+        assert !exceptionToThrow.hasErrorLabel(MongoException.SYSTEM_OVERLOADED_ERROR_LABEL)
+
+        def connectionPool = Mock(ConnectionPool)
+        connectionPool.get(_) >> { throw exceptionToThrow }
+        def serverMonitor = Mock(ServerMonitor)
+        def server = defaultServer(connectionPool, serverMonitor)
+
+        when:
+        server.getConnection(OPERATION_CONTEXT)
+
+        then:
+        def e = thrown(MongoException)
+        e.is(exceptionToThrow)
+        1 * connectionPool.invalidate(exceptionToThrow)
+        1 * serverMonitor.cancelCurrentCheck()
+    }
+
     def 'failed authentication should invalidate the connection pool'() {
         given:
         def connectionPool = Mock(ConnectionPool)

driver-core/src/test/unit/com/mongodb/internal/connection/ServerDiscoveryAndMonitoringTest.java

@@ -54,9 +54,6 @@ public class ServerDiscoveryAndMonitoringTest extends AbstractServerDiscoveryAnd
 
     public ServerDiscoveryAndMonitoringTest(final String description, final BsonDocument definition) {
         super(definition);
-        assumeFalse("https://jira.mongodb.org/browse/JAVA-5949",
-                description.equals("error_handling_handshake.json: Network timeouts before and after the handshake completes"));
-
         this.description = description;
         init(serverAddress -> NO_OP_SERVER_LISTENER, NO_OP_CLUSTER_LISTENER);
     }

driver-sync/src/test/functional/com/mongodb/client/ServerDiscoveryAndMonitoringProseTests.java

@@ -18,6 +18,7 @@
 
 import com.mongodb.ClusterFixture;
 import com.mongodb.MongoClientSettings;
+import com.mongodb.event.ConnectionCheckOutFailedEvent;
 import com.mongodb.event.ConnectionPoolClearedEvent;
 import com.mongodb.event.ConnectionPoolListener;
 import com.mongodb.event.ConnectionPoolReadyEvent;
@@ -26,6 +27,7 @@
 import com.mongodb.event.ServerHeartbeatSucceededEvent;
 import com.mongodb.event.ServerListener;
 import com.mongodb.event.ServerMonitorListener;
+import com.mongodb.internal.connection.TestConnectionPoolListener;
 import com.mongodb.internal.diagnostics.logging.Logger;
 import com.mongodb.internal.diagnostics.logging.Loggers;
 import com.mongodb.internal.time.TimePointTest;
@@ -47,6 +49,8 @@
 import java.util.Set;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.LinkedBlockingQueue;
 
 import static com.mongodb.ClusterFixture.configureFailPoint;
@@ -268,6 +272,68 @@ public void shouldEmitHeartbeatStartedBeforeSocketIsConnected() {
         // As it requires mocking and package access to `com.mongodb.internal.connection`
     }
 
+    /**
+     * See
+     * <a href="https://github.com/mongodb/specifications/blob/master/source/server-discovery-and-monitoring/server-discovery-and-monitoring-tests.md#connection-pool-backpressure">Connection Pool Backpressure</a>.
+     */
+    @Test
+    public void testConnectionPoolBackpressure() throws InterruptedException {
+        assumeTrue(serverVersionAtLeast(7, 0));
+
+        TestConnectionPoolListener connectionPoolListener = new TestConnectionPoolListener();
+
+        MongoClientSettings clientSettings = getMongoClientSettingsBuilder()
+                .applyToConnectionPoolSettings(builder -> builder
+                        .maxConnecting(100)
+                        .addConnectionPoolListener(connectionPoolListener))
+                .build();
+
+        try (MongoClient adminClient = MongoClients.create(getMongoClientSettingsBuilder().build());
+             MongoClient client = MongoClients.create(clientSettings)) {
+
+            MongoDatabase adminDatabase = adminClient.getDatabase("admin");
+            MongoDatabase database = client.getDatabase(getDefaultDatabaseName());
+            MongoCollection<Document> collection = database.getCollection("testCollection");
+
+            adminDatabase.runCommand(new Document("setParameter", 1)
+                    .append("ingressConnectionEstablishmentRateLimiterEnabled", true));
+            try {
+                adminDatabase.runCommand(new Document("setParameter", 1)
+                        .append("ingressConnectionEstablishmentRatePerSec", 20));
+                adminDatabase.runCommand(new Document("setParameter", 1)
+                        .append("ingressConnectionEstablishmentBurstCapacitySecs", 1));
+                adminDatabase.runCommand(new Document("setParameter", 1)
+                        .append("ingressConnectionEstablishmentMaxQueueDepth", 1));
+
+                collection.insertOne(Document.parse("{}"));
+
+                ExecutorService executor = Executors.newFixedThreadPool(100);
+                try {
+                    for (int i = 0; i < 100; i++) {
+                        executor.submit(() ->
+                                collection.find(new Document("$where", "function() { sleep(2000); return true; }")).first());
+                    }
+                    executor.shutdown();
+                    assertTrue("Executor did not terminate within timeout",
+                            executor.awaitTermination(20, SECONDS));
+                } finally {
+                    if (!executor.isTerminated()) {
+                        executor.shutdownNow();
+                    }
+                }
+
+                int failedCheckOutCount = connectionPoolListener.countEvents(ConnectionCheckOutFailedEvent.class);
+                assertTrue("Expected at least 10 ConnectionCheckOutFailedEvents, but got " + failedCheckOutCount,
+                        failedCheckOutCount >= 10);
+                assertEquals(0, connectionPoolListener.countEvents(ConnectionPoolClearedEvent.class));
+            } finally {
+                Thread.sleep(1000);
+                adminDatabase.runCommand(new Document("setParameter", 1)
+                        .append("ingressConnectionEstablishmentRateLimiterEnabled", false));
+            }
+        }
+    }
+
     private static void assertPoll(final BlockingQueue<?> queue, @Nullable final Class<?> allowed, final Set<Class<?>> required)
             throws InterruptedException {
         assertPoll(queue, allowed, required, Timeout.expiresIn(TEST_WAIT_TIMEOUT_MILLIS, MILLISECONDS, ZERO_DURATION_MEANS_EXPIRED));

driver-sync/src/test/functional/com/mongodb/client/unified/EventMatcher.java

@@ -26,6 +26,7 @@
 import com.mongodb.event.CommandStartedEvent;
 import com.mongodb.event.CommandSucceededEvent;
 import com.mongodb.event.ConnectionCheckOutFailedEvent;
+import com.mongodb.event.ConnectionCheckedInEvent;
 import com.mongodb.event.ConnectionClosedEvent;
 import com.mongodb.event.ConnectionCreatedEvent;
 import com.mongodb.event.ConnectionPoolClearedEvent;
@@ -208,6 +209,12 @@ public void waitForConnectionPoolEvents(final String client, final BsonDocument
             case "connectionReadyEvent":
                 eventClass = ConnectionReadyEvent.class;
                 break;
+            case "connectionClosedEvent":
+                eventClass = ConnectionClosedEvent.class;
+                break;
+            case "connectionCheckedInEvent":
+                eventClass = ConnectionCheckedInEvent.class;
+                break;
             default:
                 throw new UnsupportedOperationException("Unsupported event: " + event.getFirstKey());
         }
@@ -436,9 +443,16 @@ private static boolean serverDescriptionChangedEventMatches(final BsonDocument e
         switch (newType) {
             case "Unknown":
                 return event.getNewDescription().getType() == ServerType.UNKNOWN;
-            case "LoadBalancer": {
+            case "LoadBalancer":
                 return event.getNewDescription().getType() == ServerType.LOAD_BALANCER;
-            }
+            case "Mongos":
+                return event.getNewDescription().getType() == ServerType.SHARD_ROUTER;
+            case "Standalone":
+                return event.getNewDescription().getType() == ServerType.STANDALONE;
+            case "RSPrimary":
+                return event.getNewDescription().getType() == ServerType.REPLICA_SET_PRIMARY;
+            case "RSSecondary":
+                return event.getNewDescription().getType() == ServerType.REPLICA_SET_SECONDARY;
             default:
                 throw new UnsupportedOperationException();
         }

driver-sync/src/test/functional/com/mongodb/client/unified/UnifiedTest.java

@@ -800,6 +800,8 @@ private OperationResult executeWaitForEvent(final UnifiedTestContext context, fi
             case "poolReadyEvent":
             case "connectionCreatedEvent":
             case "connectionReadyEvent":
+            case "connectionClosedEvent":
+            case "connectionCheckedInEvent":
                 context.getEventMatcher().waitForConnectionPoolEvents(clientId, event, count, entities.getConnectionPoolListener(clientId));
                 break;
             case "serverHeartbeatStartedEvent":

driver-sync/src/test/functional/com/mongodb/client/unified/UnifiedTestModifications.java

@@ -439,15 +439,7 @@ public static void applyCustomizations(final TestDef def) {
                 .file("server-discovery-and-monitoring", "pool-clear-on-error-checkout");
         def.skipJira("https://jira.mongodb.org/browse/JAVA-5664")
                 .file("server-discovery-and-monitoring", "pool-cleared-on-min-pool-size-population-error");
-        def.skipJira("https://jira.mongodb.org/browse/JAVA-5949")
-                .file("server-discovery-and-monitoring", "backpressure-network-error-fail-single");
-        def.skipJira("https://jira.mongodb.org/browse/JAVA-5949")
-                .file("server-discovery-and-monitoring", "backpressure-network-timeout-error-single");
-        def.skipJira("https://jira.mongodb.org/browse/JAVA-5949")
-                .file("server-discovery-and-monitoring", "backpressure-network-error-fail-replicaset");
-        def.skipJira("https://jira.mongodb.org/browse/JAVA-5949")
-                .file("server-discovery-and-monitoring", "backpressure-network-timeout-error-replicaset");
-        def.skipJira("https://jira.mongodb.org/browse/JAVA-5949")
+        def.skipJira("https://jira.mongodb.org/browse/JAVA-6174")
                 .file("server-discovery-and-monitoring", "backpressure-server-description-unchanged-on-min-pool-size-population-error");
 
         // session tests