PYTHON-5504 Prototype exponential backoff in with_transaction (#2492)

PYTHON-5505 Prototype system overload retry loop for all operations (#2497)

All commands that fail with the "Retryable" error label will be retried up to 3 times.
When the error includes the "SystemOverloaded" error label we apply exponential
backoff with jitter before attempting a retry.

PYTHON-5506 Prototype adaptive token bucket retry (#2501)

Add adaptive token bucket based retry policy.
Successfully completed commands deposit 0.1 token.
Failed retry attempts consume 1 token.
A retry is only permitted if there is an available token.
Token bucket starts full with the maximum 1000 tokens.

PYTHON-5505 Use proper RetryableError and SystemOverloadedError labels

Author: ShaneHarvey

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
@@ -470,6 +472,8 @@ 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
 
 
 def _within_time_limit(start_time: float) -> bool:
@@ -703,7 +707,13 @@ 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
         while True:
+            if retry:  # Implement exponential backoff on retry.
+                jitter = random.random()  # noqa: S311
+                backoff = jitter * min(_BACKOFF_INITIAL * (2**retry), _BACKOFF_MAX)
+                await asyncio.sleep(backoff)
+            retry += 1
             await self.start_transaction(
                 read_concern, write_concern, read_preference, max_commit_time_ms
             )

pymongo/asynchronous/collection.py

@@ -58,6 +58,7 @@
     AsyncCursor,
     AsyncRawBatchCursor,
 )
+from pymongo.asynchronous.helpers import _retry_overload
 from pymongo.collation import validate_collation_or_none
 from pymongo.common import _ecoc_coll_name, _esc_coll_name
 from pymongo.errors import (
@@ -252,6 +253,7 @@ def __init__(
             unicode_decode_error_handler="replace", document_class=dict
         )
         self._timeout = database.client.options.timeout
+        self._retry_policy = database.client._retry_policy
 
         if create or kwargs:
             if _IS_SYNC:
@@ -2227,6 +2229,7 @@ async def create_indexes(
         return await self._create_indexes(indexes, session, **kwargs)
 
     @_csot.apply
+    @_retry_overload
     async def _create_indexes(
         self, indexes: Sequence[IndexModel], session: Optional[AsyncClientSession], **kwargs: Any
     ) -> list[str]:
@@ -2422,7 +2425,6 @@ async def drop_indexes(
             kwargs["comment"] = comment
         await self._drop_index("*", session=session, **kwargs)
 
-    @_csot.apply
     async def drop_index(
         self,
         index_or_name: _IndexKeyHint,
@@ -2472,6 +2474,7 @@ async def drop_index(
         await self._drop_index(index_or_name, session, comment, **kwargs)
 
     @_csot.apply
+    @_retry_overload
     async def _drop_index(
         self,
         index_or_name: _IndexKeyHint,
@@ -3072,6 +3075,7 @@ async def aggregate_raw_batches(
             )
 
     @_csot.apply
+    @_retry_overload
     async def rename(
         self,
         new_name: str,

pymongo/asynchronous/database.py

@@ -38,6 +38,7 @@
 from pymongo.asynchronous.change_stream import AsyncDatabaseChangeStream
 from pymongo.asynchronous.collection import AsyncCollection
 from pymongo.asynchronous.command_cursor import AsyncCommandCursor
+from pymongo.asynchronous.helpers import _retry_overload
 from pymongo.common import _ecoc_coll_name, _esc_coll_name
 from pymongo.database_shared import _check_name, _CodecDocumentType
 from pymongo.errors import CollectionInvalid, InvalidOperation
@@ -135,6 +136,7 @@ def __init__(
         self._name = name
         self._client: AsyncMongoClient[_DocumentType] = client
         self._timeout = client.options.timeout
+        self._retry_policy = client._retry_policy
 
     @property
     def client(self) -> AsyncMongoClient[_DocumentType]:
@@ -477,6 +479,7 @@ async def watch(
         return change_stream
 
     @_csot.apply
+    @_retry_overload
     async def create_collection(
         self,
         name: str,
@@ -819,6 +822,7 @@ async def command(
         ...
 
     @_csot.apply
+    @_retry_overload
     async def command(
         self,
         command: Union[str, MutableMapping[str, Any]],
@@ -950,6 +954,7 @@ async def command(
             )
 
     @_csot.apply
+    @_retry_overload
     async def cursor_command(
         self,
         command: Union[str, MutableMapping[str, Any]],
@@ -1265,6 +1270,7 @@ async def _drop_helper(
             )
 
     @_csot.apply
+    @_retry_overload
     async def drop_collection(
         self,
         name_or_collection: Union[str, AsyncCollection[_DocumentTypeArg]],

pymongo/asynchronous/helpers.py

@@ -16,18 +16,26 @@
 from __future__ import annotations
 
 import asyncio
+import builtins
+import functools
+import random
 import socket
+import time
+import time as time  # noqa: PLC0414 # needed in sync version
 from typing import (
     Any,
     Callable,
     TypeVar,
     cast,
 )
 
+from pymongo import _csot
 from pymongo.errors import (
     OperationFailure,
+    PyMongoError,
 )
 from pymongo.helpers_shared import _REAUTHENTICATION_REQUIRED_CODE
+from pymongo.lock import _async_create_lock
 
 _IS_SYNC = False
 
@@ -36,6 +44,7 @@
 
 
 def _handle_reauth(func: F) -> F:
+    @functools.wraps(func)
     async def inner(*args: Any, **kwargs: Any) -> Any:
         no_reauth = kwargs.pop("no_reauth", False)
         from pymongo.asynchronous.pool import AsyncConnection
@@ -68,6 +77,123 @@ async def inner(*args: Any, **kwargs: Any) -> Any:
     return cast(F, inner)
 
 
+_MAX_RETRIES = 3
+_BACKOFF_INITIAL = 0.05
+_BACKOFF_MAX = 10
+# DRIVERS-3240 will determine these defaults.
+DEFAULT_RETRY_TOKEN_CAPACITY = 1000.0
+DEFAULT_RETRY_TOKEN_RETURN = 0.1
+
+
+def _backoff(
+    attempt: int, initial_delay: float = _BACKOFF_INITIAL, max_delay: float = _BACKOFF_MAX
+) -> float:
+    jitter = random.random()  # noqa: S311
+    return jitter * min(initial_delay * (2**attempt), max_delay)
+
+
+class _TokenBucket:
+    """A token bucket implementation for rate limiting."""
+
+    def __init__(
+        self,
+        capacity: float = DEFAULT_RETRY_TOKEN_CAPACITY,
+        return_rate: float = DEFAULT_RETRY_TOKEN_RETURN,
+    ):
+        self.lock = _async_create_lock()
+        self.capacity = capacity
+        # DRIVERS-3240 will determine how full the bucket should start.
+        self.tokens = capacity
+        self.return_rate = return_rate
+
+    async def consume(self) -> bool:
+        """Consume a token from the bucket if available."""
+        async with self.lock:
+            if self.tokens >= 1:
+                self.tokens -= 1
+                return True
+            return False
+
+    async def deposit(self, retry: bool = False) -> None:
+        """Deposit a token back into the bucket."""
+        retry_token = 1 if retry else 0
+        async with self.lock:
+            self.tokens = min(self.capacity, self.tokens + retry_token + self.return_rate)
+
+
+class _RetryPolicy:
+    """A retry limiter that performs exponential backoff with jitter.
+
+    Retry attempts are limited by a token bucket to prevent overwhelming the server during
+    a prolonged outage or high load.
+    """
+
+    def __init__(
+        self,
+        token_bucket: _TokenBucket,
+        attempts: int = _MAX_RETRIES,
+        backoff_initial: float = _BACKOFF_INITIAL,
+        backoff_max: float = _BACKOFF_MAX,
+    ):
+        self.token_bucket = token_bucket
+        self.attempts = attempts
+        self.backoff_initial = backoff_initial
+        self.backoff_max = backoff_max
+
+    async def record_success(self, retry: bool) -> None:
+        """Record a successful operation."""
+        await self.token_bucket.deposit(retry)
+
+    def backoff(self, attempt: int) -> float:
+        """Return the backoff duration for the given ."""
+        return _backoff(max(0, attempt - 1), self.backoff_initial, self.backoff_max)
+
+    async def should_retry(self, attempt: int, delay: float) -> bool:
+        """Return if we have budget to retry and how long to backoff."""
+        if attempt > self.attempts:
+            return False
+
+        # If the delay would exceed the deadline, bail early before consuming a token.
+        if _csot.get_timeout():
+            if time.monotonic() + delay > _csot.get_deadline():
+                return False
+
+        # Check token bucket last since we only want to consume a token if we actually retry.
+        if not await self.token_bucket.consume():
+            # DRIVERS-3246 Improve diagnostics when this case happens.
+            # We could add info to the exception and log.
+            return False
+        return True
+
+
+def _retry_overload(func: F) -> F:
+    @functools.wraps(func)
+    async def inner(self: Any, *args: Any, **kwargs: Any) -> Any:
+        retry_policy = self._retry_policy
+        attempt = 0
+        while True:
+            try:
+                res = await func(self, *args, **kwargs)
+                await retry_policy.record_success(retry=attempt > 0)
+                return res
+            except PyMongoError as exc:
+                if not exc.has_error_label("RetryableError"):
+                    raise
+                attempt += 1
+                delay = 0
+                if exc.has_error_label("SystemOverloadedError"):
+                    delay = retry_policy.backoff(attempt)
+                if not await retry_policy.should_retry(attempt, delay):
+                    raise
+
+                # Implement exponential backoff on retry.
+                if delay:
+                    await asyncio.sleep(delay)
+                continue
+
+    return cast(F, inner)
+
+
 async def _getaddrinfo(
     host: Any, port: Any, **kwargs: Any
 ) -> list[