PYTHON-4493 - Use asyncio protocols instead of sockets for network IO (mongodb#2151)

Co-authored-by: Shane Harvey <[email protected]>

Author: NoahStapp

pymongo/asynchronous/encryption.py

@@ -64,11 +64,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,
-    _get_timeout_details,
-    _raise_connection_failure,
-)
 from pymongo.common import CONNECT_TIMEOUT
 from pymongo.daemon import _spawn_daemon
 from pymongo.encryption_options import AutoEncryptionOpts, RangeOpts
@@ -80,12 +75,17 @@
     NetworkTimeout,
     ServerSelectionTimeoutError,
 )
-from pymongo.network_layer import BLOCKING_IO_ERRORS, async_sendall
+from pymongo.network_layer import async_socket_sendall
 from pymongo.operations import UpdateOne
 from pymongo.pool_options import PoolOptions
+from pymongo.pool_shared import (
+    _async_configured_socket,
+    _get_timeout_details,
+    _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_shared import parse_host
 from pymongo.write_concern import WriteConcern
@@ -113,7 +113,7 @@
 
 async def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]:
     try:
-        return await _configured_socket(address, opts)
+        return await _async_configured_socket(address, opts)
     except Exception as exc:
         _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
 
@@ -196,7 +196,7 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
         try:
             conn = await _connect_kms(address, opts)
             try:
-                await async_sendall(conn, message)
+                await async_socket_sendall(conn, message)
                 while kms_context.bytes_needed > 0:
                     # CSOT: update timeout.
                     conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))

pymongo/asynchronous/mongo_client.py

@@ -2079,7 +2079,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/monitor.py

@@ -36,7 +36,11 @@
 from pymongo.server_description import ServerDescription
 
 if TYPE_CHECKING:
-    from pymongo.asynchronous.pool import AsyncConnection, Pool, _CancellationContext
+    from pymongo.asynchronous.pool import (  # type: ignore[attr-defined]
+        AsyncConnection,
+        Pool,
+        _CancellationContext,
+    )
     from pymongo.asynchronous.settings import TopologySettings
     from pymongo.asynchronous.topology import Topology
 

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,13 @@ 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)
+            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
@@ -301,47 +296,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)