PYTHON-5767 - Finalize client backpressure implementation for phase 1 rollout

Author: NoahStapp

pymongo/asynchronous/helpers.py

@@ -30,6 +30,7 @@
 )
 
 from pymongo import _csot
+from pymongo.common import MAX_ADAPTIVE_RETRIES
 from pymongo.errors import (
     OperationFailure,
 )
@@ -76,7 +77,6 @@ async def inner(*args: Any, **kwargs: Any) -> Any:
     return cast(F, inner)
 
 
-_MAX_RETRIES = 5
 _BACKOFF_INITIAL = 0.1
 _BACKOFF_MAX = 10
 DEFAULT_RETRY_TOKEN_CAPACITY = 1000.0
@@ -128,16 +128,15 @@ class _RetryPolicy:
     def __init__(
         self,
         token_bucket: _TokenBucket,
-        attempts: int = _MAX_RETRIES,
+        attempts: int = MAX_ADAPTIVE_RETRIES,
         backoff_initial: float = _BACKOFF_INITIAL,
         backoff_max: float = _BACKOFF_MAX,
-        adaptive_retry: bool = False,
     ):
         self.token_bucket = token_bucket
         self.attempts = attempts
         self.backoff_initial = backoff_initial
         self.backoff_max = backoff_max
-        self.adaptive_retry = adaptive_retry
+        self.adaptive_retry = False
 
     async def record_success(self, retry: bool) -> None:
         """Record a successful operation."""

pymongo/asynchronous/mongo_client.py

@@ -615,17 +615,17 @@ def __init__(
             client to use Stable API. See `versioned API <https://www.mongodb.com/docs/manual/reference/stable-api/#what-is-the-stable-api--and-should-you-use-it->`_ for
             details.
 
-          | **Adaptive retry options:**
-          | (If not enabled explicitly, adaptive retries will not be enabled.)
+          | **Overload retry options:**|
 
-          - `adaptive_retries`: (boolean) Whether the adaptive retry mechanism is enabled for this client.
-            If enabled, server overload errors will use a token-bucket based system to mitigate further overload.
+          - `max_adaptive_retries`: (int) How many retries to allow for overload errors. Defaults to ``2``.
+          - `enable_overload_retargeting`: (boolean) Whether overload retargeting is enabled for this client.
+            If enabled, server overload errors will cause retry attempts to select a server that has not yet returned an overload error, if possible.
             Defaults to ``False``.
 
         .. seealso:: The MongoDB documentation on `connections <https://dochub.mongodb.org/core/connections>`_.
 
         .. versionchanged:: 4.17
-           Added the ``adaptive_retries`` URI and keyword argument.
+           Added the ``max_adaptive_retries`` and ``enable_overload_retargeting`` URI and keyword arguments.
 
         .. versionchanged:: 4.5
            Added the ``serverMonitoringMode`` keyword argument.
@@ -894,9 +894,7 @@ def __init__(
             self._options.read_concern,
         )
 
-        self._retry_policy = _RetryPolicy(
-            _TokenBucket(), adaptive_retry=self._options.adaptive_retries
-        )
+        self._retry_policy = _RetryPolicy(_TokenBucket())
 
         self._init_based_on_options(self._seeds, srv_max_hosts, srv_service_name)
 

pymongo/client_options.py

@@ -240,6 +240,16 @@ def __init__(
             if "adaptive_retries" in options
             else options.get("adaptiveretries", common.ADAPTIVE_RETRIES)
         )
+        self.__max_adaptive_retries = (
+            options.get("max_adaptive_retries", common.MAX_ADAPTIVE_RETRIES)
+            if "max_adaptive_retries" in options
+            else options.get("maxadaptiveretries", common.MAX_ADAPTIVE_RETRIES)
+        )
+        self.__enable_overload_retargeting = (
+            options.get("enable_overload_retargeting", common.ENABLE_OVERLOAD_RETARGETING)
+            if "enable_overload_retargeting" in options
+            else options.get("enableoverloadretargeting", common.ENABLE_OVERLOAD_RETARGETING)
+        )
 
     @property
     def _options(self) -> Mapping[str, Any]:
@@ -353,9 +363,17 @@ def server_monitoring_mode(self) -> str:
         return self.__server_monitoring_mode
 
     @property
-    def adaptive_retries(self) -> bool:
-        """The configured adaptiveRetries option.
+    def max_adaptive_retries(self) -> int:
+        """The configured maxAdaptiveRetries option.
+
+        .. versionadded:: 4.17
+        """
+        return self.__max_adaptive_retries
+
+    @property
+    def enable_overload_retargeting(self) -> bool:
+        """The configured enableOverloadRetargeting option.
 
-        .. versionadded:: 4.XX
+        .. versionadded:: 4.17
         """
-        return self.__adaptive_retries
+        return self.__enable_overload_retargeting

pymongo/common.py

@@ -143,6 +143,12 @@
 # Default value for adaptiveRetries
 ADAPTIVE_RETRIES = False
 
+# Default value for max adaptive retries
+MAX_ADAPTIVE_RETRIES = 2
+
+# Default value for enableOverloadRetargeting
+ENABLE_OVERLOAD_RETARGETING = False
+
 # Auth mechanism properties that must raise an error instead of warning if they invalidate.
 _MECH_PROP_MUST_RAISE = ["CANONICALIZE_HOST_NAME"]
 
@@ -742,6 +748,8 @@ def validate_server_monitoring_mode(option: str, value: str) -> str:
     "timeoutms": validate_timeoutms,
     "servermonitoringmode": validate_server_monitoring_mode,
     "adaptiveretries": validate_boolean_or_string,
+    "maxadaptiveretries": validate_non_negative_integer,
+    "enableoverloadretargeting": validate_boolean_or_string,
 }
 
 # Dictionary where keys are the names of URI options specific to pymongo,
@@ -776,6 +784,8 @@ def validate_server_monitoring_mode(option: str, value: str) -> str:
     "auto_encryption_opts": validate_auto_encryption_opts_or_none,
     "authoidcallowedhosts": validate_list,
     "adaptive_retries": validate_boolean_or_string,
+    "max_adaptive_retries": validate_non_negative_integer,
+    "enable_overload_retargeting": validate_boolean_or_string,
 }
 
 # Dictionary where keys are any URI option name, and values are the

pymongo/synchronous/helpers.py

@@ -30,6 +30,7 @@
 )
 
 from pymongo import _csot
+from pymongo.common import MAX_ADAPTIVE_RETRIES
 from pymongo.errors import (
     OperationFailure,
 )
@@ -76,7 +77,6 @@ def inner(*args: Any, **kwargs: Any) -> Any:
     return cast(F, inner)
 
 
-_MAX_RETRIES = 5
 _BACKOFF_INITIAL = 0.1
 _BACKOFF_MAX = 10
 DEFAULT_RETRY_TOKEN_CAPACITY = 1000.0
@@ -128,16 +128,15 @@ class _RetryPolicy:
     def __init__(
         self,
         token_bucket: _TokenBucket,
-        attempts: int = _MAX_RETRIES,
+        attempts: int = MAX_ADAPTIVE_RETRIES,
         backoff_initial: float = _BACKOFF_INITIAL,
         backoff_max: float = _BACKOFF_MAX,
-        adaptive_retry: bool = False,
     ):
         self.token_bucket = token_bucket
         self.attempts = attempts
         self.backoff_initial = backoff_initial
         self.backoff_max = backoff_max
-        self.adaptive_retry = adaptive_retry
+        self.adaptive_retry = False
 
     def record_success(self, retry: bool) -> None:
         """Record a successful operation."""

pymongo/synchronous/mongo_client.py

@@ -615,17 +615,17 @@ def __init__(
             client to use Stable API. See `versioned API <https://www.mongodb.com/docs/manual/reference/stable-api/#what-is-the-stable-api--and-should-you-use-it->`_ for
             details.
 
-          | **Adaptive retry options:**
-          | (If not enabled explicitly, adaptive retries will not be enabled.)
+          | **Overload retry options:**|
 
-          - `adaptive_retries`: (boolean) Whether the adaptive retry mechanism is enabled for this client.
-            If enabled, server overload errors will use a token-bucket based system to mitigate further overload.
+          - `max_adaptive_retries`: (int) How many retries to allow for overload errors. Defaults to ``2``.
+          - `enable_overload_retargeting`: (boolean) Whether overload retargeting is enabled for this client.
+            If enabled, server overload errors will cause retry attempts to select a server that has not yet returned an overload error, if possible.
             Defaults to ``False``.
 
         .. seealso:: The MongoDB documentation on `connections <https://dochub.mongodb.org/core/connections>`_.
 
         .. versionchanged:: 4.17
-           Added the ``adaptive_retries`` URI and keyword argument.
+           Added the ``max_adaptive_retries`` and ``enable_overload_retargeting`` URI and keyword arguments.
 
         .. versionchanged:: 4.5
            Added the ``serverMonitoringMode`` keyword argument.
@@ -894,9 +894,7 @@ def __init__(
             self._options.read_concern,
         )
 
-        self._retry_policy = _RetryPolicy(
-            _TokenBucket(), adaptive_retry=self._options.adaptive_retries
-        )
+        self._retry_policy = _RetryPolicy(_TokenBucket())
 
         self._init_based_on_options(self._seeds, srv_max_hosts, srv_service_name)
 

test/asynchronous/test_client.py

@@ -652,20 +652,37 @@ async def test_detected_environment_warning(self, mock_get_hosts):
             with self.assertWarns(UserWarning):
                 self.simple_client(multi_host)
 
-    async def test_adaptive_retries(self):
-        # Assert that adaptive retries are disabled by default.
+    async def test_max_adaptive_retries(self):
+        # Assert that max adaptive retries default to 2.
         c = self.simple_client(connect=False)
-        self.assertFalse(c.options.adaptive_retries)
+        self.assertEqual(c.options.max_adaptive_retries, 2)
 
         # Assert that adaptive retries can be enabled through connection or client options.
-        c = self.simple_client(connect=False, adaptive_retries=True)
-        self.assertTrue(c.options.adaptive_retries)
+        c = self.simple_client(connect=False, max_adaptive_retries=10)
+        self.assertEqual(c.options.max_adaptive_retries, 10)
 
-        c = self.simple_client(connect=False, adaptiveRetries=True)
-        self.assertTrue(c.options.adaptive_retries)
+        c = self.simple_client(connect=False, maxAdaptiveRetries=10)
+        self.assertEqual(c.options.max_adaptive_retries, 10)
 
-        c = self.simple_client(host="mongodb://localhost/?adaptiveretries=true", connect=False)
-        self.assertTrue(c.options.adaptive_retries)
+        c = self.simple_client(host="mongodb://localhost/?maxAdaptiveRetries=10", connect=False)
+        self.assertEqual(c.options.max_adaptive_retries, 10)
+
+    async def test_enable_overload_retargeting(self):
+        # Assert that overload retargeting defaults to false.
+        c = self.simple_client(connect=False)
+        self.assertFalse(c.options.enable_overload_retargeting)
+
+        # Assert that overload retargeting can be enabled through connection or client options.
+        c = self.simple_client(connect=False, enable_overload_retargeting=True)
+        self.assertTrue(c.options.enable_overload_retargeting)
+
+        c = self.simple_client(connect=False, enableOverloadRetargeting=True)
+        self.assertTrue(c.options.enable_overload_retargeting)
+
+        c = self.simple_client(
+            host="mongodb://localhost/?enableOverloadRetargeting=true", connect=False
+        )
+        self.assertTrue(c.options.enable_overload_retargeting)
 
 
 class TestClient(AsyncIntegrationTest):