PYTHON-5506 Check CSOT deadline before consuming a token

ShaneHarvey · ShaneHarvey · commit 752d8c0f9f11 · 2025-08-21T17:05:55.000-07:00
diff --git a/pymongo/asynchronous/helpers.py b/pymongo/asynchronous/helpers.py
@@ -29,6 +29,7 @@
     cast,
 )
 
+from pymongo import _csot
 from pymongo.errors import (
     OperationFailure,
     PyMongoError,
@@ -85,12 +86,11 @@ async def inner(*args: Any, **kwargs: Any) -> Any:
 _TIME = time  # Added so synchro script doesn't remove the time import.
 
 
-async def _backoff(
+def _backoff(
     attempt: int, initial_delay: float = _BACKOFF_INITIAL, max_delay: float = _BACKOFF_MAX
-) -> None:
+) -> float:
     jitter = random.random()  # noqa: S311
-    backoff = jitter * min(initial_delay * (2**attempt), max_delay)
-    await asyncio.sleep(backoff)
+    return jitter * min(initial_delay * (2**attempt), max_delay)
 
 
 class _TokenBucket:
@@ -145,15 +145,20 @@ async def record_success(self, retry: bool) -> None:
         """Record a successful operation."""
         await self.token_bucket.deposit(retry)
 
-    async def backoff(self, attempt: int) -> None:
+    def backoff(self, attempt: int) -> float:
         """Return the backoff duration for the given ."""
-        await _backoff(max(0, attempt - 1), self.backoff_initial, self.backoff_max)
+        return _backoff(max(0, attempt - 1), self.backoff_initial, self.backoff_max)
 
-    async def should_retry(self, attempt: int) -> bool:
+    async def should_retry(self, attempt: int, delay: float) -> bool:
         """Return if we have budget to retry and how long to backoff."""
-        # TODO: Check CSOT deadline here.
         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.
@@ -176,12 +181,15 @@ async def inner(self: Any, *args: Any, **kwargs: Any) -> Any:
                 if not exc.has_error_label("Retryable"):
                     raise
                 attempt += 1
-                if not await retry_policy.should_retry(attempt):
+                delay = 0
+                if exc.has_error_label("SystemOverloaded"):
+                    delay = retry_policy.backoff(attempt)
+                if not await retry_policy.should_retry(attempt, delay):
                     raise
 
                 # Implement exponential backoff on retry.
-                if exc.has_error_label("SystemOverloaded"):
-                    await retry_policy.backoff(attempt)
+                if delay:
+                    await asyncio.sleep(delay)
                 continue
 
     return cast(F, inner)
diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py
@@ -35,6 +35,7 @@
 import asyncio
 import contextlib
 import os
+import time
 import warnings
 import weakref
 from collections import defaultdict
@@ -174,6 +175,8 @@
     UpdateMany,
 ]
 
+_TIME = time  # Added so synchro script doesn't remove the time import.
+
 
 class AsyncMongoClient(common.BaseObject, Generic[_DocumentType]):
     HOST = "localhost"
@@ -2853,13 +2856,14 @@ async def run(self) -> T:
 
                 self._always_retryable = always_retryable
                 if always_retryable:
-                    if not await self._retry_policy.should_retry(self._attempt_number):
+                    delay = self._retry_policy.backoff(self._attempt_number) if overloaded else 0
+                    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
                         else:
                             raise
                     if overloaded:
-                        await self._retry_policy.backoff(self._attempt_number)
+                        await asyncio.sleep(delay)
 
     def _is_not_eligible_for_retry(self) -> bool:
         """Checks if the exchange is not eligible for retry"""
diff --git a/pymongo/synchronous/helpers.py b/pymongo/synchronous/helpers.py
@@ -29,6 +29,7 @@
     cast,
 )
 
+from pymongo import _csot
 from pymongo.errors import (
     OperationFailure,
     PyMongoError,
@@ -87,10 +88,9 @@ def inner(*args: Any, **kwargs: Any) -> Any:
 
 def _backoff(
     attempt: int, initial_delay: float = _BACKOFF_INITIAL, max_delay: float = _BACKOFF_MAX
-) -> None:
+) -> float:
     jitter = random.random()  # noqa: S311
-    backoff = jitter * min(initial_delay * (2**attempt), max_delay)
-    time.sleep(backoff)
+    return jitter * min(initial_delay * (2**attempt), max_delay)
 
 
 class _TokenBucket:
@@ -145,15 +145,20 @@ def record_success(self, retry: bool) -> None:
         """Record a successful operation."""
         self.token_bucket.deposit(retry)
 
-    def backoff(self, attempt: int) -> None:
+    def backoff(self, attempt: int) -> float:
         """Return the backoff duration for the given ."""
-        _backoff(max(0, attempt - 1), self.backoff_initial, self.backoff_max)
+        return _backoff(max(0, attempt - 1), self.backoff_initial, self.backoff_max)
 
-    def should_retry(self, attempt: int) -> bool:
+    def should_retry(self, attempt: int, delay: float) -> bool:
         """Return if we have budget to retry and how long to backoff."""
-        # TODO: Check CSOT deadline here.
         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 self.token_bucket.consume():
             # DRIVERS-3246 Improve diagnostics when this case happens.
@@ -176,12 +181,15 @@ def inner(self: Any, *args: Any, **kwargs: Any) -> Any:
                 if not exc.has_error_label("Retryable"):
                     raise
                 attempt += 1
-                if not retry_policy.should_retry(attempt):
+                delay = 0
+                if exc.has_error_label("SystemOverloaded"):
+                    delay = retry_policy.backoff(attempt)
+                if not retry_policy.should_retry(attempt, delay):
                     raise
 
                 # Implement exponential backoff on retry.
-                if exc.has_error_label("SystemOverloaded"):
-                    retry_policy.backoff(attempt)
+                if delay:
+                    time.sleep(delay)
                 continue
 
     return cast(F, inner)
diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py
@@ -35,6 +35,7 @@
 import asyncio
 import contextlib
 import os
+import time
 import warnings
 import weakref
 from collections import defaultdict
@@ -171,6 +172,8 @@
     UpdateMany,
 ]
 
+_TIME = time  # Added so synchro script doesn't remove the time import.
+
 
 class MongoClient(common.BaseObject, Generic[_DocumentType]):
     HOST = "localhost"
@@ -2843,13 +2846,14 @@ def run(self) -> T:
 
                 self._always_retryable = always_retryable
                 if always_retryable:
-                    if not self._retry_policy.should_retry(self._attempt_number):
+                    delay = self._retry_policy.backoff(self._attempt_number) if overloaded else 0
+                    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
                         else:
                             raise
                     if overloaded:
-                        self._retry_policy.backoff(self._attempt_number)
+                        time.sleep(delay)
 
     def _is_not_eligible_for_retry(self) -> bool:
         """Checks if the exchange is not eligible for retry"""
diff --git a/test/asynchronous/test_backpressure.py b/test/asynchronous/test_backpressure.py
@@ -15,8 +15,11 @@
 """Test Client Backpressure spec."""
 from __future__ import annotations
 
+import asyncio
 import sys
 
+import pymongo
+
 sys.path[0:0] = [""]
 
 from test.asynchronous import (
@@ -187,31 +190,41 @@ async def test_retry_policy(self):
         self.assertEqual(retry_policy.backoff_initial, helpers._BACKOFF_INITIAL)
         self.assertEqual(retry_policy.backoff_max, helpers._BACKOFF_MAX)
         for i in range(1, helpers._MAX_RETRIES + 1):
-            self.assertTrue(await retry_policy.should_retry(i))
-        self.assertFalse(await retry_policy.should_retry(helpers._MAX_RETRIES + 1))
+            self.assertTrue(await retry_policy.should_retry(i, 0))
+        self.assertFalse(await retry_policy.should_retry(helpers._MAX_RETRIES + 1, 0))
         for i in range(capacity - helpers._MAX_RETRIES):
-            self.assertTrue(await retry_policy.should_retry(1))
+            self.assertTrue(await retry_policy.should_retry(1, 0))
         # No tokens left, should not retry.
-        self.assertFalse(await retry_policy.should_retry(1))
+        self.assertFalse(await retry_policy.should_retry(1, 0))
         self.assertEqual(retry_policy.token_bucket.tokens, 0)
 
         # record_success should generate tokens.
         for _ in range(int(2 / helpers.DEFAULT_RETRY_TOKEN_RETURN)):
             await retry_policy.record_success(retry=False)
         self.assertAlmostEqual(retry_policy.token_bucket.tokens, 2)
         for i in range(2):
-            self.assertTrue(await retry_policy.should_retry(1))
-        self.assertFalse(await retry_policy.should_retry(1))
+            self.assertTrue(await retry_policy.should_retry(1, 0))
+        self.assertFalse(await retry_policy.should_retry(1, 0))
 
         # Recording a successful retry should return 1 additional token.
         await retry_policy.record_success(retry=True)
         self.assertAlmostEqual(
             retry_policy.token_bucket.tokens, 1 + helpers.DEFAULT_RETRY_TOKEN_RETURN
         )
-        self.assertTrue(await retry_policy.should_retry(1))
-        self.assertFalse(await retry_policy.should_retry(1))
+        self.assertTrue(await retry_policy.should_retry(1, 0))
+        self.assertFalse(await retry_policy.should_retry(1, 0))
         self.assertAlmostEqual(retry_policy.token_bucket.tokens, helpers.DEFAULT_RETRY_TOKEN_RETURN)
 
+    async def test_retry_policy_csot(self):
+        retry_policy = _RetryPolicy(_TokenBucket())
+        self.assertTrue(await retry_policy.should_retry(1, 0.5))
+        with pymongo.timeout(0.5):
+            self.assertTrue(await retry_policy.should_retry(1, 0))
+            self.assertTrue(await retry_policy.should_retry(1, 0.1))
+            # Would exceed the timeout, should not retry.
+            self.assertFalse(await retry_policy.should_retry(1, 1.0))
+        self.assertTrue(await retry_policy.should_retry(1, 1.0))
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/test_backpressure.py b/test/test_backpressure.py
@@ -15,8 +15,11 @@
 """Test Client Backpressure spec."""
 from __future__ import annotations
 
+import asyncio
 import sys
 
+import pymongo
+
 sys.path[0:0] = [""]
 
 from test import (
@@ -187,31 +190,41 @@ def test_retry_policy(self):
         self.assertEqual(retry_policy.backoff_initial, helpers._BACKOFF_INITIAL)
         self.assertEqual(retry_policy.backoff_max, helpers._BACKOFF_MAX)
         for i in range(1, helpers._MAX_RETRIES + 1):
-            self.assertTrue(retry_policy.should_retry(i))
-        self.assertFalse(retry_policy.should_retry(helpers._MAX_RETRIES + 1))
+            self.assertTrue(retry_policy.should_retry(i, 0))
+        self.assertFalse(retry_policy.should_retry(helpers._MAX_RETRIES + 1, 0))
         for i in range(capacity - helpers._MAX_RETRIES):
-            self.assertTrue(retry_policy.should_retry(1))
+            self.assertTrue(retry_policy.should_retry(1, 0))
         # No tokens left, should not retry.
-        self.assertFalse(retry_policy.should_retry(1))
+        self.assertFalse(retry_policy.should_retry(1, 0))
         self.assertEqual(retry_policy.token_bucket.tokens, 0)
 
         # record_success should generate tokens.
         for _ in range(int(2 / helpers.DEFAULT_RETRY_TOKEN_RETURN)):
             retry_policy.record_success(retry=False)
         self.assertAlmostEqual(retry_policy.token_bucket.tokens, 2)
         for i in range(2):
-            self.assertTrue(retry_policy.should_retry(1))
-        self.assertFalse(retry_policy.should_retry(1))
+            self.assertTrue(retry_policy.should_retry(1, 0))
+        self.assertFalse(retry_policy.should_retry(1, 0))
 
         # Recording a successful retry should return 1 additional token.
         retry_policy.record_success(retry=True)
         self.assertAlmostEqual(
             retry_policy.token_bucket.tokens, 1 + helpers.DEFAULT_RETRY_TOKEN_RETURN
         )
-        self.assertTrue(retry_policy.should_retry(1))
-        self.assertFalse(retry_policy.should_retry(1))
+        self.assertTrue(retry_policy.should_retry(1, 0))
+        self.assertFalse(retry_policy.should_retry(1, 0))
         self.assertAlmostEqual(retry_policy.token_bucket.tokens, helpers.DEFAULT_RETRY_TOKEN_RETURN)
 
+    def test_retry_policy_csot(self):
+        retry_policy = _RetryPolicy(_TokenBucket())
+        self.assertTrue(retry_policy.should_retry(1, 0.5))
+        with pymongo.timeout(0.5):
+            self.assertTrue(retry_policy.should_retry(1, 0))
+            self.assertTrue(retry_policy.should_retry(1, 0.1))
+            # Would exceed the timeout, should not retry.
+            self.assertFalse(retry_policy.should_retry(1, 1.0))
+        self.assertTrue(retry_policy.should_retry(1, 1.0))
+
 
 if __name__ == "__main__":
     unittest.main()
