PYTHON-5505 Make cursor getMore retryable

Author: ShaneHarvey

pymongo/asynchronous/mongo_client.py

@@ -2737,6 +2737,7 @@ def __init__(
     ):
         self._last_error: Optional[Exception] = None
         self._retrying = False
+        self._overload = False
         self._multiple_retries = _csot.get_timeout() is not None
         self._client = mongo_client
 
@@ -2808,8 +2809,6 @@ async def run(self) -> T:
 
                 # Specialized catch on write operation
                 if not self._is_read:
-                    if not self._retryable:
-                        raise
                     if isinstance(exc, ClientBulkWriteException) and exc.error:
                         retryable_write_error_exc = isinstance(
                             exc.error, PyMongoError
@@ -2820,6 +2819,8 @@ async def run(self) -> T:
                     else:
                         retryable_write_error_exc = exc.has_error_label("RetryableWriteError")
                         overload = exc.has_error_label("Retryable")
+                    if not self._retryable and not overload:
+                        raise
                     if retryable_write_error_exc or overload:
                         assert self._session
                         await self._session._unpin()
@@ -2843,6 +2844,7 @@ async def run(self) -> T:
                 if self._client.topology_description.topology_type == TOPOLOGY_TYPE.Sharded:
                     self._deprioritized_servers.append(self._server)
 
+                self._overload = overload
                 if overload:
                     if self._attempt_number > _MAX_RETRIES:
                         if exc.has_error_label("NoWritesPerformed") and self._last_error:
@@ -2944,7 +2946,7 @@ async def _read(self) -> T:
             conn,
             read_pref,
         ):
-            if self._retrying and not self._retryable:
+            if self._retrying and not self._retryable and not self._overload:
                 self._check_last_error()
             if self._retrying:
                 _debug_log(

pymongo/synchronous/mongo_client.py

@@ -2727,6 +2727,7 @@ def __init__(
     ):
         self._last_error: Optional[Exception] = None
         self._retrying = False
+        self._overload = False
         self._multiple_retries = _csot.get_timeout() is not None
         self._client = mongo_client
 
@@ -2798,8 +2799,6 @@ def run(self) -> T:
 
                 # Specialized catch on write operation
                 if not self._is_read:
-                    if not self._retryable:
-                        raise
                     if isinstance(exc, ClientBulkWriteException) and exc.error:
                         retryable_write_error_exc = isinstance(
                             exc.error, PyMongoError
@@ -2810,6 +2809,8 @@ def run(self) -> T:
                     else:
                         retryable_write_error_exc = exc.has_error_label("RetryableWriteError")
                         overload = exc.has_error_label("Retryable")
+                    if not self._retryable and not overload:
+                        raise
                     if retryable_write_error_exc or overload:
                         assert self._session
                         self._session._unpin()
@@ -2833,6 +2834,7 @@ def run(self) -> T:
                 if self._client.topology_description.topology_type == TOPOLOGY_TYPE.Sharded:
                     self._deprioritized_servers.append(self._server)
 
+                self._overload = overload
                 if overload:
                     if self._attempt_number > _MAX_RETRIES:
                         if exc.has_error_label("NoWritesPerformed") and self._last_error:
@@ -2934,7 +2936,7 @@ def _read(self) -> T:
             conn,
             read_pref,
         ):
-            if self._retrying and not self._retryable:
+            if self._retrying and not self._retryable and not self._overload:
                 self._check_last_error()
             if self._retrying:
                 _debug_log(

test/asynchronous/test_backpressure.py

@@ -98,6 +98,37 @@ async def test_retry_overload_error_insert_one(self):
 
         self.assertIn("Retryable", str(error.exception))
 
+    @async_client_context.require_failCommand_appName
+    async def test_retry_overload_error_getMore(self):
+        coll = self.db.t
+        await coll.insert_many([{"x": 1} for _ in range(10)])
+
+        # Ensure command is retried on overload error.
+        fail_once = {
+            "configureFailPoint": "failCommand",
+            "mode": {"times": _MAX_RETRIES},
+            "data": {
+                "failCommands": ["getMore"],
+                "errorCode": 462,  # IngressRequestRateLimitExceeded
+                "errorLabels": ["Retryable"],
+            },
+        }
+        cursor = coll.find(batch_size=2)
+        await cursor.next()
+        async with self.fail_point(fail_once):
+            await cursor.to_list()
+
+        # Ensure command stops retrying after _MAX_RETRIES.
+        fail_many_times = fail_once.copy()
+        fail_many_times["mode"] = {"times": _MAX_RETRIES + 1}
+        cursor = coll.find(batch_size=2)
+        await cursor.next()
+        async with self.fail_point(fail_many_times):
+            with self.assertRaises(PyMongoError) as error:
+                await cursor.to_list()
+
+        self.assertIn("Retryable", str(error.exception))
+
 
 if __name__ == "__main__":
     unittest.main()

test/test_backpressure.py

@@ -98,6 +98,37 @@ def test_retry_overload_error_insert_one(self):
 
         self.assertIn("Retryable", str(error.exception))
 
+    @client_context.require_failCommand_appName
+    def test_retry_overload_error_getMore(self):
+        coll = self.db.t
+        coll.insert_many([{"x": 1} for _ in range(10)])
+
+        # Ensure command is retried on overload error.
+        fail_once = {
+            "configureFailPoint": "failCommand",
+            "mode": {"times": _MAX_RETRIES},
+            "data": {
+                "failCommands": ["getMore"],
+                "errorCode": 462,  # IngressRequestRateLimitExceeded
+                "errorLabels": ["Retryable"],
+            },
+        }
+        cursor = coll.find(batch_size=2)
+        cursor.next()
+        with self.fail_point(fail_once):
+            cursor.to_list()
+
+        # Ensure command stops retrying after _MAX_RETRIES.
+        fail_many_times = fail_once.copy()
+        fail_many_times["mode"] = {"times": _MAX_RETRIES + 1}
+        cursor = coll.find(batch_size=2)
+        cursor.next()
+        with self.fail_point(fail_many_times):
+            with self.assertRaises(PyMongoError) as error:
+                cursor.to_list()
+
+        self.assertIn("Retryable", str(error.exception))
+
 
 if __name__ == "__main__":
     unittest.main()