PYTHON-4324 CSOT avoid connection churn when operations timeout

ShaneHarvey · ShaneHarvey · commit 2a394e71b4d4 · 2025-04-08T16:59:39.000-07:00
diff --git a/pymongo/asynchronous/network.py b/pymongo/asynchronous/network.py
@@ -138,8 +138,9 @@ async def command(
         spec = orig = await client._encrypter.encrypt(dbname, spec, codec_options)
 
     # Support CSOT
+    applied_csot = False
     if client:
-        conn.apply_timeout(client, spec)
+        applied_csot = conn.apply_timeout(client, spec)
     _csot.apply_write_concern(spec, write_concern)
 
     if use_op_msg:
@@ -195,7 +196,7 @@ async def command(
             reply = None
             response_doc: _DocumentOut = {"ok": 1}
         else:
-            reply = await async_receive_message(conn, request_id)
+            reply = await async_receive_message(conn, request_id, enable_pending=bool(applied_csot))
             conn.more_to_come = reply.more_to_come
             unpacked_docs = reply.unpack_response(
                 codec_options=codec_options, user_fields=user_fields
diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py
@@ -34,7 +34,7 @@
 )
 
 from bson import DEFAULT_CODEC_OPTIONS
-from pymongo import _csot, helpers_shared
+from pymongo import _csot, helpers_shared, network_layer
 from pymongo.asynchronous.client_session import _validate_session_write_concern
 from pymongo.asynchronous.helpers import _handle_reauth
 from pymongo.asynchronous.network import command
@@ -188,6 +188,41 @@ def __init__(
         self.creation_time = time.monotonic()
         # For gossiping $clusterTime from the connection handshake to the client.
         self._cluster_time = None
+        self.pending_response = False
+        self.pending_bytes = 0
+        self.pending_deadline = 0.0
+
+    def mark_pending(self, nbytes: int) -> None:
+        """Mark this connection as having a pending response."""
+        # TODO: add "if self.enable_pending:"
+        self.pending_response = True
+        self.pending_bytes = nbytes
+        self.pending_deadline = time.monotonic() + 3  # 3 seconds timeout for pending response
+
+    async def complete_pending(self) -> None:
+        """Complete a pending response."""
+        if not self.pending_response:
+            return
+
+        timeout: Optional[Union[float, int]]
+        timeout = self.conn.gettimeout
+        if _csot.get_timeout():
+            deadline = min(_csot.get_deadline(), self.pending_deadline)
+        elif timeout:
+            deadline = min(time.monotonic() + timeout, self.pending_deadline)
+        else:
+            deadline = self.pending_deadline
+
+        if not _IS_SYNC:
+            # In async the reader task reads the whole message at once.
+            # TODO: respect deadline
+            await self.receive_message(None, True)
+        else:
+            # In sync we need to track the bytes left for the message.
+            network_layer.receive_data(self.conn.get_conn, self.pending_byte, deadline)
+        self.pending_response = False
+        self.pending_bytes = 0
+        self.pending_deadline = 0.0
 
     def set_conn_timeout(self, timeout: Optional[float]) -> None:
         """Cache last timeout to avoid duplicate calls to conn.settimeout."""
@@ -454,13 +489,17 @@ async def send_message(self, message: bytes, max_doc_size: int) -> None:
         except BaseException as error:
             await self._raise_connection_failure(error)
 
-    async def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _OpMsg]:
+    async def receive_message(
+        self, request_id: Optional[int], enable_pending: bool = False
+    ) -> Union[_OpReply, _OpMsg]:
         """Receive a raw BSON message or raise ConnectionFailure.
 
         If any exception is raised, the socket is closed.
         """
         try:
-            return await async_receive_message(self, request_id, self.max_message_size)
+            return await async_receive_message(
+                self, request_id, self.max_message_size, enable_pending
+            )
         # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
             await self._raise_connection_failure(error)
@@ -495,7 +534,9 @@ async def write_command(
         :param msg: bytes, the command message.
         """
         await self.send_message(msg, 0)
-        reply = await self.receive_message(request_id)
+        reply = await self.receive_message(
+            request_id, enable_pending=(_csot.get_timeout() is not None)
+        )
         result = reply.command_response(codec_options)
 
         # Raises NotPrimaryError or OperationFailure.
@@ -635,7 +676,10 @@ async def _raise_connection_failure(self, error: BaseException) -> NoReturn:
             reason = None
         else:
             reason = ConnectionClosedReason.ERROR
-        await self.close_conn(reason)
+
+        # Pending connections should be placed back in the pool.
+        if not self.pending_response:
+            await self.close_conn(reason)
         # SSLError from PyOpenSSL inherits directly from Exception.
         if isinstance(error, (IOError, OSError, SSLError)):
             details = _get_timeout_details(self.opts)
@@ -1076,7 +1120,7 @@ async def checkout(
 
         This method should always be used in a with-statement::
 
-            with pool.get_conn() as connection:
+            with pool.checkout() as connection:
                 connection.send_message(msg)
                 data = connection.receive_message(op_code, request_id)
 
@@ -1388,6 +1432,7 @@ async def _perished(self, conn: AsyncConnection) -> bool:
         pool, to keep performance reasonable - we can't avoid AutoReconnects
         completely anyway.
         """
+        await conn.complete_pending()
         idle_time_seconds = conn.idle_time_seconds()
         # If socket is idle, open a new one.
         if (
diff --git a/pymongo/asynchronous/server.py b/pymongo/asynchronous/server.py
@@ -205,7 +205,7 @@ async def run_operation(
                 reply = await conn.receive_message(None)
             else:
                 await conn.send_message(data, max_doc_size)
-                reply = await conn.receive_message(request_id)
+                reply = await conn.receive_message(request_id, operation.pending_enabled())
 
             # Unpack and check for command errors.
             if use_cmd:
diff --git a/pymongo/message.py b/pymongo/message.py
@@ -1569,6 +1569,7 @@ class _Query:
         "allow_disk_use",
         "_as_command",
         "exhaust",
+        "_pending_enabled",
     )
 
     # For compatibility with the _GetMore class.
@@ -1612,6 +1613,10 @@ def __init__(
         self.name = "find"
         self._as_command: Optional[tuple[dict[str, Any], str]] = None
         self.exhaust = exhaust
+        self._pending_enabled = False
+
+    def pending_enabled(self):
+        return self._pending_enabled
 
     def reset(self) -> None:
         self._as_command = None
@@ -1673,7 +1678,9 @@ def as_command(
         conn.send_cluster_time(cmd, self.session, self.client)  # type: ignore[arg-type]
         # Support CSOT
         if apply_timeout:
-            conn.apply_timeout(self.client, cmd=cmd)  # type: ignore[arg-type]
+            res = conn.apply_timeout(self.client, cmd=cmd)  # type: ignore[arg-type]
+            if res is not None:
+                self._pending_enabled = True
         self._as_command = cmd, self.db
         return self._as_command
 
@@ -1747,6 +1754,7 @@ class _GetMore:
         "_as_command",
         "exhaust",
         "comment",
+        "_pending_enabled",
     )
 
     name = "getMore"
@@ -1779,6 +1787,10 @@ def __init__(
         self._as_command: Optional[tuple[dict[str, Any], str]] = None
         self.exhaust = exhaust
         self.comment = comment
+        self._pending_enabled = False
+
+    def pending_enabled(self):
+        return self._pending_enabled
 
     def reset(self) -> None:
         self._as_command = None
@@ -1822,7 +1834,9 @@ def as_command(
         conn.send_cluster_time(cmd, self.session, self.client)  # type: ignore[arg-type]
         # Support CSOT
         if apply_timeout:
-            conn.apply_timeout(self.client, cmd=None)  # type: ignore[arg-type]
+            res = conn.apply_timeout(self.client, cmd=None)  # type: ignore[arg-type]
+            if res is not None:
+                self._pending_enabled = True
         self._as_command = cmd, self.db
         return self._as_command
 
diff --git a/pymongo/network_layer.py b/pymongo/network_layer.py
@@ -325,7 +325,9 @@ def wait_for_read(conn: Connection, deadline: Optional[float]) -> None:
             raise socket.timeout("timed out")
 
 
-def receive_data(conn: Connection, length: int, deadline: Optional[float]) -> memoryview:
+def receive_data(
+    conn: Connection, length: int, deadline: Optional[float], enable_pending: bool = False
+) -> memoryview:
     buf = bytearray(length)
     mv = memoryview(buf)
     bytes_read = 0
@@ -357,12 +359,16 @@ def receive_data(conn: Connection, length: int, deadline: Optional[float]) -> me
                 if conn.cancel_context.cancelled:
                     raise _OperationCancelled("operation cancelled") from None
                 # We reached the true deadline.
+                if enable_pending:
+                    conn.mark_pending(length - bytes_read)
                 raise socket.timeout("timed out") from None
             except socket.timeout:
                 if conn.cancel_context.cancelled:
                     raise _OperationCancelled("operation cancelled") from None
                 if _PYPY:
                     # We reached the true deadline.
+                    if enable_pending:
+                        conn.mark_pending(length - bytes_read)
                     raise
                 continue
             except OSError as exc:
@@ -692,6 +698,7 @@ async def async_receive_message(
     conn: AsyncConnection,
     request_id: Optional[int],
     max_message_size: int = MAX_MESSAGE_SIZE,
+    enable_pending: bool = False,
 ) -> Union[_OpReply, _OpMsg]:
     """Receive a raw BSON message or raise socket.error."""
     timeout: Optional[Union[float, int]]
@@ -721,6 +728,8 @@ async def async_receive_message(
         if pending:
             await asyncio.wait(pending)
         if len(done) == 0:
+            if enable_pending:
+                conn.mark_pending(1)
             raise socket.timeout("timed out")
         if read_task in done:
             data, op_code = read_task.result()
@@ -740,7 +749,10 @@ async def async_receive_message(
 
 
 def receive_message(
-    conn: Connection, request_id: Optional[int], max_message_size: int = MAX_MESSAGE_SIZE
+    conn: Connection,
+    request_id: Optional[int],
+    max_message_size: int = MAX_MESSAGE_SIZE,
+    enable_pending: bool = False,
 ) -> Union[_OpReply, _OpMsg]:
     """Receive a raw BSON message or raise socket.error."""
     if _csot.get_timeout():
@@ -752,7 +764,9 @@ def receive_message(
         else:
             deadline = None
     # Ignore the response's request id.
-    length, _, response_to, op_code = _UNPACK_HEADER(receive_data(conn, 16, deadline))
+    length, _, response_to, op_code = _UNPACK_HEADER(
+        receive_data(conn, 16, deadline, enable_pending)
+    )
     # No request_id for exhaust cursor "getMore".
     if request_id is not None:
         if request_id != response_to:
@@ -767,10 +781,12 @@ def receive_message(
             f"message size ({max_message_size!r})"
         )
     if op_code == 2012:
-        op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(receive_data(conn, 9, deadline))
-        data = decompress(receive_data(conn, length - 25, deadline), compressor_id)
+        op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(
+            receive_data(conn, 9, deadline, enable_pending)
+        )
+        data = decompress(receive_data(conn, length - 25, deadline), compressor_id, enable_pending)
     else:
-        data = receive_data(conn, length - 16, deadline)
+        data = receive_data(conn, length - 16, deadline, enable_pending)
 
     try:
         unpack_reply = _UNPACK_REPLY[op_code]
diff --git a/pymongo/synchronous/network.py b/pymongo/synchronous/network.py
@@ -138,8 +138,9 @@ def command(
         spec = orig = client._encrypter.encrypt(dbname, spec, codec_options)
 
     # Support CSOT
+    applied_csot = False
     if client:
-        conn.apply_timeout(client, spec)
+        applied_csot = conn.apply_timeout(client, spec)
     _csot.apply_write_concern(spec, write_concern)
 
     if use_op_msg:
@@ -195,7 +196,7 @@ def command(
             reply = None
             response_doc: _DocumentOut = {"ok": 1}
         else:
-            reply = receive_message(conn, request_id)
+            reply = receive_message(conn, request_id, enable_pending=bool(applied_csot))
             conn.more_to_come = reply.more_to_come
             unpacked_docs = reply.unpack_response(
                 codec_options=codec_options, user_fields=user_fields
diff --git a/pymongo/synchronous/pool.py b/pymongo/synchronous/pool.py
diff --git a/pymongo/synchronous/server.py b/pymongo/synchronous/server.py
