Async protocols spike, WIP exhaust cursor support

pymongo/asynchronous/encryption.py

@@ -63,7 +63,6 @@
 from pymongo.asynchronous.cursor import AsyncCursor
 from pymongo.asynchronous.database import AsyncDatabase
 from pymongo.asynchronous.mongo_client import AsyncMongoClient
-from pymongo.asynchronous.pool import _configured_socket, _raise_connection_failure
 from pymongo.common import CONNECT_TIMEOUT
 from pymongo.daemon import _spawn_daemon
 from pymongo.encryption_options import AutoEncryptionOpts, RangeOpts
@@ -75,12 +74,13 @@
     PyMongoError,
     ServerSelectionTimeoutError,
 )
-from pymongo.network_layer import BLOCKING_IO_ERRORS, async_sendall
+from pymongo.network_layer import async_sendall
 from pymongo.operations import UpdateOne
 from pymongo.pool_options import PoolOptions
+from pymongo.pool_shared import _configured_socket, _raise_connection_failure
 from pymongo.read_concern import ReadConcern
 from pymongo.results import BulkWriteResult, DeleteResult
-from pymongo.ssl_support import get_ssl_context
+from pymongo.ssl_support import BLOCKING_IO_ERRORS, get_ssl_context
 from pymongo.typings import _DocumentType, _DocumentTypeArg
 from pymongo.uri_parser import parse_host
 from pymongo.write_concern import WriteConcern

pymongo/asynchronous/mongo_client.py

@@ -1951,7 +1951,7 @@ async def _cleanup_cursor_lock(
                 # exhausted the result set we *must* close the socket
                 # to stop the server from sending more data.
                 assert conn_mgr.conn is not None
-                conn_mgr.conn.close_conn(ConnectionClosedReason.ERROR)
+                await conn_mgr.conn.close_conn(ConnectionClosedReason.ERROR)
             else:
                 await self._close_cursor_now(cursor_id, address, session=session, conn_mgr=conn_mgr)
         if conn_mgr:

pymongo/asynchronous/network.py

@@ -17,7 +17,6 @@
 
 import datetime
 import logging
-import time
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -31,20 +30,16 @@
 
 from bson import _decode_all_selective
 from pymongo import _csot, helpers_shared, message
-from pymongo.common import MAX_MESSAGE_SIZE
-from pymongo.compression_support import _NO_COMPRESSION, decompress
+from pymongo.compression_support import _NO_COMPRESSION
 from pymongo.errors import (
     NotPrimaryError,
     OperationFailure,
-    ProtocolError,
 )
 from pymongo.logger import _COMMAND_LOGGER, _CommandStatusMessage, _debug_log
-from pymongo.message import _UNPACK_REPLY, _OpMsg, _OpReply
+from pymongo.message import _OpMsg
 from pymongo.monitoring import _is_speculative_authenticate
 from pymongo.network_layer import (
-    _UNPACK_COMPRESSION_HEADER,
-    _UNPACK_HEADER,
-    async_receive_data,
+    async_receive_message,
     async_sendall,
 )
 
@@ -194,13 +189,16 @@ async def command(
         )
 
     try:
-        await async_sendall(conn.conn, msg)
+        await async_sendall(conn.conn.get_conn, msg)
         if use_op_msg and unacknowledged:
             # Unacknowledged, fake a successful command response.
             reply = None
             response_doc: _DocumentOut = {"ok": 1}
         else:
-            reply = await receive_message(conn, request_id)
+            if "dropDatabase" in spec:
+                reply = await async_receive_message(conn, request_id, debug=True)
+            else:
+                reply = await async_receive_message(conn, request_id)
             conn.more_to_come = reply.more_to_come
             unpacked_docs = reply.unpack_response(
                 codec_options=codec_options, user_fields=user_fields
@@ -297,47 +295,3 @@ async def command(
         )
 
     return response_doc  # type: ignore[return-value]
-
-
-async def receive_message(
-    conn: AsyncConnection, request_id: Optional[int], max_message_size: int = MAX_MESSAGE_SIZE
-) -> Union[_OpReply, _OpMsg]:
-    """Receive a raw BSON message or raise socket.error."""
-    if _csot.get_timeout():
-        deadline = _csot.get_deadline()
-    else:
-        timeout = conn.conn.gettimeout()
-        if timeout:
-            deadline = time.monotonic() + timeout
-        else:
-            deadline = None
-    # Ignore the response's request id.
-    length, _, response_to, op_code = _UNPACK_HEADER(await async_receive_data(conn, 16, deadline))
-    # No request_id for exhaust cursor "getMore".
-    if request_id is not None:
-        if request_id != response_to:
-            raise ProtocolError(f"Got response id {response_to!r} but expected {request_id!r}")
-    if length <= 16:
-        raise ProtocolError(
-            f"Message length ({length!r}) not longer than standard message header size (16)"
-        )
-    if length > max_message_size:
-        raise ProtocolError(
-            f"Message length ({length!r}) is larger than server max "
-            f"message size ({max_message_size!r})"
-        )
-    if op_code == 2012:
-        op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(
-            await async_receive_data(conn, 9, deadline)
-        )
-        data = decompress(await async_receive_data(conn, length - 25, deadline), compressor_id)
-    else:
-        data = await async_receive_data(conn, length - 16, deadline)
-
-    try:
-        unpack_reply = _UNPACK_REPLY[op_code]
-    except KeyError:
-        raise ProtocolError(
-            f"Got opcode {op_code!r} but expected {_UNPACK_REPLY.keys()!r}"
-        ) from None
-    return unpack_reply(data)