JAVA-5950 Update Transactions Convenient API with exponential backoff on retries (#1899)

Author: nhachicha

driver-benchmarks/src/main/com/mongodb/benchmark/benchmarks/RawBsonArrayEncodingBenchmark.java

@@ -17,7 +17,8 @@
 
 package com.mongodb.benchmark.benchmarks;
 
-import org.bson.BsonArray;import org.bson.BsonDocument;
+import org.bson.BsonArray;
+import org.bson.BsonDocument;
 import org.bson.RawBsonDocument;
 import org.bson.codecs.BsonDocumentCodec;
 
@@ -52,4 +53,4 @@ public void setUp() throws IOException {
     public int getBytesPerRun() {
         return documentBytes.length * NUM_INTERNAL_ITERATIONS;
     }
-}
+}

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

@@ -0,0 +1,43 @@
+/*
+ * 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;
+
+import com.mongodb.lang.Nullable;
+
+/**
+ * An exception indicating that the convenient transactions API ({@code withTransaction})
+ * exceeded its overall timeout while retrying the user-supplied callback or the commit loop.
+ * The last encountered error (if any) is attached as the
+ * {@linkplain Throwable#getCause() cause}.
+ *
+ * @since 5.7
+ * @mongodb.driver.manual core/transactions-in-applications/#callback-api withTransaction
+ */
+public final class WithTransactionTimeoutException extends MongoClientException {
+
+    private static final long serialVersionUID = 1L;
+
+    /**
+     * Construct a new instance
+     * @param message the message
+     * @param cause the cause
+     * @since 5.7
+     */
+    public WithTransactionTimeoutException(final String message, @Nullable final Throwable cause) {
+        super(message, cause);
+    }
+}

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

@@ -113,10 +113,6 @@ public TimeoutContext(final TimeoutSettings timeoutSettings) {
         this(false, timeoutSettings, startTimeout(timeoutSettings.getTimeoutMS()));
     }
 
-    private TimeoutContext(final TimeoutSettings timeoutSettings, @Nullable final Timeout timeout) {
-        this(false, timeoutSettings, timeout);
-    }
-
     private TimeoutContext(final boolean isMaintenanceContext,
                            final TimeoutSettings timeoutSettings,
                            @Nullable final Timeout timeout) {
@@ -181,6 +177,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;
@@ -389,11 +386,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{"

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;
 

driver-core/src/main/com/mongodb/internal/time/ExponentialBackoff.java

@@ -0,0 +1,76 @@
+/*
+ * 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;
+
+import com.mongodb.internal.VisibleForTesting;
+
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.DoubleSupplier;
+
+import static com.mongodb.assertions.Assertions.assertTrue;
+import static com.mongodb.internal.VisibleForTesting.AccessModifier.PRIVATE;
+
+/**
+ * Provides exponential backoff calculations with jitter for retry scenarios.
+ */
+public final class ExponentialBackoff {
+
+    private static final double TRANSACTION_BASE_MS = 5.0;
+    @VisibleForTesting(otherwise = PRIVATE)
+    static final double TRANSACTION_MAX_MS = 500.0;
+    private static final double TRANSACTION_GROWTH = 1.5;
+
+    // TODO-JAVA-6079
+    private static DoubleSupplier testJitterSupplier = null;
+
+    private ExponentialBackoff() {
+    }
+
+    /**
+     * Calculate the backoff in milliseconds for transaction retries.
+     *
+     * @param attemptNumber attempt number > 0
+     * @return The calculated backoff in milliseconds.
+     */
+    public static long calculateTransactionBackoffMs(final int attemptNumber) {
+        assertTrue(attemptNumber > 0, "Attempt number must be at least 1 (1-based) in the context of transaction backoff calculation");
+        double jitter = testJitterSupplier != null
+                ? testJitterSupplier.getAsDouble()
+                : ThreadLocalRandom.current().nextDouble();
+        return Math.round(jitter * Math.min(
+                TRANSACTION_BASE_MS * Math.pow(TRANSACTION_GROWTH, attemptNumber - 1),
+                TRANSACTION_MAX_MS));
+    }
+
+    /**
+     * Set a custom jitter supplier for testing purposes.
+     *
+     * @param supplier A DoubleSupplier that returns values in [0, 1] range.
+     */
+    @VisibleForTesting(otherwise = PRIVATE)
+    public static void setTestJitterSupplier(final DoubleSupplier supplier) {
+        testJitterSupplier = supplier;
+    }
+
+    /**
+     * Clear the test jitter supplier, reverting to default ThreadLocalRandom behavior.
+     */
+    @VisibleForTesting(otherwise = PRIVATE)
+    public static void clearTestJitterSupplier() {
+        testJitterSupplier = null;
+    }
+}

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());
     }
 }

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();
+    }
+}

driver-core/src/main/com/mongodb/internal/time/TimePoint.java

@@ -62,14 +62,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());
     }
 
     /**

driver-core/src/test/unit/com/mongodb/internal/time/ExponentialBackoffTest.java

@@ -0,0 +1,82 @@
+/*
+ * 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;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.function.DoubleSupplier;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class ExponentialBackoffTest {
+    /**
+     * Expected {@linkplain ExponentialBackoff#calculateTransactionBackoffMs(int) backoffs} with 1.0 as
+     * {@link ExponentialBackoff#setTestJitterSupplier(DoubleSupplier) jitter}.
+     */
+    private static final double[] EXPECTED_BACKOFFS_MAX_VALUES = {5.0, 7.5, 11.25, 16.875, 25.3125, 37.96875, 56.953125, 85.4296875, 128.14453125,
+            192.21679688, 288.32519531, 432.48779297, 500.0};
+
+    @Test
+    void testCalculateTransactionBackoffMs() {
+        for (int attemptNumber = 1; attemptNumber <= EXPECTED_BACKOFFS_MAX_VALUES.length; attemptNumber++) {
+            long backoff = ExponentialBackoff.calculateTransactionBackoffMs(attemptNumber);
+            long expectedBackoff = Math.round(EXPECTED_BACKOFFS_MAX_VALUES[attemptNumber - 1]);
+            assertTrue(backoff >= 0 && backoff <= expectedBackoff,
+                    String.format("Attempt %d: backoff should be between 0 ms and %d ms, got: %d", attemptNumber,
+                            expectedBackoff, backoff));
+        }
+    }
+
+    @Test
+    void testCalculateTransactionBackoffMsRespectsMaximum() {
+        for (int attemptNumber = 1; attemptNumber < EXPECTED_BACKOFFS_MAX_VALUES.length * 2; attemptNumber++) {
+            long backoff = ExponentialBackoff.calculateTransactionBackoffMs(attemptNumber);
+            assertTrue(backoff >= 0 && backoff <= ExponentialBackoff.TRANSACTION_MAX_MS,
+                String.format("Attempt %d: backoff should be capped at %f ms, got: %d ms",
+                        attemptNumber, ExponentialBackoff.TRANSACTION_MAX_MS, backoff));
+        }
+    }
+
+    @Test
+    void testCalculateTransactionBackoffMsWithJitterOne() {
+        ExponentialBackoff.setTestJitterSupplier(() -> 1.0);
+        try {
+            for (int attemptNumber = 1; attemptNumber <= EXPECTED_BACKOFFS_MAX_VALUES.length; attemptNumber++) {
+                long backoff = ExponentialBackoff.calculateTransactionBackoffMs(attemptNumber);
+                long expected = Math.round(EXPECTED_BACKOFFS_MAX_VALUES[attemptNumber - 1]);
+                assertEquals(expected, backoff,
+                        String.format("Attempt %d: with jitter=1.0, backoff should be %d ms", attemptNumber, expected));
+            }
+        } finally {
+            ExponentialBackoff.clearTestJitterSupplier();
+        }
+    }
+
+    @Test
+    void testCalculateTransactionBackoffMsWithJitterZero() {
+        ExponentialBackoff.setTestJitterSupplier(() -> 0.0);
+        try {
+            for (int attemptNumber = 1; attemptNumber <= EXPECTED_BACKOFFS_MAX_VALUES.length; attemptNumber++) {
+                long backoff = ExponentialBackoff.calculateTransactionBackoffMs(attemptNumber);
+                assertEquals(0, backoff, String.format("Attempt %d: with jitter=0, backoff should always be 0 ms", attemptNumber));
+            }
+        } finally {
+            ExponentialBackoff.clearTestJitterSupplier();
+        }
+    }
+}

driver-scala/src/main/scala/org/mongodb/scala/package.scala

@@ -407,6 +407,14 @@ package object scala extends ClientSessionImplicits with ObservableImplicits wit
    */
   type MongoOperationTimeoutException = com.mongodb.MongoOperationTimeoutException
 
+  /**
+   * An exception indicating that the convenient transactions API (`withTransaction`) exceeded its overall timeout
+   * while retrying the user-supplied callback or the commit loop.
+   *
+   * @since 5.7
+   */
+  type WithTransactionTimeoutException = com.mongodb.WithTransactionTimeoutException
+
   /**
    * An exception indicating a failure to apply the write concern to the requested write operation
    *