PYTHON-5885 Client Backpressure with baseBackoffMS (#2877)

NoahStapp · web-flow · commit 7d1af1cf7e54 · 2026-07-24T13:32:51.000-04:00
diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -1,13 +1,16 @@
 Changelog
 =========
 
-Changes in Version 4.18.0
--------------------------
+Changes in Version 4.18.0 (2026/XX/XX)
+--------------------------------------
+
+PyMongo 4.18 brings a number of changes including:
 
 - Improved TLS connection performance by reusing TLS sessions across connections
   to the same server, avoiding a full handshake on each new connection.
   Session resumption is supported on all Python versions for synchronous clients
   and on Python 3.11+ for async clients.
+- Improved performance for MongoDB 9.0's Intelligent Workload Management (IWM) by only retrying overload errors when doing so is expected to not worsen server conditions.
 - Redacted potentially sensitive authentication mechanism properties, including
   AWS session tokens, from the representations of
   :class:`~pymongo.synchronous.mongo_client.MongoClient` and
diff --git a/pymongo/asynchronous/helpers.py b/pymongo/asynchronous/helpers.py
@@ -26,6 +26,7 @@
 from typing import (
     Any,
     Callable,
+    Optional,
     TypeVar,
     cast,
 )
@@ -84,10 +85,12 @@ async def inner(*args: Any, **kwargs: Any) -> Any:
 
 
 def _backoff(
-    attempt: int, initial_delay: float = _BACKOFF_INITIAL, max_delay: float = _BACKOFF_MAX
+    attempt: int,
+    base_backoff: float,
+    max_delay: float = _BACKOFF_MAX,
 ) -> float:
     jitter = random.random()  # noqa: S311
-    return jitter * min(initial_delay * (2**attempt), max_delay)
+    return jitter * min(base_backoff * (2**attempt), max_delay)
 
 
 class _RetryPolicy:
@@ -103,9 +106,13 @@ def __init__(
         self.backoff_initial = backoff_initial
         self.backoff_max = backoff_max
 
-    def backoff(self, attempt: int) -> float:
-        """Return the backoff duration for the given attempt."""
-        return _backoff(max(0, attempt - 1), self.backoff_initial, self.backoff_max)
+    def backoff(self, attempt: int, base_backoff: Optional[float] = None) -> float:
+        """Return the actual backoff duration for the given attempt and base backoff."""
+        return _backoff(
+            max(0, attempt),
+            self.backoff_initial if base_backoff is None or base_backoff < 0 else base_backoff,
+            self.backoff_max,
+        )
 
     async def should_retry(self, attempt: int, delay: float) -> bool:
         """Return if we have retry attempts remaining and the next backoff would not exceed a timeout."""
diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py
@@ -2807,6 +2807,7 @@ def __init__(
         self._attempt_number = 0
         self._is_run_command = is_run_command
         self._is_aggregate_write = is_aggregate_write
+        self._base_backoff_ms: Optional[float] = None
 
     async def run(self) -> T:
         """Runs the supplied func() and attempts a retry
@@ -2856,26 +2857,29 @@ async def run(self) -> T:
 
                 # Execute specialized catch on read
                 if self._is_read:
-                    if isinstance(exc, (ConnectionFailure, OperationFailure)):
+                    if isinstance(exc_to_check, (ConnectionFailure, OperationFailure)):
                         # ConnectionFailures do not supply a code property
-                        exc_code = getattr(exc, "code", None)
-                        overloaded = exc.has_error_label("SystemOverloadedError")
+                        exc_code = getattr(exc_to_check, "code", None)
+                        overloaded = exc_to_check.has_error_label("SystemOverloadedError")
                         if overloaded:
                             self._max_retries = self._client.options.max_adaptive_retries
-                        always_retryable = exc.has_error_label("RetryableError") and overloaded
+                            self._base_backoff_ms = getattr(exc_to_check, "_base_backoff_ms", None)
+                        always_retryable = (
+                            exc_to_check.has_error_label("RetryableError") and overloaded
+                        )
                         if not self._client.options.retry_reads or (
                             not always_retryable
                             and (
                                 self._is_not_eligible_for_retry()
                                 or (
-                                    isinstance(exc, OperationFailure)
+                                    isinstance(exc_to_check, OperationFailure)
                                     and exc_code not in helpers_shared._RETRYABLE_ERROR_CODES
                                 )
                             )
                         ):
                             raise
                         self._retrying = True
-                        self._last_error = exc
+                        self._last_error = exc_to_check
                         self._attempt_number += 1
 
                         # Revert back to starting state only if the first
@@ -2902,6 +2906,7 @@ async def run(self) -> T:
                     overloaded = exc_to_check.has_error_label("SystemOverloadedError")
                     if overloaded:
                         self._max_retries = self._client.options.max_adaptive_retries
+                        self._base_backoff_ms = getattr(exc_to_check, "_base_backoff_ms", None)
                     always_retryable = exc_to_check.has_error_label("RetryableError") and overloaded
 
                     # Always retry abortTransaction and commitTransaction up to once
@@ -2945,7 +2950,10 @@ async def run(self) -> T:
 
                 self._always_retryable = always_retryable
                 if overloaded:
-                    delay = self._retry_policy.backoff(self._attempt_number)
+                    delay = self._retry_policy.backoff(
+                        self._attempt_number,
+                        self._base_backoff_ms / 1000 if self._base_backoff_ms else None,
+                    )
                     if not await self._retry_policy.should_retry(self._attempt_number, delay):
                         if exc_to_check.has_error_label("NoWritesPerformed") and self._last_error:
                             raise self._last_error from exc
diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py
@@ -241,7 +241,7 @@ async def _hello(
         cmd = self.hello_cmd()
         performing_handshake = not self.performed_handshake
         awaitable = False
-        cmd["backpressure"] = True
+        cmd["backpressure"] = "2"
         if performing_handshake:
             self.performed_handshake = True
             cmd["client"] = self.opts.metadata
diff --git a/pymongo/errors.py b/pymongo/errors.py
@@ -190,12 +190,15 @@ def __init__(
         max_wire_version: Optional[int] = None,
     ) -> None:
         error_labels = None
+        base_backoff_ms = None
         if details is not None:
             error_labels = details.get("errorLabels")
+            base_backoff_ms = details.get("baseBackoffMS")
         super().__init__(_format_detailed_error(error, details), error_labels=error_labels)
         self.__code = code
         self.__details = details
         self.__max_wire_version = max_wire_version
+        self.__base_backoff_ms = base_backoff_ms
 
     @property
     def _max_wire_version(self) -> Optional[int]:
@@ -222,6 +225,10 @@ def details(self) -> Optional[Mapping[str, Any]]:
     def timeout(self) -> bool:
         return self.__code in (50,)
 
+    @property
+    def _base_backoff_ms(self) -> Optional[float]:
+        return self.__base_backoff_ms
+
 
 class CursorNotFound(OperationFailure):
     """Raised while iterating query results if the cursor is
diff --git a/pymongo/synchronous/helpers.py b/pymongo/synchronous/helpers.py
@@ -26,6 +26,7 @@
 from typing import (
     Any,
     Callable,
+    Optional,
     TypeVar,
     cast,
 )
@@ -84,10 +85,12 @@ def inner(*args: Any, **kwargs: Any) -> Any:
 
 
 def _backoff(
-    attempt: int, initial_delay: float = _BACKOFF_INITIAL, max_delay: float = _BACKOFF_MAX
+    attempt: int,
+    base_backoff: float,
+    max_delay: float = _BACKOFF_MAX,
 ) -> float:
     jitter = random.random()  # noqa: S311
-    return jitter * min(initial_delay * (2**attempt), max_delay)
+    return jitter * min(base_backoff * (2**attempt), max_delay)
 
 
 class _RetryPolicy:
@@ -103,9 +106,13 @@ def __init__(
         self.backoff_initial = backoff_initial
         self.backoff_max = backoff_max
 
-    def backoff(self, attempt: int) -> float:
-        """Return the backoff duration for the given attempt."""
-        return _backoff(max(0, attempt - 1), self.backoff_initial, self.backoff_max)
+    def backoff(self, attempt: int, base_backoff: Optional[float] = None) -> float:
+        """Return the actual backoff duration for the given attempt and base backoff."""
+        return _backoff(
+            max(0, attempt),
+            self.backoff_initial if base_backoff is None or base_backoff < 0 else base_backoff,
+            self.backoff_max,
+        )
 
     def should_retry(self, attempt: int, delay: float) -> bool:
         """Return if we have retry attempts remaining and the next backoff would not exceed a timeout."""
diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py
@@ -2798,6 +2798,7 @@ def __init__(
         self._attempt_number = 0
         self._is_run_command = is_run_command
         self._is_aggregate_write = is_aggregate_write
+        self._base_backoff_ms: Optional[float] = None
 
     def run(self) -> T:
         """Runs the supplied func() and attempts a retry
@@ -2847,26 +2848,29 @@ def run(self) -> T:
 
                 # Execute specialized catch on read
                 if self._is_read:
-                    if isinstance(exc, (ConnectionFailure, OperationFailure)):
+                    if isinstance(exc_to_check, (ConnectionFailure, OperationFailure)):
                         # ConnectionFailures do not supply a code property
-                        exc_code = getattr(exc, "code", None)
-                        overloaded = exc.has_error_label("SystemOverloadedError")
+                        exc_code = getattr(exc_to_check, "code", None)
+                        overloaded = exc_to_check.has_error_label("SystemOverloadedError")
                         if overloaded:
                             self._max_retries = self._client.options.max_adaptive_retries
-                        always_retryable = exc.has_error_label("RetryableError") and overloaded
+                            self._base_backoff_ms = getattr(exc_to_check, "_base_backoff_ms", None)
+                        always_retryable = (
+                            exc_to_check.has_error_label("RetryableError") and overloaded
+                        )
                         if not self._client.options.retry_reads or (
                             not always_retryable
                             and (
                                 self._is_not_eligible_for_retry()
                                 or (
-                                    isinstance(exc, OperationFailure)
+                                    isinstance(exc_to_check, OperationFailure)
                                     and exc_code not in helpers_shared._RETRYABLE_ERROR_CODES
                                 )
                             )
                         ):
                             raise
                         self._retrying = True
-                        self._last_error = exc
+                        self._last_error = exc_to_check
                         self._attempt_number += 1
 
                         # Revert back to starting state only if the first
@@ -2893,6 +2897,7 @@ def run(self) -> T:
                     overloaded = exc_to_check.has_error_label("SystemOverloadedError")
                     if overloaded:
                         self._max_retries = self._client.options.max_adaptive_retries
+                        self._base_backoff_ms = getattr(exc_to_check, "_base_backoff_ms", None)
                     always_retryable = exc_to_check.has_error_label("RetryableError") and overloaded
 
                     # Always retry abortTransaction and commitTransaction up to once
@@ -2936,7 +2941,10 @@ def run(self) -> T:
 
                 self._always_retryable = always_retryable
                 if overloaded:
-                    delay = self._retry_policy.backoff(self._attempt_number)
+                    delay = self._retry_policy.backoff(
+                        self._attempt_number,
+                        self._base_backoff_ms / 1000 if self._base_backoff_ms else None,
+                    )
                     if not self._retry_policy.should_retry(self._attempt_number, delay):
                         if exc_to_check.has_error_label("NoWritesPerformed") and self._last_error:
                             raise self._last_error from exc
diff --git a/pymongo/synchronous/pool.py b/pymongo/synchronous/pool.py
@@ -241,7 +241,7 @@ def _hello(
         cmd = self.hello_cmd()
         performing_handshake = not self.performed_handshake
         awaitable = False
-        cmd["backpressure"] = True
+        cmd["backpressure"] = "2"
         if performing_handshake:
             self.performed_handshake = True
             cmd["client"] = self.opts.metadata
diff --git a/test/asynchronous/test_client_backpressure.py b/test/asynchronous/test_client_backpressure.py
@@ -223,9 +223,9 @@ async def test_01_operation_retry_uses_exponential_backoff(self, random_func):
             end1 = perf_counter()
 
             # f. Compare the times between the two runs.
-            # The sum of 2 backoffs is 0.3 seconds. There is a 0.3-second window to account for potential variance between the two
+            # The sum of 2 backoffs is 0.6 seconds. There is a 0.6-second window to account for potential variance between the two
             # runs.
-            self.assertTrue(abs((end1 - start1) - (end0 - start0 + 0.3)) < 0.3)
+            self.assertTrue(abs((end1 - start1) - (end0 - start0 + 0.6)) < 0.6)
 
     @async_client_context.require_failCommand_appName
     async def test_03_overload_retries_limited(self):
@@ -294,6 +294,70 @@ async def test_04_overload_retries_limited_configured(self):
         # 6. Assert that the total number of started commands is max_retries + 1.
         self.assertEqual(len(self.listener.started_events), max_retries + 1)
 
+    @unittest.skipIf(
+        sys.platform == "darwin",
+        "externalClientBaseBackoffMS is not supported on macOS",
+    )
+    @patch("random.random")
+    @async_client_context.require_version_min(9, 0, 0, -1)
+    @async_client_context.require_failCommand_appName
+    async def test_05_overload_errors_with_basebackoffms_override_backoff(self, random_func):
+        # Drivers should test that overload errors with `baseBackoffMS` override the default backoff duration.
+
+        # 1. Let `client` be a `MongoClient`.
+        client = self.client
+
+        # 2. Let `coll` be a collection.
+        coll = client.test.test
+
+        # 3. Configure the random number generator used for exponential backoff jitter to always return a number as
+        # close as possible to `1`.
+        random_func.return_value = 1
+
+        # 4. Configure the following failPoint:
+        fail_point = dict(
+            mode="alwaysOn",
+            data=dict(
+                failCommands=["insert"],
+                errorCode=462,
+                errorLabels=["SystemOverloadedError", "RetryableError"],
+                appName=self.app_name,
+            ),
+        )
+        async with self.fail_point(fail_point):
+            # 5. Insert the document `{ a: 1 }`. Expect that the command errors. Measure the duration of the command
+            # execution.
+            start0 = perf_counter()
+            with self.assertRaises(OperationFailure):
+                await coll.insert_one({"a": 1})
+            end0 = perf_counter()
+            exponential_backoff_time = end0 - start0
+
+            # 6. Run the following command to set up `baseBackoffMS` on overload errors.
+            try:
+                await client.admin.command("setParameter", 1, externalClientBaseBackoffMS=50)
+
+                # 7. Execute step 5 again.
+                start1 = perf_counter()
+                with self.assertRaises(OperationFailure) as ctx:
+                    await coll.insert_one({"a": 1})
+                end1 = perf_counter()
+                with_base_backoff_ms_time = end1 - start1
+
+                # 8. Assert the server attached `baseBackoffMS` to the error and the driver parsed it.
+                self.assertEqual(ctx.exception._base_backoff_ms, 50)
+            finally:
+                # 9. Run the following command to disable `baseBackoffMS` on overload errors.
+                await client.admin.command("setParameter", 1, externalClientBaseBackoffMS=0)
+
+        # 10. Assert absolute bounds on each run's duration.
+        # A run can never be faster than the sum of its backoffs.
+        # With jitter pinned to 1, the default backoffs are 0.2 + 0.4 = 0.6s
+        # and the baseBackoffMS=50 backoffs are 0.1 + 0.2 = 0.3s.
+        self.assertGreaterEqual(exponential_backoff_time, 0.6)
+        self.assertGreaterEqual(with_base_backoff_ms_time, 0.3)
+        self.assertLess(with_base_backoff_ms_time, 0.6)
+
 
 # Location of JSON test specifications.
 if _IS_SYNC:
diff --git a/test/asynchronous/test_client_metadata.py b/test/asynchronous/test_client_metadata.py
@@ -229,8 +229,8 @@ async def test_handshake_documents_include_backpressure(self):
         await client.admin.command("ping")
 
         # Assert that for every handshake document intercepted:
-        # the document has a field `backpressure` whose value is `true`.
-        self.assertEqual(self.handshake_req["backpressure"], True)
+        # the document has a field `backpressure` whose value is `"2"`.
+        self.assertEqual(self.handshake_req["backpressure"], "2")
 
 
 if __name__ == "__main__":
diff --git a/test/asynchronous/test_transactions.py b/test/asynchronous/test_transactions.py
@@ -710,7 +710,9 @@ async def callback(session):
             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
+        self.assertLess(
+            abs(end - start - (no_backoff_time + 2.3)), 0.5
+        )  # sum of 13 backoffs is ~2.3
 
 
 class TestOptionsInsideTransactionProse(AsyncTransactionsBase):
diff --git a/test/test_client_backpressure.py b/test/test_client_backpressure.py
diff --git a/test/test_client_metadata.py b/test/test_client_metadata.py
diff --git a/test/test_transactions.py b/test/test_transactions.py
