PYTHON-5518: withTransaction API retries too frequently

graph.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import matplotlib.cm as mplcm
+import matplotlib.colors as colors
+import matplotlib.pyplot as plt
+import numpy as np
+
+threads = [1, 2, 4, 8, 16, 32, 64, 128, 256]
+
+data = {
+    "local_main": {
+        "avg": [[0.01], [0.03], [0.05], [0.09], [0.21], [0.38], [0.95], [4.05], [11.11]],
+        "p50": [[0.01], [0.02], [0.03], [0.05], [0.12], [0.25], [0.67], [3.22], [9.8]],
+        "p90": [[0.02], [0.05], [0.11], [0.23], [0.52], [0.94], [2.01], [8.88], [23.47]],
+        "p99": [[0.02], [0.05], [0.11], [0.23], [0.52], [0.94], [2.01], [8.88], [23.47]],
+        "p100": [[0.04], [0.21], [0.3], [0.58], [1.22], [2.59], [7.25], [17.03], [25.38]],
+    },
+    "local_original_1.5": {
+        "avg": [[0.01], [0.02], [0.04], [0.09], [0.2], [0.35], [0.65], [1.14], [2.21]],
+        "p50": [[0.01], [0.01], [0.01], [0.02], [0.06], [0.09], [0.2], [0.45], [1.91]],
+        "p90": [[0.02], [0.04], [0.11], [0.22], [0.49], [0.81], [1.87], [3.72], [4.76]],
+        "p99": [[0.02], [0.04], [0.11], [0.22], [0.49], [0.81], [1.87], [3.72], [4.76]],
+        "p100": [[0.04], [0.62], [1.11], [3.33], [4.71], [6.07], [5.64], [6.05], [6.76]],
+    },
+    "local_original_2": {
+        "avg": [[0.01], [0.02], [0.04], [0.09], [0.18], [0.39], [0.63], [1.23], [2.28]],
+        "p50": [[0.01], [0.01], [0.01], [0.02], [0.02], [0.06], [0.17], [0.41], [1.9]],
+        "p90": [[0.01], [0.02], [0.08], [0.21], [0.33], [1.24], [1.83], [3.82], [4.91]],
+        "p99": [[0.01], [0.02], [0.08], [0.21], [0.33], [1.24], [1.83], [3.82], [4.91]],
+        "p100": [[0.04], [1.3], [1.54], [3.07], [3.72], [5.55], [5.44], [6.42], [7.06]],
+    },
+    "local_server_algo": {
+        "avg": [[0.01], [0.02], [0.05], [0.1], [0.19], [0.36], [0.73], [1.23], [2.19]],
+        "p50": [[0.01], [0.01], [0.01], [0.01], [0.03], [0.09], [0.19], [0.59], [2.04]],
+        "p90": [[0.02], [0.04], [0.1], [0.22], [0.51], [1.07], [2.37], [3.58], [4.74]],
+        "p99": [[0.02], [0.04], [0.1], [0.22], [0.51], [1.07], [2.37], [3.58], [4.74]],
+        "p100": [[0.09], [0.65], [1.35], [2.87], [3.31], [4.4], [6.55], [5.84], [6.88]],
+    },
+}
+
+
+metrics = ["avg", "p90", "p100"]
+metric_titles = {
+    "avg": "Average Latency",
+    "p50": "p50 Latency",
+    "p90": "p90 Latency",
+    "p99": "p99 Latency",
+    "p100": "p100 (Max) Latency",
+}
+
+plt.figure(figsize=(16, 4 * len(metrics)))
+NUM_COLORS = len(data.keys()) + 1
+cm = plt.get_cmap("gist_rainbow")
+cNorm = colors.Normalize(vmin=0, vmax=NUM_COLORS - 1)
+scalarMap = mplcm.ScalarMappable(norm=cNorm, cmap=cm)
+for i, metric in enumerate(metrics, 1):
+    if metric in ["avg"]:
+        ax = plt.subplot(len(metrics), 2, i)
+    else:
+        ax = plt.subplot(len(metrics), 2, (i - 1) * (2) + 1)
+    ax.set_prop_cycle(color=[scalarMap.to_rgba(i) for i in range(NUM_COLORS)])
+    order = []
+    for label, vals in data.items():
+        if metric not in vals:
+            continue
+        arr = np.concatenate(np.around(np.array(vals[metric]), decimals=2))
+        order.append(plt.plot(threads, arr, "o-", label=label))
+
+    plt.title(metric_titles[metric])
+    plt.xscale("log", base=2)
+    plt.xlabel("Threads")
+    plt.ylabel("Seconds")
+    plt.xticks(threads, threads)
+    plt.grid(True, which="both", axis="x", linestyle="--", alpha=0.5)
+    plt.axhline(y=0, color="gray", linestyle="-")
+    if metric != "p90":
+        plt.legend().set_visible(False)
+    else:
+        plt.legend(loc=(1.01, 0.5), fontsize=8)
+
+plt.tight_layout()
+plt.show()

pymongo/asynchronous/client_session.py

@@ -472,15 +472,20 @@ def _max_time_expired_error(exc: PyMongoError) -> bool:
 # This limit is non-configurable and was chosen to be twice the 60 second
 # default value of MongoDB's `transactionLifetimeLimitSeconds` parameter.
 _WITH_TRANSACTION_RETRY_TIME_LIMIT = 120
-_BACKOFF_MAX = 1
-_BACKOFF_INITIAL = 0.050  # 50ms initial backoff
+_BACKOFF_MAX = 0.500  # 500ms max backoff
+_BACKOFF_INITIAL = 0.005  # 5ms initial backoff
 
 
 def _within_time_limit(start_time: float) -> bool:
     """Are we within the with_transaction retry limit?"""
     return time.monotonic() - start_time < _WITH_TRANSACTION_RETRY_TIME_LIMIT
 
 
+def _would_exceed_time_limit(start_time: float, backoff: float) -> bool:
+    """Is the backoff within the with_transaction retry limit?"""
+    return time.monotonic() + backoff - start_time >= _WITH_TRANSACTION_RETRY_TIME_LIMIT
+
+
 _T = TypeVar("_T")
 
 if TYPE_CHECKING:
@@ -708,10 +713,14 @@ async def callback(session, custom_arg, custom_kwarg=None):
         """
         start_time = time.monotonic()
         retry = 0
+        last_error: Optional[BaseException] = None
         while True:
             if retry:  # Implement exponential backoff on retry.
                 jitter = random.random()  # noqa: S311
-                backoff = jitter * min(_BACKOFF_INITIAL * (2**retry), _BACKOFF_MAX)
+                backoff = jitter * min(_BACKOFF_INITIAL * (1.5**retry), _BACKOFF_MAX)
+                if _would_exceed_time_limit(start_time, backoff):
+                    assert last_error is not None
+                    raise last_error
                 await asyncio.sleep(backoff)
             retry += 1
             await self.start_transaction(
@@ -721,6 +730,7 @@ async def callback(session, custom_arg, custom_kwarg=None):
                 ret = await callback(self)
             # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
             except BaseException as exc:
+                last_error = exc
                 if self.in_transaction:
                     await self.abort_transaction()
                 if (

pymongo/synchronous/client_session.py

@@ -470,15 +470,20 @@ def _max_time_expired_error(exc: PyMongoError) -> bool:
 # This limit is non-configurable and was chosen to be twice the 60 second
 # default value of MongoDB's `transactionLifetimeLimitSeconds` parameter.
 _WITH_TRANSACTION_RETRY_TIME_LIMIT = 120
-_BACKOFF_MAX = 1
-_BACKOFF_INITIAL = 0.050  # 50ms initial backoff
+_BACKOFF_MAX = 0.500  # 500ms max backoff
+_BACKOFF_INITIAL = 0.005  # 5ms initial backoff
 
 
 def _within_time_limit(start_time: float) -> bool:
     """Are we within the with_transaction retry limit?"""
     return time.monotonic() - start_time < _WITH_TRANSACTION_RETRY_TIME_LIMIT
 
 
+def _would_exceed_time_limit(start_time: float, backoff: float) -> bool:
+    """Is the backoff within the with_transaction retry limit?"""
+    return time.monotonic() + backoff - start_time >= _WITH_TRANSACTION_RETRY_TIME_LIMIT
+
+
 _T = TypeVar("_T")
 
 if TYPE_CHECKING:
@@ -706,17 +711,22 @@ def callback(session, custom_arg, custom_kwarg=None):
         """
         start_time = time.monotonic()
         retry = 0
+        last_error: Optional[BaseException] = None
         while True:
             if retry:  # Implement exponential backoff on retry.
                 jitter = random.random()  # noqa: S311
-                backoff = jitter * min(_BACKOFF_INITIAL * (2**retry), _BACKOFF_MAX)
+                backoff = jitter * min(_BACKOFF_INITIAL * (1.5**retry), _BACKOFF_MAX)
+                if _would_exceed_time_limit(start_time, backoff):
+                    assert last_error is not None
+                    raise last_error
                 time.sleep(backoff)
             retry += 1
             self.start_transaction(read_concern, write_concern, read_preference, max_commit_time_ms)
             try:
                 ret = callback(self)
             # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
             except BaseException as exc:
+                last_error = exc
                 if self.in_transaction:
                     self.abort_transaction()
                 if (

test/asynchronous/test_transactions.py

@@ -16,7 +16,9 @@
 from __future__ import annotations
 
 import asyncio
+import random
 import sys
+import time
 from io import BytesIO
 from test.asynchronous.utils_spec_runner import AsyncSpecRunner
 
@@ -441,7 +443,7 @@ async def set_fail_point(self, command_args):
             await self.configure_fail_point(client, command_args)
 
     @async_client_context.require_transactions
-    async def test_callback_raises_custom_error(self):
+    async def test_1_callback_raises_custom_error(self):
         class _MyException(Exception):
             pass
 
@@ -453,7 +455,7 @@ async def raise_error(_):
                 await s.with_transaction(raise_error)
 
     @async_client_context.require_transactions
-    async def test_callback_returns_value(self):
+    async def test_2_callback_returns_value(self):
         async def callback(_):
             return "Foo"
 
@@ -481,7 +483,7 @@ def callback(_):
             self.assertEqual(await s.with_transaction(callback), "Foo")
 
     @async_client_context.require_transactions
-    async def test_callback_not_retried_after_timeout(self):
+    async def test_3_1_callback_not_retried_after_timeout(self):
         listener = OvertCommandListener()
         client = await self.async_rs_client(event_listeners=[listener])
         coll = client[self.db.name].test
@@ -509,7 +511,7 @@ async def callback(session):
 
     @async_client_context.require_test_commands
     @async_client_context.require_transactions
-    async def test_callback_not_retried_after_commit_timeout(self):
+    async def test_3_2_callback_not_retried_after_commit_timeout(self):
         listener = OvertCommandListener()
         client = await self.async_rs_client(event_listeners=[listener])
         coll = client[self.db.name].test
@@ -543,7 +545,7 @@ async def callback(session):
 
     @async_client_context.require_test_commands
     @async_client_context.require_transactions
-    async def test_commit_not_retried_after_timeout(self):
+    async def test_3_3_commit_not_retried_after_timeout(self):
         listener = OvertCommandListener()
         client = await self.async_rs_client(event_listeners=[listener])
         coll = client[self.db.name].test
@@ -613,6 +615,72 @@ async def callback(session):
             await s.with_transaction(callback)
             self.assertFalse(s.in_transaction)
 
+    @async_client_context.require_test_commands
+    @async_client_context.require_transactions
+    async def test_4_retry_backoff_is_enforced(self):
+        client = async_client_context.client
+        coll = client[self.db.name].test
+        # patch random to make it deterministic -- once to effectively have
+        # no backoff and the second time with "max" backoff (always waiting the longest
+        # possible time)
+        _original_random_random = random.random
+
+        def always_one():
+            return 1
+
+        def always_zero():
+            return 0
+
+        random.random = always_zero
+        # set fail point to trigger transaction failure and trigger backoff
+        await self.set_fail_point(
+            {
+                "configureFailPoint": "failCommand",
+                "mode": {"times": 13},
+                "data": {
+                    "failCommands": ["commitTransaction"],
+                    "errorCode": 251,
+                },
+            }
+        )
+        self.addAsyncCleanup(
+            self.set_fail_point, {"configureFailPoint": "failCommand", "mode": "off"}
+        )
+
+        async def callback(session):
+            await coll.insert_one({}, session=session)
+
+        start = time.monotonic()
+        async with self.client.start_session() as s:
+            await s.with_transaction(callback)
+        end = time.monotonic()
+        no_backoff_time = end - start
+
+        random.random = always_one
+        # set fail point to trigger transaction failure and trigger backoff
+        await self.set_fail_point(
+            {
+                "configureFailPoint": "failCommand",
+                "mode": {
+                    "times": 13
+                },  # sufficiently high enough such that the time effect of backoff is noticeable
+                "data": {
+                    "failCommands": ["commitTransaction"],
+                    "errorCode": 251,
+                },
+            }
+        )
+        self.addAsyncCleanup(
+            self.set_fail_point, {"configureFailPoint": "failCommand", "mode": "off"}
+        )
+        start = time.monotonic()
+        async with self.client.start_session() as s:
+            await s.with_transaction(callback)
+        end = time.monotonic()
+        self.assertLess(abs(end - start - (no_backoff_time + 2.2)), 1)  # sum of 13 backoffs is 2.2
+
+        random.random = _original_random_random
+
 
 class TestOptionsInsideTransactionProse(AsyncTransactionsBase):
     @async_client_context.require_transactions