Spike

Author: NoahStapp

pymongo/asynchronous/network.py

@@ -15,6 +15,7 @@
 """Internal network layer helper methods."""
 from __future__ import annotations
 
+import asyncio
 import datetime
 import logging
 import statistics
@@ -200,23 +201,23 @@ async def command_stream(
 
     try:
         write_start = time.monotonic()
-        await async_sendall_stream(conn.conn[1], msg)
+        await async_sendall_stream(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_stream(conn.conn[0], request_id)
+            reply = await receive_message_stream(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)}")
+            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)}")
             conn.more_to_come = reply.more_to_come
             unpacked_docs = reply.unpack_response(
                 codec_options=codec_options, user_fields=user_fields
@@ -316,7 +317,7 @@ async def command_stream(
 
 
 async def receive_message_stream(
-    conn: StreamReader, request_id: Optional[int], max_message_size: int = MAX_MESSAGE_SIZE
+    conn: AsyncConnectionStream, 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():
@@ -329,39 +330,40 @@ async def receive_message_stream(
     #         deadline = None
     deadline = None
     # Ignore the response's request id.
-    read_start = time.monotonic()
-    length, _, response_to, op_code = _UNPACK_HEADER(await async_receive_data_stream(conn, 16, deadline))
-    read_elapsed = time.monotonic() - read_start
-    # print(f"Read header in {read_elapsed}")
-    # 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_stream(conn, 9, deadline)
-        )
-        data = decompress(await async_receive_data_stream(conn, length - 25, deadline), compressor_id)
-    else:
-        read_start = time.monotonic()
-        data = await async_receive_data_stream(conn, length - 16, deadline)
-        read_elapsed = time.monotonic() - read_start
-        # print(f"Read body in {read_elapsed}")
+    loop = asyncio.get_running_loop()
+    done = loop.create_future()
+    mv = memoryview(bytearray(max_message_size))
+    conn.conn[1].reset(mv, done)
+    await asyncio.wait_for(done, timeout=None)
+    length, op_code = done.result()
+
+    # length, _, response_to, op_code = _UNPACK_HEADER(await async_receive_data_stream(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_stream(conn, 9, deadline)
+    #     )
+    #     data = decompress(await async_receive_data_stream(conn, length - 25, deadline), compressor_id)
+    # else:
+    #     data = await async_receive_data_stream(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)
+    return unpack_reply(mv[16:length])
 

pymongo/asynchronous/pool.py

@@ -80,7 +80,7 @@
     ConnectionCheckOutFailedReason,
     ConnectionClosedReason,
 )
-from pymongo.network_layer import async_sendall_stream
+from pymongo.network_layer import async_sendall_stream, _UNPACK_HEADER, PyMongoProtocol
 from pymongo.pool_options import PoolOptions
 from pymongo.read_preferences import ReadPreference
 from pymongo.server_api import _add_to_command
@@ -794,7 +794,7 @@ class AsyncConnectionStream:
     """
 
     def __init__(
-        self, conn: tuple[asyncio.StreamReader, asyncio.StreamWriter], 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
@@ -1107,7 +1107,7 @@ async def send_message(self, message: bytes, max_doc_size: int) -> None:
             )
 
         try:
-            await async_sendall_stream(self.conn[1], message)
+            await async_sendall_stream(self.conn, message)
         except BaseException as error:
             self._raise_connection_failure(error)
 
@@ -1117,7 +1117,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_stream(self.conn[0], request_id, self.max_message_size)
+            return await receive_message_stream(self.conn, request_id, self.max_message_size)
         except BaseException as error:
             self._raise_connection_failure(error)
 
@@ -1235,15 +1235,15 @@ def _close_conn(self) -> None:
         # Note: We catch exceptions to avoid spurious errors on interpreter
         # shutdown.
         try:
-            self.conn[1].close()
+            self.conn[0].close()
         except asyncio.CancelledError:
             raise
         except Exception:  # noqa: S110
             pass
 
     def conn_closed(self) -> bool:
         """Return True if we know socket has been closed, False otherwise."""
-        return self.conn[1].is_closing()
+        return self.conn[0].is_closing()
 
     def send_cluster_time(
         self,
@@ -1315,11 +1315,6 @@ def __repr__(self) -> str:
         )
 
 
-async def _create_connection_stream(address: _Address, options: PoolOptions) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
-    """Given (host, port) and PoolOptions, connect and return a paired StreamReader and StreamWriter.
-    """
-    sock = _create_connection(address, options)
-    return await asyncio.open_connection(sock=sock)
 
 
 def _create_connection(address: _Address, options: PoolOptions) -> socket.socket:
@@ -1395,34 +1390,31 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
 
 async def _configured_stream(
     address: _Address, options: PoolOptions
-) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+) -> tuple[asyncio.BaseTransport, PyMongoProtocol]:
     """Given (host, port) and PoolOptions, return a configured socket.
 
     Can raise socket.error, ConnectionFailure, or _CertificateError.
 
     Sets socket's SSL and timeout options.
     """
-    (reader, writer) = await _create_connection_stream(address, options)
+    sock = _create_connection(address, options)
     ssl_context = options._ssl_context
 
     if ssl_context is None:
-        # sock.settimeout(options.socket_timeout)
-        return reader, writer
+        return await asyncio.get_running_loop().create_connection(lambda: PyMongoProtocol(), sock=sock)
 
     host = address[0]
     try:
         # We have to pass hostname / ip address to wrap_socket
         # to use SSLContext.check_hostname.
-        await writer.start_tls(ssl_context, server_hostname=host)
+        transport, protocol = await asyncio.get_running_loop().create_connection(lambda: PyMongoProtocol(), sock=sock, server_hostname=host, ssl=ssl_context)
     except _CertificateError:
-        writer.close()
-        await writer.wait_closed()
+        transport.close()
         # Raise _CertificateError directly like we do after match_hostname
         # below.
         raise
     except (OSError, SSLError) as exc:
-        writer.close()
-        await writer.wait_closed()
+        transport.close()
         # We raise AutoReconnect for transient and permanent SSL handshake
         # failures alike. Permanent handshake failures, like protocol
         # mismatch, will be turned into ServerSelectionTimeoutErrors later.
@@ -1434,13 +1426,12 @@ async def _configured_stream(
         and not options.tls_allow_invalid_hostnames
     ):
         try:
-            ssl.match_hostname(writer.get_extra_info("peercert"), hostname=host)  # type:ignore[attr-defined]
+            ssl.match_hostname(transport.get_extra_info("peercert"), hostname=host)  # type:ignore[attr-defined]
         except _CertificateError:
-            writer.close()
-            await writer.wait_closed()
+            transport.close()
             raise
 
-    return reader, writer
+    return transport, protocol
 
 
 async def _configured_socket(

pymongo/network_layer.py

@@ -70,10 +70,50 @@
 BLOCKING_IO_ERRORS = (BlockingIOError, BLOCKING_IO_LOOKUP_ERROR, *ssl_support.BLOCKING_IO_ERRORS)
 
 
-async def async_sendall_stream(stream: asyncio.StreamWriter, buf: bytes) -> None:
+class PyMongoProtocol(asyncio.Protocol):
+    def __init__(self):
+        self.transport = None
+        self.done = None
+        self.buffer = None
+        self.expected_length = 0
+        self.expecting_header = False
+        self.bytes_read = 0
+        self.op_code = None
+
+    def connection_made(self, transport):
+        self.transport = transport
+
+    def write(self, message: bytes):
+        self.transport.write(message)
+
+    def data_received(self, data):
+        size = len(data)
+        if size == 0:
+            raise OSError("connection closed")
+        self.buffer[self.bytes_read:self.bytes_read + size] = data
+        self.bytes_read += size
+        if self.expecting_header:
+            self.expected_length, _, response_to, self.op_code = _UNPACK_HEADER(self.buffer[:16])
+            self.expecting_header = False
+
+        if self.bytes_read == self.expected_length:
+            self.done.set_result((self.expected_length, self.op_code))
+
+    def connection_lost(self, exc):
+        if self.done and not self.done.done():
+            self.done.set_result(True)
+
+    def reset(self, buffer: memoryview, done: asyncio.Future):
+        self.buffer = buffer
+        self.done = done
+        self.bytes_read = 0
+        self.expecting_header = True
+        self.op_code = None
+
+
+async def async_sendall_stream(stream: AsyncConnectionStream, buf: bytes) -> None:
     try:
-        stream.write(buf)
-        await asyncio.wait_for(stream.drain(), timeout=None)
+        stream.conn[1].write(buf)
     except asyncio.TimeoutError as exc:
         # Convert the asyncio.wait_for timeout error to socket.timeout which pool.py understands.
         raise socket.timeout("timed out") from exc
@@ -92,7 +132,7 @@ async def _poll_cancellation(conn: AsyncConnection) -> None:
 
 
 async def async_receive_data_stream(
-    conn: StreamReader, length: int, deadline: Optional[float]
+    conn: AsyncConnectionStream, length: int, deadline: Optional[float]
 ) -> memoryview:
     # sock = conn.conn
     # sock_timeout = sock.gettimeout()
@@ -104,9 +144,13 @@ async def async_receive_data_stream(
     #     timeout = max(deadline - time.monotonic(), 0)
     # else:
     #     timeout = sock_timeout
+    loop = asyncio.get_running_loop()
 
+    done = loop.create_future()
+    conn.conn[1].setup(done, length)
     try:
-        return await asyncio.wait_for(_async_receive_stream(conn, length), timeout=None)
+        await asyncio.wait_for(done, timeout=None)
+        return done.result()
         # read_task = create_task(_async_receive_stream(conn, length))
         # tasks = [read_task, cancellation_task]
         # done, pending = await asyncio.wait(