PYTHON-5668 - Merge backpressure branch into mainline

.evergreen/resync-specs.sh

@@ -94,6 +94,9 @@ do
     change-streams|change_streams)
       cpjson change-streams/tests/ change_streams/
       ;;
+    client-backpressure|client_backpressure)
+      cpjson client-backpressure/tests client-backpressure
+      ;;
     client-side-encryption|csfle|fle)
       cpjson client-side-encryption/tests/ client-side-encryption/spec
       cpjson client-side-encryption/corpus/ client-side-encryption/corpus

doc/changelog.rst

@@ -14,6 +14,9 @@ PyMongo 4.17 brings a number of changes including:
 - Added the :meth:`~pymongo.asynchronous.client_session.AsyncClientSession.bind` and :meth:`~pymongo.client_session.ClientSession.bind` methods
   that allow users to bind a session to all database operations within the scope of a context manager instead of having to explicitly pass the session to each individual operation.
   See <PLACEHOLDER> for examples and more information.
+- Added support for MongoDB's Intelligent Workload Management (IWM) and ingress connection rate limiting features.
+  The driver now gracefully handles write-blocking scenarios and optimizes connection establishment during high-load conditions to maintain application availability.
+  See <DOCSP-55426> and <DOCSP-57078> for more information.
 
 Changes in Version 4.16.0 (2026/01/07)
 --------------------------------------

pymongo/asynchronous/client_bulk.py

@@ -563,9 +563,21 @@ async def _execute_command(
                         error, ConnectionFailure
                     ) and not isinstance(error, (NotPrimaryError, WaitQueueTimeoutError))
 
+                    retryable_label_error = (
+                        hasattr(error, "details")
+                        and isinstance(error.details, dict)
+                        and "errorLabels" in error.details
+                        and isinstance(error.details["errorLabels"], list)
+                        and "RetryableError" in error.details["errorLabels"]
+                    )
+
                     # Synthesize the full bulk result without modifying the
                     # current one because this write operation may be retried.
-                    if retryable and (retryable_top_level_error or retryable_network_error):
+                    if retryable and (
+                        retryable_top_level_error
+                        or retryable_network_error
+                        or retryable_label_error
+                    ):
                         full = copy.deepcopy(full_result)
                         _merge_command(self.ops, self.idx_offset, full, result)
                         _throw_client_bulk_write_exception(full, self.verbose_results)

pymongo/asynchronous/client_session.py

@@ -135,7 +135,9 @@
 
 from __future__ import annotations
 
+import asyncio
 import collections
+import random
 import time
 import uuid
 from collections.abc import Mapping as _Mapping
@@ -162,7 +164,9 @@
 from pymongo.errors import (
     ConfigurationError,
     ConnectionFailure,
+    ExecutionTimeout,
     InvalidOperation,
+    NetworkTimeout,
     OperationFailure,
     PyMongoError,
     WTimeoutError,
@@ -427,13 +431,17 @@ def __init__(self, opts: Optional[TransactionOptions], client: AsyncMongoClient[
         self.recovery_token = None
         self.attempt = 0
         self.client = client
+        self.has_completed_command = False
 
     def active(self) -> bool:
         return self.state in (_TxnState.STARTING, _TxnState.IN_PROGRESS)
 
     def starting(self) -> bool:
         return self.state == _TxnState.STARTING
 
+    def set_starting(self) -> None:
+        self.state = _TxnState.STARTING
+
     @property
     def pinned_conn(self) -> Optional[AsyncConnection]:
         if self.active() and self.conn_mgr:
@@ -459,6 +467,7 @@ async def reset(self) -> None:
         self.sharded = False
         self.recovery_token = None
         self.attempt = 0
+        self.has_completed_command = False
 
     def __del__(self) -> None:
         if self.conn_mgr:
@@ -493,11 +502,29 @@ 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 = 0.500  # 500ms max backoff
+_BACKOFF_INITIAL = 0.005  # 5ms initial backoff
 
 
-def _within_time_limit(start_time: float) -> bool:
+def _within_time_limit(start_time: float, backoff: float = 0) -> bool:
     """Are we within the with_transaction retry limit?"""
-    return time.monotonic() - start_time < _WITH_TRANSACTION_RETRY_TIME_LIMIT
+    remaining = _csot.remaining()
+    if remaining is not None and remaining <= 0:
+        return False
+    return time.monotonic() + backoff - start_time < _WITH_TRANSACTION_RETRY_TIME_LIMIT
+
+
+def _make_timeout_error(error: BaseException) -> PyMongoError:
+    """Convert error to a NetworkTimeout or ExecutionTimeout as appropriate."""
+    if _csot.remaining() is not None:
+        timeout_error: PyMongoError = ExecutionTimeout(
+            str(error), 50, {"ok": 0, "errmsg": str(error), "code": 50}
+        )
+    else:
+        timeout_error = NetworkTimeout(str(error))
+    if isinstance(error, PyMongoError):
+        timeout_error._error_labels = error._error_labels.copy()
+    return timeout_error
 
 
 _T = TypeVar("_T")
@@ -744,23 +771,34 @@ async def callback(session, custom_arg, custom_kwarg=None):
             https://github.com/mongodb/specifications/blob/master/source/transactions-convenient-api/transactions-convenient-api.md#handling-errors-inside-the-callback
         """
         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 * (1.5**retry), _BACKOFF_MAX)
+                if not _within_time_limit(start_time, backoff):
+                    assert last_error is not None
+                    raise _make_timeout_error(last_error) from last_error
+                await asyncio.sleep(backoff)
+            retry += 1
             await self.start_transaction(
                 read_concern, write_concern, read_preference, max_commit_time_ms
             )
             try:
                 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 (
-                    isinstance(exc, PyMongoError)
-                    and exc.has_error_label("TransientTransactionError")
-                    and _within_time_limit(start_time)
+                if isinstance(exc, PyMongoError) and exc.has_error_label(
+                    "TransientTransactionError"
                 ):
-                    # Retry the entire transaction.
-                    continue
+                    if _within_time_limit(start_time):
+                        # Retry the entire transaction.
+                        continue
+                    raise _make_timeout_error(last_error) from exc
                 raise
 
             if not self.in_transaction:
@@ -771,17 +809,18 @@ async def callback(session, custom_arg, custom_kwarg=None):
                 try:
                     await self.commit_transaction()
                 except PyMongoError as exc:
-                    if (
-                        exc.has_error_label("UnknownTransactionCommitResult")
-                        and _within_time_limit(start_time)
-                        and not _max_time_expired_error(exc)
-                    ):
+                    last_error = exc
+                    if exc.has_error_label(
+                        "UnknownTransactionCommitResult"
+                    ) and not _max_time_expired_error(exc):
+                        if not _within_time_limit(start_time):
+                            raise _make_timeout_error(last_error) from exc
                         # Retry the commit.
                         continue
 
-                    if exc.has_error_label("TransientTransactionError") and _within_time_limit(
-                        start_time
-                    ):
+                    if exc.has_error_label("TransientTransactionError"):
+                        if not _within_time_limit(start_time):
+                            raise _make_timeout_error(last_error) from exc
                         # Retry the entire transaction.
                         break
                     raise