Replace Timeout methods with branching "run" methods

config/spotbugs/exclude.xml

@@ -280,17 +280,14 @@
     <!-- Ignoring await return; intended to be used in a loop -->
     <Match>
         <Class name="com.mongodb.internal.time.Timeout"/>
-        <Method name="awaitOn"/>
         <Bug pattern="RV_RETURN_VALUE_IGNORED_BAD_PRACTICE"/>
     </Match>
     <Match>
         <Class name="com.mongodb.internal.time.Timeout"/>
-        <Method name="awaitOn"/>
         <Bug pattern="RV_RETURN_VALUE_IGNORED"/>
     </Match>
     <Match>
         <Class name="com.mongodb.internal.time.Timeout"/>
-        <Method name="awaitOn"/>
         <Bug pattern="WA_AWAIT_NOT_IN_LOOP"/>
     </Match>
 

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

@@ -44,6 +44,7 @@ public final class MongoOperationTimeoutException extends MongoTimeoutException
      * @param message the message
      */
     public MongoOperationTimeoutException(final String message) {
+        // TODO (CSOT) - JAVA-5248 move all messages here as constants to ensure consistency?
         super(message);
     }
 

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

@@ -17,6 +17,7 @@
 package com.mongodb.internal;
 
 import com.mongodb.MongoInterruptedException;
+import com.mongodb.internal.function.CheckedSupplier;
 
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;

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

@@ -23,14 +23,14 @@
 import com.mongodb.session.ClientSession;
 
 import java.util.Objects;
-import java.util.Optional;
-import java.util.function.Supplier;
+import java.util.function.LongConsumer;
 
 import static com.mongodb.assertions.Assertions.assertNotNull;
 import static com.mongodb.assertions.Assertions.assertNull;
 import static com.mongodb.assertions.Assertions.isTrue;
-import static com.mongodb.assertions.Assertions.notNull;
+import static com.mongodb.internal.VisibleForTesting.AccessModifier.*;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.NANOSECONDS;
 
 /**
  * Timeout Context.
@@ -48,16 +48,21 @@ public class TimeoutContext {
     private Timeout computedServerSelectionTimeout;
     private long minRoundTripTimeMS = 0;
 
-    private MaxTimeSupplier maxTimeSupplier = this::calculateMaxTimeMS;
+    @Nullable
+    private MaxTimeSupplier maxTimeSupplier = null;
 
-    public static MongoOperationTimeoutException createMongoTimeoutException() {
-        return createMongoTimeoutException("Remaining timeoutMS is less than the servers minimum round trip time.");
+    public static MongoOperationTimeoutException createMongoRoundTripTimeoutException() {
+        return createMongoTimeoutException("Remaining timeoutMS is less than the server's minimum round trip time.");
     }
 
     public static MongoOperationTimeoutException createMongoTimeoutException(final String message) {
         return new MongoOperationTimeoutException(message);
     }
 
+    public static void throwMongoTimeoutException(final String message) {
+        throw new MongoOperationTimeoutException(message);
+    }
+
     public static MongoOperationTimeoutException createMongoTimeoutException(final Throwable cause) {
         return createMongoTimeoutException("Operation timed out: " + cause.getMessage(), cause);
     }
@@ -70,7 +75,7 @@ public static MongoOperationTimeoutException createMongoTimeoutException(final S
     }
 
     public static TimeoutContext createMaintenanceTimeoutContext(final TimeoutSettings timeoutSettings) {
-        return new TimeoutContext(true, timeoutSettings, calculateTimeout(timeoutSettings.getTimeoutMS()));
+        return new TimeoutContext(true, timeoutSettings, startTimeout(timeoutSettings.getTimeoutMS()));
     }
 
     public static TimeoutContext createTimeoutContext(final ClientSession session, final TimeoutSettings timeoutSettings) {
@@ -101,7 +106,7 @@ public static TimeoutContext createTimeoutContext(final ClientSession session, f
     }
 
     public TimeoutContext(final TimeoutSettings timeoutSettings) {
-        this(false, timeoutSettings, calculateTimeout(timeoutSettings.getTimeoutMS()));
+        this(false, timeoutSettings, startTimeout(timeoutSettings.getTimeoutMS()));
     }
 
     public TimeoutContext(final TimeoutSettings timeoutSettings, @Nullable final Timeout timeout) {
@@ -129,8 +134,8 @@ public boolean hasTimeoutMS() {
      * @return true if the timeout has been set and it has expired
      */
     public boolean hasExpired() {
-        // Use timeout.remaining instead of timeout.hasExpired that measures in nanoseconds.
-        return timeout != null && !timeout.isInfinite() && timeout.remaining(MILLISECONDS) <= 0;
+        // TODO (CSOT) this method leaks Timeout internals, should be removed (not inlined, but inverted using lambdas)
+        return Timeout.nullAsInfinite(timeout).call(NANOSECONDS, () -> false, (ns) -> false, () -> true);
     }
 
     /**
@@ -144,19 +149,9 @@ public TimeoutContext minRoundTripTimeMS(final long minRoundTripTimeMS) {
         return this;
     }
 
-    public Optional<MongoOperationTimeoutException> validateHasTimedOutForCommandExecution() {
-        if (hasTimedOutForCommandExecution()) {
-            return Optional.of(createMongoTimeoutException());
-        }
-        return Optional.empty();
-    }
-
-    private boolean hasTimedOutForCommandExecution() {
-        if (timeout == null || timeout.isInfinite()) {
-            return false;
-        }
-        long remaining = timeout.remaining(MILLISECONDS);
-        return remaining <= 0 || minRoundTripTimeMS > remaining;
+    @Nullable
+    public Timeout timeoutIncludingRoundTrip() {
+        return timeout == null ? null : timeout.shortenBy(minRoundTripTimeMS, MILLISECONDS);
     }
 
     /**
@@ -166,29 +161,15 @@ private boolean hasTimedOutForCommandExecution() {
      * @return timeout to use.
      */
     public long timeoutOrAlternative(final long alternativeTimeoutMS) {
-        Long timeoutMS = timeoutSettings.getTimeoutMS();
-        if (timeoutMS == null) {
-            return alternativeTimeoutMS;
-        }
-        return timeoutRemainingMS();
-    }
-
-    /**
-     * Calculates the minimum timeout value between two possible timeouts.
-     *
-     * @param alternativeTimeoutMS the alternative timeout
-     * @return the minimum value to use.
-     */
-    public long calculateMin(final long alternativeTimeoutMS) {
-        Long timeoutMS = timeoutSettings.getTimeoutMS();
-        if (timeoutMS == null) {
-            return alternativeTimeoutMS;
-        } else if (timeoutMS == 0) {
+        if (timeout == null) {
             return alternativeTimeoutMS;
-        } else if (alternativeTimeoutMS == 0) {
-            return timeoutRemainingMS();
         } else {
-            return Math.min(timeoutRemainingMS(), alternativeTimeoutMS);
+            return timeout.call(MILLISECONDS,
+                    () -> 0L,
+                    (ms) -> ms,
+                    () -> {
+                        throw createMongoTimeoutException("The operation timeout has expired.");
+                    });
         }
     }
 
@@ -200,30 +181,43 @@ public long getMaxAwaitTimeMS() {
         return timeoutSettings.getMaxAwaitTimeMS();
     }
 
-    public long getMaxTimeMS() {
-        return notNull("Should never be null", maxTimeSupplier.get());
-    }
+    public void runMaxTimeMSTimeout(final Runnable onInfinite, final LongConsumer onRemaining,
+            final Runnable onExpired) {
+        if (maxTimeSupplier != null) {
+            runWithFixedTimout(maxTimeSupplier.get(), onInfinite, onRemaining);
+            return;
+        }
 
-    private long calculateMaxTimeMS() {
-        long maxTimeMS = timeoutOrAlternative(timeoutSettings.getMaxTimeMS());
-        if (timeout == null || timeout.isInfinite()) {
-            return maxTimeMS;
+        if (timeout != null) {
+            timeout.shortenBy(minRoundTripTimeMS, MILLISECONDS)
+                    .run(MILLISECONDS, onInfinite, onRemaining, onExpired);
+        } else {
+            runWithFixedTimout(timeoutSettings.getMaxTimeMS(), onInfinite, onRemaining);
         }
-        if (minRoundTripTimeMS >= maxTimeMS) {
-            throw createMongoTimeoutException();
+    }
+
+    private static void runWithFixedTimout(final long ms, final Runnable onInfinite, final LongConsumer onRemaining) {
+        if (ms == 0) {
+            onInfinite.run();
+        } else {
+            onRemaining.accept(ms);
         }
-        return maxTimeMS - minRoundTripTimeMS;
     }
 
-    public void setMaxTimeSupplier(final MaxTimeSupplier maxTimeSupplier) {
-        this.maxTimeSupplier = maxTimeSupplier;
+    public void resetToDefaultMaxTime() {
+        this.maxTimeSupplier = null;
     }
 
-    public void resetToDefaultMaxTimeSupplier() {
-        this.maxTimeSupplier = this::calculateMaxTimeMS;
+    public void setMaxTimeOverride(final long maxTimeMS) {
+        this.maxTimeSupplier = () -> maxTimeMS;
     }
 
-    public Long getMaxCommitTimeMS() {
+    public void setMaxTimeOverrideToMaxCommitTime() {
+        this.maxTimeSupplier = () -> getMaxCommitTimeMS();
+    }
+
+    @VisibleForTesting(otherwise = PRIVATE)
+    public long getMaxCommitTimeMS() {
         Long maxCommitTimeMS = timeoutSettings.getMaxCommitTimeMS();
         return timeoutOrAlternative(maxCommitTimeMS != null ? maxCommitTimeMS : 0);
     }
@@ -236,22 +230,30 @@ public long getWriteTimeoutMS() {
         return timeoutOrAlternative(0);
     }
 
-    public int getConnectTimeoutMs() {
-        return (int) calculateMin(getTimeoutSettings().getConnectTimeoutMS());
+    public Timeout createConnectTimeoutMs() {
+        // null timeout treated as infinite will be later than the other
+
+        return Timeout.earliest(
+                Timeout.expiresIn(getTimeoutSettings().getConnectTimeoutMS(), MILLISECONDS, Timeout.ZeroDurationIs.INFINITE),
+                Timeout.nullAsInfinite(timeout));
     }
 
     public void resetTimeout() {
         assertNotNull(timeout);
-        timeout = calculateTimeout(timeoutSettings.getTimeoutMS());
+        timeout = startTimeout(timeoutSettings.getTimeoutMS());
     }
 
     /**
      * Resets the timeout if this timeout context is being used by pool maintenance
      */
     public void resetMaintenanceTimeout() {
-        if (isMaintenanceContext && timeout != null && !timeout.isInfinite()) {
-            timeout = calculateTimeout(timeoutSettings.getTimeoutMS());
+        if (!isMaintenanceContext) {
+            return;
         }
+        timeout = Timeout.nullAsInfinite(timeout).call(NANOSECONDS,
+                () -> timeout,
+                (ms) -> startTimeout(timeoutSettings.getTimeoutMS()),
+                () -> startTimeout(timeoutSettings.getTimeoutMS()));
     }
 
     public TimeoutContext withAdditionalReadTimeout(final int additionalReadTimeout) {
@@ -267,14 +269,6 @@ public TimeoutContext withAdditionalReadTimeout(final int additionalReadTimeout)
         return new TimeoutContext(timeoutSettings.withReadTimeoutMS(newReadTimeout > 0 ? newReadTimeout : Long.MAX_VALUE));
     }
 
-    private long timeoutRemainingMS() {
-        assertNotNull(timeout);
-        if (timeout.hasExpired()) {
-            throw createMongoTimeoutException("The operation timeout has expired.");
-        }
-        return timeout.isInfinite() ? 0 : timeout.remaining(MILLISECONDS);
-    }
-
     @Override
     public String toString() {
         return "TimeoutContext{"
@@ -306,9 +300,9 @@ public int hashCode() {
     }
 
     @Nullable
-    public static Timeout calculateTimeout(@Nullable final Long timeoutMS) {
+    public static Timeout startTimeout(@Nullable final Long timeoutMS) {
         if (timeoutMS != null) {
-            return timeoutMS == 0 ? Timeout.infinite() : Timeout.expiresIn(timeoutMS, MILLISECONDS);
+            return Timeout.expiresIn(timeoutMS, MILLISECONDS, Timeout.ZeroDurationIs.INFINITE);
         }
         return null;
     }
@@ -334,7 +328,7 @@ public Timeout computeServerSelectionTimeout() {
             return serverSelectionTimeout;
         }
 
-        if (serverSelectionTimeout.orEarlier(timeout) == timeout) {
+        if (timeout != null && Timeout.earliest(serverSelectionTimeout, timeout) == timeout) {
             return timeout;
         }
 
@@ -364,8 +358,7 @@ public Timeout getTimeout() {
         return timeout;
     }
 
-    public interface MaxTimeSupplier extends Supplier<Long> {
-        @Override
-        Long get();
+    public interface MaxTimeSupplier {
+        long get();
     }
 }

driver-core/src/main/com/mongodb/internal/async/package-info.java

@@ -15,7 +15,6 @@
  */
 
 /**
- * This package contains cluster and connection event related classes
  */
 
 @NonNullApi

driver-core/src/main/com/mongodb/internal/authentication/package-info.java

@@ -15,7 +15,6 @@
  */
 
 /**
- * This package contains cluster and connection event related classes
  */
 
 @NonNullApi