Store computedServerSelectionTimeout so it can be reused.

Needed for use with authentication

JAVA-5211

Author: rozza

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

@@ -40,6 +40,8 @@ public class TimeoutContext {
 
     @Nullable
     private Timeout timeout;
+    @Nullable
+    private Timeout computedServerSelectionTimeout;
     private long minRoundTripTimeMS = 0;
 
     public static MongoOperationTimeoutException createMongoTimeoutException() {
@@ -191,14 +193,17 @@ public long getWriteTimeoutMS() {
         return timeoutOrAlternative(0);
     }
 
+    public int getConnectTimeoutMs() {
+        return (int) calculateMin(getTimeoutSettings().getConnectTimeoutMS());
+    }
 
     public void resetTimeout() {
         assertNotNull(timeout);
         timeout = calculateTimeout(timeoutSettings.getTimeoutMS());
     }
 
     /**
-     * Resest the timeout if this timeout context is being used by pool maintenance
+     * Resets the timeout if this timeout context is being used by pool maintenance
      */
     public void resetMaintenanceTimeout() {
         if (isMaintenanceContext && timeout != null && !timeout.isInfinite()) {
@@ -266,22 +271,47 @@ public static Timeout calculateTimeout(@Nullable final Long timeoutMS) {
     }
 
     public Timeout computedServerSelectionTimeout() {
-        long ms = getTimeoutSettings().getServerSelectionTimeoutMS();
-        Timeout serverSelectionTimeout = StartTime.now().timeoutAfterOrInfiniteIfNegative(ms, MILLISECONDS);
-        return serverSelectionTimeout.orEarlier(timeout);
+        if (computedServerSelectionTimeout == null) {
+            computedServerSelectionTimeout = StartTime.now()
+                    .timeoutAfterOrInfiniteIfNegative(getTimeoutSettings().getServerSelectionTimeoutMS(), MILLISECONDS)
+                    .orEarlier(timeout);
+        }
+        return computedServerSelectionTimeout;
+    }
+
+    /**
+     * Returns the timeout context to use for the handshake process
+     *
+     * <p>Uses the current context if:
+     * <ul>
+     *     <li>If in a maintenance context</li>
+     *     <li>If has no timeoutMS, so to keep the same legacy behavior.</li>
+     *     <li>If there is no computedServerSelectionTimeout.</li>
+     *     <li>If the computedServerSelectionTimeout is the same as the timeout.</li>
+     * </ul>
+     *
+     * @return the timeout context
+     */
+    public TimeoutContext withComputedServerSelectionTimeoutContext() {
+        if (isMaintenanceContext || !hasTimeoutMS() || computedServerSelectionTimeout == null
+                || Objects.equals(computedServerSelectionTimeout, timeout)) {
+            computedServerSelectionTimeout = null;
+            return this;
+        }
+
+        TimeoutContext timeoutContext = new TimeoutContext(false, timeoutSettings, computedServerSelectionTimeout);
+        computedServerSelectionTimeout = null;
+        return timeoutContext;
     }
 
     public Timeout startWaitQueueTimeout(final StartTime checkoutStart) {
         final long ms = getTimeoutSettings().getMaxWaitTimeMS();
         return checkoutStart.timeoutAfterOrInfiniteIfNegative(ms, MILLISECONDS);
     }
 
-    public int getConnectTimeoutMs() {
-        return (int) getTimeoutSettings().getConnectTimeoutMS();
-    }
-
     @Nullable
     public Timeout getTimeout() {
         return timeout;
     }
+
 }

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

@@ -196,10 +196,13 @@ public int getGeneration() {
     }
 
     @Override
-    public void open(final OperationContext operationContext) {
+    public void open(final OperationContext originalOperationContext) {
         isTrue("Open already called", stream == null);
         stream = streamFactory.create(serverId.getAddress());
         try {
+            OperationContext operationContext = originalOperationContext
+                    .withTimeoutContext(originalOperationContext.getTimeoutContext().withComputedServerSelectionTimeoutContext());
+
             stream.open(operationContext);
 
             InternalConnectionInitializationDescription initializationDescription = connectionInitializer.startHandshake(this, operationContext);
@@ -218,9 +221,13 @@ public void open(final OperationContext operationContext) {
     }
 
     @Override
-    public void openAsync(final OperationContext operationContext, final SingleResultCallback<Void> callback) {
+    public void openAsync(final OperationContext originalOperationContext, final SingleResultCallback<Void> callback) {
         isTrue("Open already called", stream == null, callback);
         try {
+
+            OperationContext operationContext = originalOperationContext
+                    .withTimeoutContext(originalOperationContext.getTimeoutContext().withComputedServerSelectionTimeoutContext());
+
             stream = streamFactory.create(serverId.getAddress());
             stream.openAsync(operationContext, new AsyncCompletionHandler<Void>() {