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 move all messages here as constants, to ensure consistency?
         super(message);
     }
 

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

@@ -0,0 +1,32 @@
+/*
+ * 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;
+
+/**
+ * <p>This class is not part of the public API and may be removed or changed at any time</p>
+ */
+@FunctionalInterface
+public interface CheckedConsumer<T, E extends Exception> {
+
+    /**
+     * Performs this operation on the given argument.
+     *
+     * @param t the input argument
+     * @throws E the checked exception to throw
+     */
+    void accept(T t) throws E;
+}

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

@@ -0,0 +1,33 @@
+/*
+ * 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;
+
+/**
+ * <p>This class is not part of the public API and may be removed or changed at any time</p>
+ */
+@FunctionalInterface
+public interface CheckedFunction<T, R, E extends Exception> {
+
+    /**
+     * Applies the function to the given argument.
+     *
+     * @param t the function argument
+     * @return the function result
+     * @throws E the checked exception to throw
+     */
+    R apply(T t) throws E;
+}

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

@@ -0,0 +1,31 @@
+/*
+ * 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;
+
+/**
+ * <p>This class is not part of the public API and may be removed or changed at any time</p>
+ */
+@FunctionalInterface
+public interface CheckedRunnable<E extends Exception> {
+
+    /**
+     * Checked run.
+     *
+     * @throws E the checked exception to throw
+     */
+    void run() throws E;
+}

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

@@ -23,12 +23,12 @@
 import com.mongodb.session.ClientSession;
 
 import java.util.Objects;
-import java.util.Optional;
 
 import static com.mongodb.assertions.Assertions.assertNotNull;
 import static com.mongodb.assertions.Assertions.assertNull;
 import static com.mongodb.assertions.Assertions.isTrue;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.NANOSECONDS;
 
 /**
  * Timeout Context.
@@ -46,7 +46,7 @@ public class TimeoutContext {
     private Timeout computedServerSelectionTimeout;
     private long minRoundTripTimeMS = 0;
 
-    public static MongoOperationTimeoutException createMongoTimeoutException() {
+    public static MongoOperationTimeoutException createMongoRoundTripTimeoutException() {
         return createMongoTimeoutException("Remaining timeoutMS is less than the servers minimum round trip time.");
     }
 
@@ -125,8 +125,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.run(timeout, NANOSECONDS, () -> false, (ns) -> false, () -> true, () -> false);
     }
 
     /**
@@ -140,19 +140,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,9 +156,17 @@ public long timeoutOrAlternative(final long alternativeTimeoutMS) {
         if (timeoutMS == null) {
             return alternativeTimeoutMS;
         } else if (timeoutMS == 0) {
+            // TODO-CSOT is this correct? shouldn't this check/return the remaining timeout ms?
             return timeoutMS;
         } else {
-            return timeoutRemainingMS();
+            assertNotNull(timeout);
+            // TODO-CSOT refactor once above addressed? Unclear how to avoid extracting numbers.
+            return timeout.run(MILLISECONDS,
+                    () -> 0L,
+                    (ms) -> ms,
+                    () -> {
+                        throw createMongoTimeoutException("The operation timeout has expired.");
+                    });
         }
     }
 
@@ -182,12 +180,23 @@ public long calculateMin(final long alternativeTimeoutMS) {
         Long timeoutMS = timeoutSettings.getTimeoutMS();
         if (timeoutMS == null) {
             return alternativeTimeoutMS;
-        } else if (timeoutMS == 0) {
+        } else if (timeoutMS == 0) { // infinite
+            // TODO-CSOT is this correct? alt might be 0=infinite here, but later is min'd against remaining?
             return alternativeTimeoutMS;
-        } else if (alternativeTimeoutMS == 0) {
-            return timeoutRemainingMS();
         } else {
-            return Math.min(timeoutRemainingMS(), alternativeTimeoutMS);
+            assertNotNull(timeout);
+            // TODO-CSOT refactor once above addressed? Unclear how to avoid extracting numbers.
+            long remaining = timeout.run(MILLISECONDS,
+                    () -> 0L,
+                    (ms) -> ms,
+                    () -> {
+                        throw createMongoTimeoutException("The operation timeout has expired.");
+                    });
+            if (alternativeTimeoutMS == 0) {
+                return remaining;
+            } else {
+                return Math.min(remaining, alternativeTimeoutMS);
+            }
         }
     }
 
@@ -201,11 +210,18 @@ public long getMaxAwaitTimeMS() {
 
     public long getMaxTimeMS() {
         long maxTimeMS = timeoutOrAlternative(timeoutSettings.getMaxTimeMS());
-        if (timeout == null || timeout.isInfinite()) {
-            return maxTimeMS;
-        }
+        return Timeout.run(timeout, MILLISECONDS,
+                // TODO-CSOT why does infinite translate to getMaxTime?
+                () -> maxTimeMS,
+                // TODO-CSOT why is timeout's ms not used anywhere?
+                (ms) -> calculateMaxTimeMs(maxTimeMS),
+                () -> calculateMaxTimeMs(maxTimeMS),
+                () -> maxTimeMS);
+    }
+
+    private long calculateMaxTimeMs(final long maxTimeMS) {
         if (minRoundTripTimeMS >= maxTimeMS) {
-            throw createMongoTimeoutException();
+            throw createMongoRoundTripTimeoutException();
         }
         return maxTimeMS - minRoundTripTimeMS;
     }
@@ -224,7 +240,7 @@ public long getWriteTimeoutMS() {
     }
 
     public int getConnectTimeoutMs() {
-        return (int) calculateMin(getTimeoutSettings().getConnectTimeoutMS());
+        return Math.toIntExact(calculateMin(getTimeoutSettings().getConnectTimeoutMS()));
     }
 
     public void resetTimeout() {
@@ -236,9 +252,14 @@ 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.run(timeout, NANOSECONDS,
+                () -> timeout,
+                (ms) -> calculateTimeout(timeoutSettings.getTimeoutMS()),
+                () -> calculateTimeout(timeoutSettings.getTimeoutMS()),
+                () -> timeout);
     }
 
     public TimeoutContext withAdditionalReadTimeout(final int additionalReadTimeout) {
@@ -254,14 +275,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{"
@@ -294,6 +307,7 @@ 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);
         }

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

@@ -138,9 +138,9 @@ public ServerTuple selectServer(final ServerSelector serverSelector, final Opera
                         serverSelector, currentDescription);
                 return serverTuple;
             }
-            if (computedServerSelectionTimeout.hasExpired()) {
+            Timeout.ifExistsAndExpired(computedServerSelectionTimeout, () -> {
                 throw createAndLogTimeoutException(operationContext.getId(), serverSelector, currentDescription);
-            }
+            });
             if (!selectionWaitingLogged) {
                 logServerSelectionWaiting(clusterId, operationContext.getId(), computedServerSelectionTimeout, serverSelector, currentDescription);
                 selectionWaitingLogged = true;
@@ -266,13 +266,13 @@ private boolean handleServerSelectionRequest(
                 }
             }
 
-            if (request.getTimeout().hasExpired()) {
+            boolean[] result = {false};
+            Timeout.ifExistsAndExpired(request.getTimeout(), () -> {
                 request.onResult(null, createAndLogTimeoutException(request.getOperationId(),
                         request.originalSelector, description));
-                return true;
-            }
-
-            return false;
+                result[0] = true;
+            });
+            return result[0];
         } catch (Exception e) {
             request.onResult(null, e);
             return true;
@@ -448,20 +448,21 @@ public void run() {
                 ClusterDescription curDescription = description;
 
                 Timeout timeout = Timeout.infinite();
-
+                boolean someWaitersNotSatisfied = false;
                 for (Iterator<ServerSelectionRequest> iter = waitQueue.iterator(); iter.hasNext();) {
                     ServerSelectionRequest currentRequest = iter.next();
                     if (handleServerSelectionRequest(currentRequest, currentPhase, curDescription)) {
                         iter.remove();
                     } else {
+                        someWaitersNotSatisfied = true;
                         timeout = timeout
                                 .orEarlier(currentRequest.getTimeout())
                                 .orEarlier(startMinWaitHeartbeatTimeout());
                     }
                 }
 
-                // if there are any waiters that were not satisfied, connect
-                if (!timeout.isInfinite()) {
+                // TODO-CSOT because min heartbeat cannot be infinite, infinite is being used to mark the second branch
+                if (someWaitersNotSatisfied) {
                     connect();
                 }
 
@@ -508,7 +509,8 @@ private static void logServerSelectionWaiting(
                     asList(
                             new Entry(OPERATION, null),
                             new Entry(OPERATION_ID, operationId),
-                            new Entry(REMAINING_TIME_MS, timeout.remainingOrNegativeForInfinite(MILLISECONDS)),
+                            new Entry(REMAINING_TIME_MS, timeout.run(MILLISECONDS, () -> "infinite", (ms) -> ms, () -> 0L)),
+                            // TODO-CSOT the above extracts values, but the alternative seems worse
                             new Entry(SELECTOR, serverSelector.toString()),
                             new Entry(TOPOLOGY_DESCRIPTION, clusterDescription.getShortDescription())),
                     "Waiting for server to become available for operation[ {}] with ID {}.[ Remaining time: {} ms.]"