Fix/improve ClientSessionImpl, remove ClientSessionClock

stIncMale · nhachicha · commit 88071e03710a · 2026-02-25T13:00:51.000Z
diff --git a/driver-core/src/main/com/mongodb/internal/TimeoutContext.java b/driver-core/src/main/com/mongodb/internal/TimeoutContext.java
@@ -168,6 +168,10 @@ public Timeout timeoutIncludingRoundTrip() {
         return timeout == null ? null : timeout.shortenBy(minRoundTripTimeMS, MILLISECONDS);
     }
 
+    public Timeout timeoutOrAlternative(final Timeout alternative) {
+        return timeout == null ? alternative : timeout;
+    }
+
     /**
      * Returns the remaining {@code timeoutMS} if set or the {@code alternativeTimeoutMS}.
      *
@@ -176,6 +180,7 @@ public Timeout timeoutIncludingRoundTrip() {
      * @param alternativeTimeoutMS the alternative timeout.
      * @return timeout to use.
      */
+    @VisibleForTesting(otherwise = PRIVATE)
     public long timeoutOrAlternative(final long alternativeTimeoutMS) {
         if (timeout == null) {
             return alternativeTimeoutMS;
@@ -380,11 +385,6 @@ public TimeoutContext withAdditionalReadTimeout(final int additionalReadTimeout)
         return new TimeoutContext(timeoutSettings.withReadTimeoutMS(newReadTimeout > 0 ? newReadTimeout : Long.MAX_VALUE));
     }
 
-    // Creates a copy of the timeout context that can be reset without resetting the original.
-    public TimeoutContext copyTimeoutContext() {
-        return new TimeoutContext(getTimeoutSettings(), getTimeout());
-    }
-
     @Override
     public String toString() {
         return "TimeoutContext{"
diff --git a/driver-core/src/main/com/mongodb/internal/connection/Time.java b/driver-core/src/main/com/mongodb/internal/connection/Time.java
@@ -20,7 +20,11 @@
  * To enable unit testing of classes that rely on System.nanoTime
  *
  * <p>This class is not part of the public API and may be removed or changed at any time</p>
+ *
+ * @deprecated Use {@link com.mongodb.internal.time.SystemNanoTime} in production code,
+ * and {@code Mockito.mockStatic} in test code to tamper with it.
  */
+@Deprecated
 public final class Time {
     static final long CONSTANT_TIME = 42;
 
diff --git a/driver-core/src/main/com/mongodb/internal/time/StartTime.java b/driver-core/src/main/com/mongodb/internal/time/StartTime.java
@@ -59,6 +59,6 @@ public interface StartTime {
      * @return a StartPoint, as of now
      */
     static StartTime now() {
-        return TimePoint.at(System.nanoTime());
+        return TimePoint.at(SystemNanoTime.get());
     }
 }
diff --git a/driver-core/src/main/com/mongodb/internal/time/SystemNanoTime.java b/driver-core/src/main/com/mongodb/internal/time/SystemNanoTime.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.time;
+
+/**
+ * Avoid using this class directly and prefer using other program elements from {@link com.mongodb.internal.time}, if possible.
+ * <p>
+ * We do not use {@link System#nanoTime()} directly in the rest of the {@link com.mongodb.internal.time} package,
+ * and use {@link SystemNanoTime#get()} instead because we need to tamper with it via {@code Mockito.mockStatic},
+ * and mocking methods of {@link System} class is both impossible and unwise.
+ */
+public final class SystemNanoTime {
+    private SystemNanoTime() {
+    }
+
+    public static long get() {
+        return System.nanoTime();
+    }
+}
diff --git a/driver-core/src/main/com/mongodb/internal/time/TimePoint.java b/driver-core/src/main/com/mongodb/internal/time/TimePoint.java
@@ -61,14 +61,14 @@ static TimePoint at(@Nullable final Long nanos) {
 
     @VisibleForTesting(otherwise = PRIVATE)
     long currentNanos() {
-        return System.nanoTime();
+        return SystemNanoTime.get();
     }
 
     /**
      * Returns the current {@link TimePoint}.
      */
     static TimePoint now() {
-        return at(System.nanoTime());
+        return at(SystemNanoTime.get());
     }
 
     /**
diff --git a/driver-sync/src/main/com/mongodb/client/internal/ClientSessionClock.java b/driver-sync/src/main/com/mongodb/client/internal/ClientSessionClock.java
diff --git a/driver-sync/src/main/com/mongodb/client/internal/ClientSessionImpl.java b/driver-sync/src/main/com/mongodb/client/internal/ClientSessionImpl.java
@@ -39,8 +39,12 @@
 import com.mongodb.internal.session.BaseClientSessionImpl;
 import com.mongodb.internal.session.ServerSessionPool;
 import com.mongodb.internal.time.ExponentialBackoff;
+import com.mongodb.internal.time.Timeout;
 import com.mongodb.lang.Nullable;
 
+import java.util.concurrent.TimeUnit;
+import java.util.function.BooleanSupplier;
+
 import static com.mongodb.MongoException.TRANSIENT_TRANSACTION_ERROR_LABEL;
 import static com.mongodb.MongoException.UNKNOWN_TRANSACTION_COMMIT_RESULT_LABEL;
 import static com.mongodb.assertions.Assertions.assertNotNull;
@@ -51,7 +55,7 @@
 
 final class ClientSessionImpl extends BaseClientSessionImpl implements ClientSession {
 
-    private static final int MAX_RETRY_TIME_LIMIT_MS = 120000;
+    private static final long MAX_RETRY_TIME_LIMIT_MS = 120000;
 
     private final OperationExecutor operationExecutor;
     private TransactionState transactionState = TransactionState.NONE;
@@ -251,21 +255,23 @@ public <T> T withTransaction(final TransactionBody<T> transactionBody) {
     @Override
     public <T> T withTransaction(final TransactionBody<T> transactionBody, final TransactionOptions options) {
         notNull("transactionBody", transactionBody);
-        long startTime = ClientSessionClock.INSTANCE.now();
         TimeoutContext withTransactionTimeoutContext = createTimeoutContext(options);
-        ExponentialBackoff transactionBackoff = ExponentialBackoff.TRANSACTION;
+        Timeout withTransactionTimeout = withTransactionTimeoutContext.timeoutOrAlternative(
+                assertNotNull(TimeoutContext.startTimeout(MAX_RETRY_TIME_LIMIT_MS)));
+        BooleanSupplier withTransactionTimeoutExpired = () -> withTransactionTimeout.call(TimeUnit.MILLISECONDS,
+                () -> false, ms -> false, () -> true);
         int transactionAttempt = 0;
         MongoException lastError = null;
 
         try {
             outer:
             while (true) {
                 if (transactionAttempt > 0) {
-                    backoff(transactionBackoff, transactionAttempt, startTime, lastError);
+                    backoff(transactionAttempt, withTransactionTimeout, assertNotNull(lastError));
                 }
                 T retVal;
                 try {
-                    startTransaction(options, withTransactionTimeoutContext.copyTimeoutContext());
+                    startTransaction(options, withTransactionTimeoutContext);
                     transactionAttempt++;
 
                     if (transactionSpan != null) {
@@ -281,7 +287,7 @@ public <T> T withTransaction(final TransactionBody<T> transactionBody, final Tra
                         if (!(e instanceof MongoOperationTimeoutException)) {
                             MongoException exceptionToHandle = OperationHelper.unwrap((MongoException) e);
                             if (exceptionToHandle.hasErrorLabel(TRANSIENT_TRANSACTION_ERROR_LABEL)
-                                    && ClientSessionClock.INSTANCE.now() - startTime < MAX_RETRY_TIME_LIMIT_MS) {
+                                    && !withTransactionTimeoutExpired.getAsBoolean()) {
                                 if (transactionSpan != null) {
                                     transactionSpan.spanFinalizing(false);
                                 }
@@ -297,9 +303,10 @@ public <T> T withTransaction(final TransactionBody<T> transactionBody, final Tra
                             commitTransaction(false);
                             break;
                         } catch (MongoException e) {
+                            lastError = e;
                             clearTransactionContextOnError(e);
                             if (!(e instanceof MongoOperationTimeoutException)
-                                    && ClientSessionClock.INSTANCE.now() - startTime < MAX_RETRY_TIME_LIMIT_MS) {
+                                    && !withTransactionTimeoutExpired.getAsBoolean()) {
                                 applyMajorityWriteConcernToTransactionOptions();
 
                                 if (!(e instanceof MongoExecutionTimeoutException)
@@ -309,7 +316,6 @@ public <T> T withTransaction(final TransactionBody<T> transactionBody, final Tra
                                     if (transactionSpan != null) {
                                         transactionSpan.spanFinalizing(true);
                                     }
-                                    lastError = e;
                                     continue outer;
                                 }
                             }
@@ -374,15 +380,12 @@ private TimeoutContext createTimeoutContext(final TransactionOptions transaction
                 operationExecutor.getTimeoutSettings()));
     }
 
-    private static void backoff(final ExponentialBackoff exponentialBackoff, final int transactionAttempt, final long startTime,
-            final MongoException lastError) {
-        long backoffMs = exponentialBackoff.calculateDelayBeforeNextRetryMs(transactionAttempt - 1);
-        if (ClientSessionClock.INSTANCE.now() + backoffMs - startTime >= MAX_RETRY_TIME_LIMIT_MS) {
-            if (lastError != null) {
-                throw lastError;
-            }
-            throw new MongoClientException("Transaction retry timeout exceeded");
-        }
+    private static void backoff(final int transactionAttempt,
+            final Timeout withTransactionTimeout, final MongoException lastError) {
+        long backoffMs = ExponentialBackoff.TRANSACTION.calculateDelayBeforeNextRetryMs(transactionAttempt - 1);
+        withTransactionTimeout.shortenBy(backoffMs, TimeUnit.MILLISECONDS).onExpired(() -> {
+            throw lastError;
+        });
         try {
             if (backoffMs > 0) {
                 Thread.sleep(backoffMs);
diff --git a/driver-sync/src/test/functional/com/mongodb/client/ClientSideOperationTimeoutTest.java b/driver-sync/src/test/functional/com/mongodb/client/ClientSideOperationTimeoutTest.java
@@ -23,7 +23,6 @@
 
 import static org.junit.jupiter.api.Assumptions.assumeFalse;
 
-
 // See https://github.com/mongodb/specifications/tree/master/source/client-side-operation-timeout/tests
 public class ClientSideOperationTimeoutTest extends UnifiedSyncTest {
 
diff --git a/driver-sync/src/test/functional/com/mongodb/client/WithTransactionProseTest.java b/driver-sync/src/test/functional/com/mongodb/client/WithTransactionProseTest.java
@@ -20,17 +20,22 @@
 import com.mongodb.MongoClientException;
 import com.mongodb.MongoException;
 import com.mongodb.TransactionOptions;
-import com.mongodb.client.internal.ClientSessionClock;
 import com.mongodb.client.model.Sorts;
 import com.mongodb.internal.time.ExponentialBackoff;
+import com.mongodb.internal.time.StartTime;
+import com.mongodb.internal.time.SystemNanoTime;
 import org.bson.BsonDocument;
 import org.bson.Document;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 
+import java.time.Duration;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 
 import static com.mongodb.ClusterFixture.TIMEOUT;
 import static com.mongodb.ClusterFixture.isDiscoverableReplicaSet;
@@ -44,8 +49,7 @@
 
 // See https://github.com/mongodb/specifications/blob/master/source/transactions-convenient-api/tests/README.md#prose-tests
 public class WithTransactionProseTest extends DatabaseTestCase {
-    private static final long START_TIME_MS = 1L;
-    private static final long ERROR_GENERATING_INTERVAL = 121000L;
+    private static final Duration ERROR_GENERATING_INTERVAL = Duration.ofSeconds(120);
 
     @BeforeEach
     @Override
@@ -66,7 +70,7 @@ public void setUp() {
     public void testCallbackRaisesCustomError() {
         final String exceptionMessage = "NotTransientOrUnknownError";
         try (ClientSession session = client.startSession()) {
-            session.withTransaction((TransactionBody<Void>) () -> {
+            session.withTransaction(() -> {
                 throw new MongoException(exceptionMessage);
             });
             // should not get here
@@ -101,13 +105,13 @@ public void testRetryTimeoutEnforcedTransientTransactionError() {
         final String errorMessage = "transient transaction error";
 
         try (ClientSession session = client.startSession()) {
-            ClientSessionClock.INSTANCE.setTime(START_TIME_MS);
-            session.withTransaction((TransactionBody<Void>) () -> {
-                ClientSessionClock.INSTANCE.setTime(ERROR_GENERATING_INTERVAL);
-                MongoException e = new MongoException(112, errorMessage);
-                e.addLabel(MongoException.TRANSIENT_TRANSACTION_ERROR_LABEL);
-                throw e;
-            });
+            doWithSystemNanoTimeHandle(systemNanoTimeHandle ->
+                session.withTransaction(() -> {
+                    systemNanoTimeHandle.setRelativeToStart(ERROR_GENERATING_INTERVAL);
+                    MongoException e = new MongoException(112, errorMessage);
+                    e.addLabel(MongoException.TRANSIENT_TRANSACTION_ERROR_LABEL);
+                    throw e;
+                }));
             fail("Test should have thrown an exception.");
         } catch (Exception e) {
             assertEquals(errorMessage, e.getMessage());
@@ -127,12 +131,12 @@ public void testRetryTimeoutEnforcedUnknownTransactionCommit() {
                         + "'data': {'failCommands': ['commitTransaction'], 'errorCode': 91, 'closeConnection': false}}"));
 
         try (ClientSession session = client.startSession()) {
-            ClientSessionClock.INSTANCE.setTime(START_TIME_MS);
-            session.withTransaction((TransactionBody<Void>) () -> {
-                ClientSessionClock.INSTANCE.setTime(ERROR_GENERATING_INTERVAL);
-                collection.insertOne(session, new Document("_id", 2));
-                return null;
-            });
+            doWithSystemNanoTimeHandle(systemNanoTimeHandle ->
+                session.withTransaction(() -> {
+                    systemNanoTimeHandle.setRelativeToStart(ERROR_GENERATING_INTERVAL);
+                    collection.insertOne(session, new Document("_id", 2));
+                    return null;
+                }));
             fail("Test should have thrown an exception.");
         } catch (Exception e) {
             assertEquals(91, ((MongoException) e).getCode());
@@ -156,12 +160,12 @@ public void testRetryTimeoutEnforcedTransientTransactionErrorOnCommit() {
                         + "'errmsg': 'Transaction 0 has been aborted', 'closeConnection': false}}"));
 
         try (ClientSession session = client.startSession()) {
-            ClientSessionClock.INSTANCE.setTime(START_TIME_MS);
-            session.withTransaction((TransactionBody<Void>) () -> {
-                ClientSessionClock.INSTANCE.setTime(ERROR_GENERATING_INTERVAL);
-                collection.insertOne(session, Document.parse("{ _id : 1 }"));
-                return null;
-            });
+            doWithSystemNanoTimeHandle(systemNanoTimeHandle ->
+                session.withTransaction(() -> {
+                    systemNanoTimeHandle.setRelativeToStart(ERROR_GENERATING_INTERVAL);
+                    collection.insertOne(session, Document.parse("{ _id : 1 }"));
+                    return null;
+                }));
             fail("Test should have thrown an exception.");
         } catch (Exception e) {
             assertEquals(251, ((MongoException) e).getCode());
@@ -224,9 +228,9 @@ public void testRetryBackoffIsEnforced() throws InterruptedException {
         long noBackoffTime;
         try (ClientSession session = client.startSession();
              FailPoint ignored = FailPoint.enable(failPointDocument, getPrimary())) {
-            long startNanos = System.nanoTime();
+            StartTime startTime = StartTime.now();
             session.withTransaction(() -> collection.insertOne(session, Document.parse("{}")));
-            noBackoffTime = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos);
+            noBackoffTime = startTime.elapsed().toMillis();
         } finally {
             // Clear the test jitter supplier to avoid affecting other tests
             ExponentialBackoff.clearTestJitterSupplier();
@@ -241,9 +245,9 @@ public void testRetryBackoffIsEnforced() throws InterruptedException {
         long withBackoffTime;
         try (ClientSession session = client.startSession();
              FailPoint ignored = FailPoint.enable(failPointDocument, getPrimary())) {
-            long startNanos = System.nanoTime();
+            StartTime startTime = StartTime.now();
             session.withTransaction(() -> collection.insertOne(session, Document.parse("{}")));
-            withBackoffTime = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos);
+            withBackoffTime = startTime.elapsed().toMillis();
         } finally {
             ExponentialBackoff.clearTestJitterSupplier();
         }
@@ -278,7 +282,18 @@ public void testExponentialBackoffOnTransientError() throws InterruptedException
         }
     }
 
-    private boolean canRunTests() {
+    private static boolean canRunTests() {
         return isSharded() || isDiscoverableReplicaSet();
     }
+
+    private static void doWithSystemNanoTimeHandle(final Consumer<SystemNanoTimeHandle> action) {
+        long startNanos = SystemNanoTime.get();
+        try (MockedStatic<SystemNanoTime> mockedStaticSystem = Mockito.mockStatic(SystemNanoTime.class)) {
+            action.accept(change -> mockedStaticSystem.when(SystemNanoTime::get).thenReturn(startNanos + change.toNanos()));
+        }
+    }
+
+    private interface SystemNanoTimeHandle {
+        void setRelativeToStart(Duration change);
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -59,6 +59,6 @@ public interface StartTime {`
`59`	`59`	`* @return a StartPoint, as of now`
`60`	`60`	`*/`
`61`	`61`	`static StartTime now() {`
`62`		`- return TimePoint.at(System.nanoTime());`
	`62`	`+ return TimePoint.at(SystemNanoTime.get());`
`63`	`63`	`}`
`64`	`64`	`}`
Original file line number	Diff line number	Diff line change
`@@ -61,14 +61,14 @@ static TimePoint at(@Nullable final Long nanos) {`
`61`	`61`
`62`	`62`	`@VisibleForTesting(otherwise = PRIVATE)`
`63`	`63`	`long currentNanos() {`
`64`		`- return System.nanoTime();`
	`64`	`+ return SystemNanoTime.get();`
`65`	`65`	`}`
`66`	`66`
`67`	`67`	`/**`
`68`	`68`	`* Returns the current {@link TimePoint}.`
`69`	`69`	`*/`
`70`	`70`	`static TimePoint now() {`
`71`		`- return at(System.nanoTime());`
	`71`	`+ return at(SystemNanoTime.get());`
`72`	`72`	`}`
`73`	`73`
`74`	`74`	`/**`