PYTHON-6262 getMore retries and multi-mongos retryable writes fail in load-balanced mode

.evergreen/scripts/run_tests.py

@@ -85,6 +85,16 @@ def handle_green_framework() -> None:
     LOGGER.info(f"Running tests with {GREEN_FRAMEWORK}...")
 
 
+def handle_load_balancer_marker() -> None:
+    """Select the full default suite plus the load-balancer-marked tests.
+
+    The last ``-m`` wins over the one setup_tests.py wrote, and
+    pytest_conf.py grants the implicit default markers only to unmarked
+    tests, so the expression needs an explicit ``load_balancer`` term.
+    """
+    TEST_ARGS.extend(["-m", "default or default_async or load_balancer"])
+
+
 def handle_c_ext() -> None:
     if platform.python_implementation() != "CPython":
         return
@@ -194,6 +204,8 @@ def run() -> None:
     # Handle green framework first so they can patch modules.
     if GREEN_FRAMEWORK:
         handle_green_framework()
+    elif TEST_NAME == "load_balancer":
+        handle_load_balancer_marker()
 
     # Ensure C extensions if applicable.
     if not os.environ.get("NO_EXT"):

.evergreen/scripts/test_generate_config.py

@@ -282,6 +282,24 @@ def test_non_standard_cov_tasks_have_no_cov_companions():
             assert companion in no_cov, f"missing no-cov companion for {task.name}"
 
 
+def test_load_balancer_marker_selection(monkeypatch):
+    """The load balancer tasks select the load_balancer marker explicitly:
+    pytest_conf.py only grants the implicit default markers to unmarked
+    tests, so the marker expression must add the full default suite
+    alongside the load balancer tests (PYTHON-6262).
+    """
+    import run_tests
+
+    monkeypatch.setattr(run_tests, "TEST_NAME", "load_balancer")
+    monkeypatch.setattr(run_tests, "GREEN_FRAMEWORK", None)
+    monkeypatch.setattr(run_tests, "TEST_ARGS", ["-m", "load_balancer"])
+    run_tests.handle_load_balancer_marker()
+    assert run_tests.TEST_ARGS[-2:] == [
+        "-m",
+        "default or default_async or load_balancer",
+    ]
+
+
 def test_sanitizer_task_wiring():
     """PYTHON-6048: one pr-tagged sanitizer task, selected by exactly one
     variant that pins MAX_VERSION and runs on a large host for ASan's memory

doc/changelog.rst

@@ -52,6 +52,9 @@ Bug fixes
 - Fixed undefined behavior in the C extension's BSON buffer growth: growing
   an encode buffer past 1GiB doubled a signed integer past its limit.
   Over-limit encodes still raise ``ValueError`` (`PYTHON-6252`_).
+- Fixed a bug where ``getMore`` commands in load-balanced mode were not
+  retried when the server returned a retryable overload error
+  (`PYTHON-6262`_).
 
 .. _PYTHON-5929: https://jira.mongodb.org/browse/PYTHON-5929
 .. _PYTHON-5959: https://jira.mongodb.org/browse/PYTHON-5959
@@ -62,6 +65,7 @@ Bug fixes
 .. _PYTHON-6182: https://jira.mongodb.org/browse/PYTHON-6182
 .. _PYTHON-6226: https://jira.mongodb.org/browse/PYTHON-6226
 .. _PYTHON-6252: https://jira.mongodb.org/browse/PYTHON-6252
+.. _PYTHON-6262: https://jira.mongodb.org/browse/PYTHON-6262
 
 Changes in Version 4.19.0 (2026/10/14)
 --------------------------------------

pymongo/asynchronous/mongo_client.py

@@ -1979,24 +1979,6 @@ async def _run_operation(
         :param address: Optional address when sending a message
             to a specific server, used for getMore.
         """
-        if operation.conn_mgr:
-            server = await self._select_server(
-                operation.read_preference,
-                operation.session,  # type: ignore[arg-type]
-                operation.name,
-                address=address,
-            )
-
-            async with operation.conn_mgr._lock:
-                async with _ClientCheckout.for_existing_conn(
-                    self,
-                    server,
-                    operation.session,  # type: ignore[arg-type]
-                    operation.conn_mgr.conn,
-                ):
-                    return await run_with_conn(
-                        operation.conn_mgr.conn, operation, operation.read_preference
-                    )
 
         async def _cmd(
             _session: AsyncClientSession | None,
@@ -2014,6 +1996,7 @@ async def _cmd(
             address=address,
             retryable=isinstance(operation, _Query),
             operation=operation.name,
+            conn_mgr=operation.conn_mgr,
         )
 
     async def _retry_with_session(
@@ -2060,6 +2043,7 @@ async def _retry_internal(
         operation_id: int | None = None,
         is_run_command: bool = False,
         is_aggregate_write: bool = False,
+        conn_mgr: Any = None,
     ) -> T:
         """Internal retryable helper for all client transactions.
 
@@ -2074,6 +2058,7 @@ async def _retry_internal(
         :param is_run_command: If this is a runCommand operation, defaults to False
         :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
         :param operation_id: Stable operation id shared across retries, defaults to None
+        :param conn_mgr: Pinned connection manager, defaults to None
 
         :return: Output of the calling func()
         """
@@ -2090,6 +2075,7 @@ async def _retry_internal(
             operation_id=operation_id,
             is_run_command=is_run_command,
             is_aggregate_write=is_aggregate_write,
+            conn_mgr=conn_mgr,
         ).run()
 
     async def _retryable_read(
@@ -2103,6 +2089,7 @@ async def _retryable_read(
         operation_id: int | None = None,
         is_run_command: bool = False,
         is_aggregate_write: bool = False,
+        conn_mgr: Any = None,
     ) -> T:
         """Execute an operation with consecutive retries if possible
 
@@ -2121,6 +2108,7 @@ async def _retryable_read(
         :param is_run_command: If this is a runCommand operation, defaults to False.
         :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
         :param operation_id: Stable operation id shared across retries, defaults to None
+        :param conn_mgr: Pinned connection manager, defaults to None
         """
 
         # Ensure that the client supports retrying on reads and there is no session in
@@ -2141,6 +2129,7 @@ async def _retryable_read(
                 operation_id=operation_id,
                 is_run_command=is_run_command,
                 is_aggregate_write=is_aggregate_write,
+                conn_mgr=conn_mgr,
             )
 
     async def _retryable_write(
@@ -2954,6 +2943,11 @@ async def run(self) -> T:
                             )
                         ):
                             raise
+                        if self._conn_mgr is not None and self._conn_mgr.more_to_come:
+                            # A receive-only exhaust getMore has no request
+                            # in flight to retry: the error reply ends the
+                            # stream.
+                            raise
                         self._retrying = True
                         self._last_error = exc_to_check
                         self._attempt_number += 1
@@ -3092,19 +3086,30 @@ async def _read(self) -> T:
         """
         self._server = await self._get_server()
         assert self._read_pref is not None, "Read Preference required on read calls"
+        if self._conn_mgr is not None:
+            # The cursor is pinned to a connection (e.g. load-balanced mode):
+            # retry on that same connection.
+            async with self._conn_mgr._lock:
+                async with _ClientCheckout.for_existing_conn(
+                    self._client, self._server, self._session, self._conn_mgr.conn
+                ):
+                    return await self._run_read_func(self._conn_mgr.conn, self._read_pref)
         async with self._client._conn_from_server(self._read_pref, self._server, self._session) as (
             conn,
             read_pref,
         ):
-            if self._retrying and not self._retryable and not self._always_retryable:
-                self._check_last_error()
-            if self._retrying:
-                self._log_retry(is_write=False)
-            # One operation id across all attempts of this operation if APM/logging is enabled
-            if self._operation_id is None:
-                return await self._func(self._session, self._server, conn, read_pref)  # type: ignore
-            with _op_id._OpIdContext(self._operation_id):
-                return await self._func(self._session, self._server, conn, read_pref)  # type: ignore
+            return await self._run_read_func(conn, read_pref)
+
+    async def _run_read_func(self, conn: AsyncConnection, read_pref: _ServerMode) -> T:
+        if self._retrying and not self._retryable and not self._always_retryable:
+            self._check_last_error()
+        if self._retrying:
+            self._log_retry(is_write=False)
+        # One operation id across all attempts of this operation if APM/logging is enabled
+        if self._operation_id is None:
+            return await self._func(self._session, self._server, conn, read_pref)  # type: ignore
+        with _op_id._OpIdContext(self._operation_id):
+            return await self._func(self._session, self._server, conn, read_pref)  # type: ignore
 
 
 if _HAS_REGISTER_AT_FORK:

pymongo/mongo_client_shared.py

@@ -140,6 +140,7 @@ class _AgnosticClientConnectionRetryableBase(Generic[T, _ServerT, _FuncT]):
         "_base_backoff_ms",
         "_bulk",
         "_client",
+        "_conn_mgr",
         "_deprioritized_servers",
         "_func",
         "_is_aggregate_write",
@@ -172,6 +173,7 @@ def __init__(
         operation_id: int | None = None,
         is_run_command: bool = False,
         is_aggregate_write: bool = False,
+        conn_mgr: Any = None,
     ):
         self._last_error: Exception | None = None
         self._retrying = False
@@ -200,6 +202,7 @@ def __init__(
         self._is_run_command = is_run_command
         self._is_aggregate_write = is_aggregate_write
         self._base_backoff_ms: float | None = None
+        self._conn_mgr = conn_mgr
 
     def _is_not_eligible_for_retry(self) -> bool:
         """Checks if the exchange is not eligible for retry"""

pymongo/synchronous/mongo_client.py

@@ -1973,24 +1973,6 @@ def _run_operation(
         :param address: Optional address when sending a message
             to a specific server, used for getMore.
         """
-        if operation.conn_mgr:
-            server = self._select_server(
-                operation.read_preference,
-                operation.session,  # type: ignore[arg-type]
-                operation.name,
-                address=address,
-            )
-
-            with operation.conn_mgr._lock:
-                with _ClientCheckout.for_existing_conn(
-                    self,
-                    server,
-                    operation.session,  # type: ignore[arg-type]
-                    operation.conn_mgr.conn,
-                ):
-                    return run_with_conn(
-                        operation.conn_mgr.conn, operation, operation.read_preference
-                    )
 
         def _cmd(
             _session: ClientSession | None,
@@ -2008,6 +1990,7 @@ def _cmd(
             address=address,
             retryable=isinstance(operation, _Query),
             operation=operation.name,
+            conn_mgr=operation.conn_mgr,
         )
 
     def _retry_with_session(
@@ -2054,6 +2037,7 @@ def _retry_internal(
         operation_id: int | None = None,
         is_run_command: bool = False,
         is_aggregate_write: bool = False,
+        conn_mgr: Any = None,
     ) -> T:
         """Internal retryable helper for all client transactions.
 
@@ -2068,6 +2052,7 @@ def _retry_internal(
         :param is_run_command: If this is a runCommand operation, defaults to False
         :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
         :param operation_id: Stable operation id shared across retries, defaults to None
+        :param conn_mgr: Pinned connection manager, defaults to None
 
         :return: Output of the calling func()
         """
@@ -2084,6 +2069,7 @@ def _retry_internal(
             operation_id=operation_id,
             is_run_command=is_run_command,
             is_aggregate_write=is_aggregate_write,
+            conn_mgr=conn_mgr,
         ).run()
 
     def _retryable_read(
@@ -2097,6 +2083,7 @@ def _retryable_read(
         operation_id: int | None = None,
         is_run_command: bool = False,
         is_aggregate_write: bool = False,
+        conn_mgr: Any = None,
     ) -> T:
         """Execute an operation with consecutive retries if possible
 
@@ -2115,6 +2102,7 @@ def _retryable_read(
         :param is_run_command: If this is a runCommand operation, defaults to False.
         :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
         :param operation_id: Stable operation id shared across retries, defaults to None
+        :param conn_mgr: Pinned connection manager, defaults to None
         """
 
         # Ensure that the client supports retrying on reads and there is no session in
@@ -2135,6 +2123,7 @@ def _retryable_read(
                 operation_id=operation_id,
                 is_run_command=is_run_command,
                 is_aggregate_write=is_aggregate_write,
+                conn_mgr=conn_mgr,
             )
 
     def _retryable_write(
@@ -2940,6 +2929,11 @@ def run(self) -> T:
                             )
                         ):
                             raise
+                        if self._conn_mgr is not None and self._conn_mgr.more_to_come:
+                            # A receive-only exhaust getMore has no request
+                            # in flight to retry: the error reply ends the
+                            # stream.
+                            raise
                         self._retrying = True
                         self._last_error = exc_to_check
                         self._attempt_number += 1
@@ -3078,19 +3072,30 @@ def _read(self) -> T:
         """
         self._server = self._get_server()
         assert self._read_pref is not None, "Read Preference required on read calls"
+        if self._conn_mgr is not None:
+            # The cursor is pinned to a connection (e.g. load-balanced mode):
+            # retry on that same connection.
+            with self._conn_mgr._lock:
+                with _ClientCheckout.for_existing_conn(
+                    self._client, self._server, self._session, self._conn_mgr.conn
+                ):
+                    return self._run_read_func(self._conn_mgr.conn, self._read_pref)
         with self._client._conn_from_server(self._read_pref, self._server, self._session) as (
             conn,
             read_pref,
         ):
-            if self._retrying and not self._retryable and not self._always_retryable:
-                self._check_last_error()
-            if self._retrying:
-                self._log_retry(is_write=False)
-            # One operation id across all attempts of this operation if APM/logging is enabled
-            if self._operation_id is None:
-                return self._func(self._session, self._server, conn, read_pref)  # type: ignore
-            with _op_id._OpIdContext(self._operation_id):
-                return self._func(self._session, self._server, conn, read_pref)  # type: ignore
+            return self._run_read_func(conn, read_pref)
+
+    def _run_read_func(self, conn: Connection, read_pref: _ServerMode) -> T:
+        if self._retrying and not self._retryable and not self._always_retryable:
+            self._check_last_error()
+        if self._retrying:
+            self._log_retry(is_write=False)
+        # One operation id across all attempts of this operation if APM/logging is enabled
+        if self._operation_id is None:
+            return self._func(self._session, self._server, conn, read_pref)  # type: ignore
+        with _op_id._OpIdContext(self._operation_id):
+            return self._func(self._session, self._server, conn, read_pref)  # type: ignore
 
 
 if _HAS_REGISTER_AT_FORK:

test/__init__.py

@@ -787,6 +787,18 @@ def require_exhaust_cursors(self, func):
             func=func,
         )
 
+    def require_exhaust_stream(self, func):
+        """Run a test only if the server streams exhaust getMores.
+
+        mongos streams exhaust getMores from 7.1 (SERVER-57297); older
+        mongos reply per batch, so the receive-only state is never reached.
+        """
+        return self._require(
+            lambda: not (self.is_mongos or self.load_balancer) or self.version.at_least(7, 1),
+            "Exhaust streams require mongos 7.1+",
+            func=func,
+        )
+
     def supports_transactions(self):
         return self.is_mongos or self.is_rs
 

test/asynchronous/__init__.py

@@ -787,6 +787,18 @@ def require_exhaust_cursors(self, func):
             func=func,
         )
 
+    def require_exhaust_stream(self, func):
+        """Run a test only if the server streams exhaust getMores.
+
+        mongos streams exhaust getMores from 7.1 (SERVER-57297); older
+        mongos reply per batch, so the receive-only state is never reached.
+        """
+        return self._require(
+            lambda: not (self.is_mongos or self.load_balancer) or self.version.at_least(7, 1),
+            "Exhaust streams require mongos 7.1+",
+            func=func,
+        )
+
     def supports_transactions(self):
         return self.is_mongos or self.is_rs