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

@@ -17,20 +17,21 @@
 
 import com.mongodb.MongoClientException;
 import com.mongodb.MongoOperationTimeoutException;
+import com.mongodb.internal.function.CheckedFunction;
+import com.mongodb.internal.function.CheckedSupplier;
 import com.mongodb.internal.time.StartTime;
 import com.mongodb.internal.time.Timeout;
 import com.mongodb.lang.Nullable;
 import com.mongodb.session.ClientSession;
 
 import java.util.Objects;
-import java.util.Optional;
-import java.util.function.Supplier;
+import java.util.concurrent.TimeUnit;
 
 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 java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.NANOSECONDS;
 
 /**
  * Timeout Context.
@@ -48,16 +49,21 @@ public class TimeoutContext {
     private Timeout computedServerSelectionTimeout;
     private long minRoundTripTimeMS = 0;
 
-    private MaxTimeSupplier maxTimeSupplier = this::calculateMaxTimeMS;
+    @Nullable
+    private Long maxTimeMSOverride = 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);
     }
@@ -129,8 +135,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 remove this method (do not inline, but use lambdas)
+        return Timeout.nullAsInfinite(timeout).run(NANOSECONDS, () -> false, (ns) -> false, () -> true);
     }
 
     /**
@@ -144,19 +150,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 +162,15 @@ private boolean hasTimedOutForCommandExecution() {
      * @return timeout to use.
      */
     public long timeoutOrAlternative(final long alternativeTimeoutMS) {
-        Long timeoutMS = timeoutSettings.getTimeoutMS();
-        if (timeoutMS == null) {
+        if (timeout == 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) {
-            return alternativeTimeoutMS;
-        } else if (alternativeTimeoutMS == 0) {
-            return timeoutRemainingMS();
         } else {
-            return Math.min(timeoutRemainingMS(), alternativeTimeoutMS);
+            return timeout.run(MILLISECONDS,
+                    () -> 0L,
+                    (ms) -> ms,
+                    () -> {
+                        throw createMongoTimeoutException("The operation timeout has expired.");
+                    });
         }
     }
 
@@ -200,27 +182,33 @@ public long getMaxAwaitTimeMS() {
         return timeoutSettings.getMaxAwaitTimeMS();
     }
 
-    public long getMaxTimeMS() {
-        return notNull("Should never be null", maxTimeSupplier.get());
+    public Timeout getMaxTimeMSTimeout() {
+        if (maxTimeMSOverride != null) {
+            return new SingleUseTimeout(maxTimeMSOverride);
+        } else {
+            long maxTimeMS = timeoutSettings.getMaxTimeMS();
+            return timeout != null
+                    ? timeout.shortenBy(minRoundTripTimeMS, MILLISECONDS)
+                    : new SingleUseTimeout(maxTimeMS);
+        }
     }
 
-    private long calculateMaxTimeMS() {
-        long maxTimeMS = timeoutOrAlternative(timeoutSettings.getMaxTimeMS());
-        if (timeout == null || timeout.isInfinite()) {
-            return maxTimeMS;
-        }
-        if (minRoundTripTimeMS >= maxTimeMS) {
-            throw createMongoTimeoutException();
-        }
-        return maxTimeMS - minRoundTripTimeMS;
+    // TODO (CSOT) JAVA-5385 used only by tests?
+    public long getMaxTimeMS() {
+        return getMaxTimeMSTimeout().run(MILLISECONDS,
+                () -> 0L,
+                (ms) -> ms,
+                () -> {
+                    throw createMongoRoundTripTimeoutException();
+                });
     }
 
-    public void setMaxTimeSupplier(final MaxTimeSupplier maxTimeSupplier) {
-        this.maxTimeSupplier = maxTimeSupplier;
+    public void resetToDefaultMaxTime() {
+        this.maxTimeMSOverride = null;
     }
 
-    public void resetToDefaultMaxTimeSupplier() {
-        this.maxTimeSupplier = this::calculateMaxTimeMS;
+    public void setMaxTimeOverride(final long maxTimeMS) {
+        this.maxTimeMSOverride = maxTimeMS;
     }
 
     public Long getMaxCommitTimeMS() {
@@ -236,8 +224,52 @@ 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.expiresInWithZeroAsInfinite(getTimeoutSettings().getConnectTimeoutMS(), MILLISECONDS),
+                Timeout.nullAsInfinite(timeout));
+    }
+
+    /**
+     * A timeout that begins at the moment that it is used.
+     */
+    static class SingleUseTimeout implements Timeout {
+        @Nullable
+        private final Long nanos;
+        private boolean used = false;
+
+        public SingleUseTimeout(final long ms) {
+            if (ms == 0) {
+                nanos = null;
+            } else {
+                nanos = NANOSECONDS.convert(ms, MILLISECONDS);
+            }
+        }
+
+        @Override
+        public Timeout shortenBy(final long amount, final TimeUnit timeUnit) {
+            throw new AssertionError("Single-use timeout must not be shortened");
+        }
+
+        @Override
+        public <T, E extends Exception> T checkedRun(final TimeUnit timeUnit, final CheckedSupplier<T, E> onInfinite,
+                final CheckedFunction<Long, T, E> onHasRemaining, final CheckedSupplier<T, E> onExpired) throws E {
+            if (used) {
+                throw new AssertionError("Single-use timeout must not be used multiple times");
+            }
+            used = true;
+            if (this.nanos == null) {
+                return onInfinite.get();
+            }
+            long remaining = timeUnit.convert(nanos, NANOSECONDS);
+            if (remaining == 0) {
+                return onExpired.get();
+            } else {
+                return onHasRemaining.apply(remaining);
+            }
+        }
     }
 
     public void resetTimeout() {
@@ -249,9 +281,13 @@ public void resetTimeout() {
      * 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).run(NANOSECONDS,
+                () -> timeout,
+                (ms) -> calculateTimeout(timeoutSettings.getTimeoutMS()),
+                () -> calculateTimeout(timeoutSettings.getTimeoutMS()));
     }
 
     public TimeoutContext withAdditionalReadTimeout(final int additionalReadTimeout) {
@@ -267,14 +303,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{"
@@ -307,8 +335,9 @@ public int hashCode() {
 
     @Nullable
     public static Timeout calculateTimeout(@Nullable final Long timeoutMS) {
+        // TODO-CSOT rename to startTimeout?
         if (timeoutMS != null) {
-            return timeoutMS == 0 ? Timeout.infinite() : Timeout.expiresIn(timeoutMS, MILLISECONDS);
+            return Timeout.expiresInWithZeroAsInfinite(timeoutMS, MILLISECONDS);
         }
         return null;
     }
@@ -334,7 +363,7 @@ public Timeout computeServerSelectionTimeout() {
             return serverSelectionTimeout;
         }
 
-        if (serverSelectionTimeout.orEarlier(timeout) == timeout) {
+        if (timeout != null && Timeout.earliest(serverSelectionTimeout, timeout) == timeout) {
             return timeout;
         }
 
@@ -363,9 +392,4 @@ public Timeout startWaitQueueTimeout(final StartTime checkoutStart) {
     public Timeout getTimeout() {
         return timeout;
     }
-
-    public interface MaxTimeSupplier extends Supplier<Long> {
-        @Override
-        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