Use connectTimeoutMS for connection establishment in maintenance Time…

driver-core/src/main/com/mongodb/internal/TimeoutContext.java

@@ -289,6 +289,10 @@ public long getWriteTimeoutMS() {
 
     public int getConnectTimeoutMs() {
         final long connectTimeoutMS = getTimeoutSettings().getConnectTimeoutMS();
+        if (isMaintenanceContext) {
+            return (int) connectTimeoutMS;
+        }
+
         return Math.toIntExact(Timeout.nullAsInfinite(timeout).call(MILLISECONDS,
                 () -> connectTimeoutMS,
                 (ms) -> connectTimeoutMS == 0 ? ms : Math.min(ms, connectTimeoutMS),

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

@@ -81,11 +81,11 @@ public SocksSocket(@Nullable final Socket socket, final ProxySettings proxySetti
     }
 
     @Override
-    public void connect(final SocketAddress endpoint, final int timeoutMs) throws IOException {
+    public void connect(final SocketAddress endpoint, final int connectTimeoutMs) throws IOException {
         // `Socket` requires `IllegalArgumentException`
-        isTrueArgument("timeoutMs", timeoutMs >= 0);
+        isTrueArgument("connectTimeoutMs", connectTimeoutMs >= 0);
         try {
-            Timeout timeout = Timeout.expiresIn(timeoutMs, MILLISECONDS, ZERO_DURATION_MEANS_INFINITE);
+            Timeout timeout = Timeout.expiresIn(connectTimeoutMs, MILLISECONDS, ZERO_DURATION_MEANS_INFINITE);
             InetSocketAddress unresolvedAddress = (InetSocketAddress) endpoint;
             assertTrue(unresolvedAddress.isUnresolved());
             this.remoteAddress = unresolvedAddress;

driver-core/src/test/functional/com/mongodb/internal/connection/SocketStreamHelperSpecification.groovy

@@ -18,8 +18,11 @@ package com.mongodb.internal.connection
 
 import com.mongodb.ClusterFixture
 import com.mongodb.MongoInternalException
+import com.mongodb.MongoOperationTimeoutException
 import com.mongodb.connection.SocketSettings
 import com.mongodb.connection.SslSettings
+import com.mongodb.internal.TimeoutContext
+import com.mongodb.internal.TimeoutSettings
 import jdk.net.ExtendedSocketOptions
 import spock.lang.IgnoreIf
 import spock.lang.Specification
@@ -78,6 +81,30 @@ class SocketStreamHelperSpecification extends Specification {
         socket?.close()
     }
 
+    def 'should throw MongoOperationTimeoutException during initialization when timeoutMS expires'() {
+        given:
+        Socket socket = SocketFactory.default.createSocket()
+
+        when:
+        SocketStreamHelper.initialize(
+                OPERATION_CONTEXT.withTimeoutContext(new TimeoutContext(
+                                new TimeoutSettings(
+                                        1,
+                                        100,
+                                        100,
+                                        1,
+                                        100))),
+                socket, getSocketAddresses(getPrimary(), new DefaultInetAddressResolver()).get(0),
+                SocketSettings.builder().build(), SslSettings.builder().build())
+
+        then:
+        thrown(MongoOperationTimeoutException)
+
+        cleanup:
+        socket?.close()
+    }
+
+
     def 'should connect socket()'() {
         given:
         Socket socket = SocketFactory.default.createSocket()

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

@@ -47,6 +47,7 @@
 import com.mongodb.event.ConnectionClosedEvent;
 import com.mongodb.event.ConnectionCreatedEvent;
 import com.mongodb.event.ConnectionReadyEvent;
+import com.mongodb.internal.connection.InternalStreamConnection;
 import com.mongodb.internal.connection.ServerHelper;
 import com.mongodb.internal.connection.TestCommandListener;
 import com.mongodb.internal.connection.TestConnectionPoolListener;
@@ -1018,6 +1019,44 @@ public void shouldThrowTimeoutExceptionForSubsequentCommitTransaction() {
         assertEquals(1, failedEvents.size());
     }
 
+    /**
+     * Not a prose spec test. However, it is additional test case for better coverage.
+     * <p>
+     * From the spec:
+     * - When doing `minPoolSize` maintenance, `connectTimeoutMS` is used as the timeout for socket establishment.
+     */
+    @Test
+    @DisplayName("Should use connectTimeoutMS when establishing connection in background")
+    public void shouldUseConnectTimeoutMsWhenEstablishingConnectionInBackground() {
+        assumeTrue(serverVersionAtLeast(4, 4));
+
+        collectionHelper.runAdminCommand("{"
+                + "configureFailPoint: \"" + FAIL_COMMAND_NAME + "\","
+                + "mode: \"alwaysOn\","
+                + "  data: {"
+                + "    failCommands: [\"hello\", \"isMaster\"],"
+                + "    blockConnection: true,"
+                + "    blockTimeMS: " + 500
+                + "  }"
+                + "}");
+
+        try (MongoClient ignored = createMongoClient(getMongoClientSettingsBuilder()
+                .applyToConnectionPoolSettings(builder -> builder.minSize(1))
+                // Use a very short timeout to ensure that the connection establishment will fail on the first handshake command.
+                .timeout(10, TimeUnit.MILLISECONDS))) {
+            InternalStreamConnection.setRecordEverything(true);
+
+            // Wait for the connection to start establishment in the background.
+            sleep(1000);
+        } finally {
+            InternalStreamConnection.setRecordEverything(false);
+        }
+
+        List<CommandFailedEvent> commandFailedEvents = commandListener.getCommandFailedEvents("isMaster");
+        assertEquals(1, commandFailedEvents.size());
+        assertInstanceOf(MongoOperationTimeoutException.class, commandFailedEvents.get(0).getThrowable());
+    }
+
     private static Stream<Arguments> test8ServerSelectionArguments() {
         return Stream.of(
                 Arguments.of(Named.of("serverSelectionTimeoutMS honored if timeoutMS is not set",