Final POC for Wire Protocol-handling protocol

Author: NoahStapp

pymongo/asynchronous/network.py

@@ -15,12 +15,8 @@
 """Internal network layer helper methods."""
 from __future__ import annotations
 
-import asyncio
 import datetime
 import logging
-import statistics
-import time
-from asyncio import streams, StreamReader
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -34,26 +30,24 @@
 
 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_sendall, async_receive_data,
+    async_receive_message,
+    async_sendall,
 )
 
 if TYPE_CHECKING:
     from bson import CodecOptions
     from pymongo.asynchronous.client_session import AsyncClientSession
     from pymongo.asynchronous.mongo_client import AsyncMongoClient
-    from pymongo.asynchronous.pool import AsyncConnection, AsyncStreamConnection, AsyncConnectionProtocol
+    from pymongo.asynchronous.pool import AsyncConnectionProtocol
     from pymongo.compression_support import SnappyContext, ZlibContext, ZstdContext
     from pymongo.monitoring import _EventListeners
     from pymongo.read_concern import ReadConcern
@@ -63,13 +57,8 @@
 
 _IS_SYNC = False
 
-# TOTAL = []
-# TOTAL_WRITE = []
-# TOTAL_READ = []
-# print(f"TOTALS: {TOTAL, TOTAL_WRITE, TOTAL_READ}")
 
-
-async def command_stream(
+async def command(
     conn: AsyncConnectionProtocol,
     dbname: str,
     spec: MutableMapping[str, Any],
@@ -200,24 +189,13 @@ async def command_stream(
         )
 
     try:
-        write_start = time.monotonic()
         await async_sendall(conn, msg)
-        write_elapsed = time.monotonic() - write_start
         if use_op_msg and unacknowledged:
             # Unacknowledged, fake a successful command response.
             reply = None
             response_doc: _DocumentOut = {"ok": 1}
         else:
-            read_start = time.monotonic()
-            reply = await receive_message(conn, request_id)
-            read_elapsed = time.monotonic() - read_start
-            # if name == "insert":
-            #     TOTAL.append(write_elapsed + read_elapsed)
-            #     TOTAL_READ.append(read_elapsed)
-            #     TOTAL_WRITE.append(write_elapsed)
-            # if name == "endSessions":
-            #     print(
-            #         f"AVERAGE READ: {statistics.mean(TOTAL_READ)}, AVERAGE WRITE: {statistics.mean(TOTAL_WRITE)}, AVERAGE ELAPSED: {statistics.mean(TOTAL)}")
+            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
@@ -314,48 +292,3 @@ async def command_stream(
         )
 
     return response_doc  # type: ignore[return-value]
-
-
-async def receive_message(
-    conn: AsyncConnectionProtocol, 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
-    deadline = None
-    # Ignore the response's request id.
-    data, op_code = await async_receive_data(conn, 0, deadline)
-    # 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)

pymongo/asynchronous/pool.py

@@ -37,12 +37,11 @@
     Union,
 )
 
-from asyncio import streams
 from bson import DEFAULT_CODEC_OPTIONS
 from pymongo import _csot, helpers_shared
 from pymongo.asynchronous.client_session import _validate_session_write_concern
 from pymongo.asynchronous.helpers import _handle_reauth
-from pymongo.asynchronous.network import command_stream, receive_message
+from pymongo.asynchronous.network import command
 from pymongo.common import (
     MAX_BSON_SIZE,
     MAX_MESSAGE_SIZE,
@@ -80,7 +79,7 @@
     ConnectionCheckOutFailedReason,
     ConnectionClosedReason,
 )
-from pymongo.network_layer import async_sendall, _UNPACK_HEADER, PyMongoProtocol
+from pymongo.network_layer import PyMongoProtocol, async_receive_message, async_sendall
 from pymongo.pool_options import PoolOptions
 from pymongo.read_preferences import ReadPreference
 from pymongo.server_api import _add_to_command
@@ -534,7 +533,7 @@ async def command(
         if self.op_msg_enabled:
             self._raise_if_not_writable(unacknowledged)
         try:
-            return await command_stream(
+            return await command(
                 self,
                 dbname,
                 spec,
@@ -578,7 +577,6 @@ async def send_message(self, message: bytes, max_doc_size: int) -> None:
         try:
             await async_sendall(self.conn, message)
         except BaseException as error:
-            print(error)
             self._raise_connection_failure(error)
 
     async def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _OpMsg]:
@@ -587,7 +585,7 @@ async def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _O
         If any exception is raised, the socket is closed.
         """
         try:
-            return await receive_message(self, request_id, self.max_message_size)
+            return await async_receive_message(self, request_id, self.max_message_size)
         except BaseException as error:
             self._raise_connection_failure(error)
 
@@ -795,7 +793,11 @@ class AsyncConnectionProtocol:
     """
 
     def __init__(
-        self, conn: tuple[asyncio.BaseTransport, PyMongoProtocol], pool: Pool, address: tuple[str, int], id: int
+        self,
+        conn: tuple[asyncio.BaseTransport, PyMongoProtocol],
+        pool: Pool,
+        address: tuple[str, int],
+        id: int,
     ):
         self.pool_ref = weakref.ref(pool)
         self.conn = conn
@@ -1066,7 +1068,7 @@ async def command(
         if self.op_msg_enabled:
             self._raise_if_not_writable(unacknowledged)
         try:
-            return await command_stream(
+            return await command(
                 self,
                 dbname,
                 spec,
@@ -1118,7 +1120,7 @@ async def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _O
         If any exception is raised, the socket is closed.
         """
         try:
-            return await receive_message(self, request_id, self.max_message_size)
+            return await async_receive_message(self, request_id, self.max_message_size)
         except BaseException as error:
             self._raise_connection_failure(error)
 
@@ -1316,8 +1318,6 @@ def __repr__(self) -> str:
         )
 
 
-
-
 def _create_connection(address: _Address, options: PoolOptions) -> socket.socket:
     """Given (host, port) and PoolOptions, connect and return a socket object.
 
@@ -1400,15 +1400,23 @@ async def _configured_stream(
     """
     sock = _create_connection(address, options)
     ssl_context = options._ssl_context
+    timeout = sock.gettimeout()
 
     if ssl_context is None:
-        return await asyncio.get_running_loop().create_connection(lambda: PyMongoProtocol(), sock=sock)
+        return await asyncio.get_running_loop().create_connection(
+            lambda: PyMongoProtocol(timeout=timeout, buffer_size=2**16), sock=sock
+        )
 
     host = address[0]
     try:
         # We have to pass hostname / ip address to wrap_socket
         # to use SSLContext.check_hostname.
-        transport, protocol = await asyncio.get_running_loop().create_connection(lambda: PyMongoProtocol(), sock=sock, server_hostname=host, ssl=ssl_context)
+        transport, protocol = await asyncio.get_running_loop().create_connection(
+            lambda: PyMongoProtocol(timeout=timeout, buffer_size=2**14),
+            sock=sock,
+            server_hostname=host,
+            ssl=ssl_context,
+        )
     except _CertificateError:
         transport.close()
         # Raise _CertificateError directly like we do after match_hostname
@@ -1819,7 +1827,9 @@ async def remove_stale_sockets(self, reference_generation: int) -> None:
                     self.requests -= 1
                     self.size_cond.notify()
 
-    async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> AsyncConnectionProtocol:
+    async def connect(
+        self, handler: Optional[_MongoClientErrorHandler] = None
+    ) -> AsyncConnectionProtocol:
         """Connect to Mongo and return a new AsyncConnection.
 
         Can raise ConnectionFailure.
@@ -1849,7 +1859,7 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
             )
 
         try:
-            sock = await _configured_stream(self.address, self.opts)
+            transport, protocol = await _configured_stream(self.address, self.opts)
         except BaseException as error:
             async with self.lock:
                 self.active_contexts.discard(tmp_context)
@@ -1875,7 +1885,7 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
 
             raise
 
-        conn = AsyncConnectionProtocol(sock, self, self.address, conn_id)  # type: ignore[arg-type]
+        conn = AsyncConnectionProtocol((transport, protocol), self, self.address, conn_id)  # type: ignore[arg-type]
         async with self.lock:
             self.active_contexts.add(conn.cancel_context)
             self.active_contexts.discard(tmp_context)

pymongo/asynchronous/server.py

@@ -16,8 +16,6 @@
 from __future__ import annotations
 
 import logging
-import statistics
-import time
 from datetime import datetime
 from typing import (
     TYPE_CHECKING,
@@ -60,12 +58,6 @@
 _CURSOR_DOC_FIELDS = {"cursor": {"firstBatch": 1, "nextBatch": 1}}
 
 
-# TOTAL = []
-# TOTAL_WRITE = []
-# TOTAL_READ = []
-# print(f"TOTALS: {TOTAL, TOTAL_WRITE, TOTAL_READ}")
-
-
 class Server:
     def __init__(
         self,
@@ -212,19 +204,8 @@ async def run_operation(
             if more_to_come:
                 reply = await conn.receive_message(None)
             else:
-                write_start = time.monotonic()
                 await conn.send_message(data, max_doc_size)
-                write_elapsed = time.monotonic() - write_start
-
-                read_start = time.monotonic()
                 reply = await conn.receive_message(request_id)
-                read_elapsed = time.monotonic() - read_start
-
-                # TOTAL.append(write_elapsed + read_elapsed)
-                # TOTAL_READ.append(read_elapsed)
-                # TOTAL_WRITE.append(write_elapsed)
-                # print(
-                #     f"AVERAGE READ: {statistics.mean(TOTAL_READ)}, AVERAGE WRITE: {statistics.mean(TOTAL_WRITE)}, AVERAGE ELAPSED: {statistics.mean(TOTAL)}")
 
             # Unpack and check for command errors.
             if use_cmd:

pymongo/connection.py

@@ -21,7 +21,6 @@
 """
 from __future__ import annotations
 
-import asyncio
 import datetime
 import random
 import struct
@@ -1547,7 +1546,9 @@ def unpack(cls, msg: bytes) -> _OpMsg:
             raise ProtocolError(f"Unsupported OP_MSG payload type: 0x{first_payload_type:x}")
 
         if len(msg) != first_payload_size + 5:
-            raise ProtocolError(f"Unsupported OP_MSG reply: >1 section, {len(msg)} vs {first_payload_size + 5}")
+            raise ProtocolError(
+                f"Unsupported OP_MSG reply: >1 section, {len(msg)} vs {first_payload_size + 5}"
+            )
 
         payload_document = msg[5:]
         return cls(flags, payload_document)