debugging

NoahStapp · NoahStapp · commit 39dcb1ef0950 · 2024-12-04T12:04:23.000-05:00
diff --git a/pymongo/asynchronous/network.py b/pymongo/asynchronous/network.py
@@ -17,6 +17,7 @@
 
 import datetime
 import logging
+import statistics
 import time
 from asyncio import streams, StreamReader
 from typing import (
@@ -61,6 +62,11 @@
 
 _IS_SYNC = False
 
+TOTAL = []
+TOTAL_WRITE = []
+TOTAL_READ = []
+# print(f"TOTALS: {TOTAL, TOTAL_WRITE, TOTAL_READ}")
+
 
 async def command_stream(
     conn: AsyncConnectionStream,
@@ -113,7 +119,6 @@ async def command_stream(
         bson._decode_all_selective.
     :param exhaust_allowed: True if we should enable OP_MSG exhaustAllowed.
     """
-    # print("Running stream command!")
     name = next(iter(spec))
     ns = dbname + ".$cmd"
     speculative_hello = False
@@ -194,13 +199,24 @@ async def command_stream(
         )
 
     try:
+        write_start = time.monotonic()
         await async_sendall_stream(conn.conn[1], 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)
+            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)}")
             conn.more_to_come = reply.more_to_come
             unpacked_docs = reply.unpack_response(
                 codec_options=codec_options, user_fields=user_fields
@@ -313,7 +329,10 @@ 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:
@@ -333,7 +352,10 @@ async def receive_message_stream(
         )
         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}")
 
     try:
         unpack_reply = _UNPACK_REPLY[op_code]
diff --git a/pymongo/asynchronous/server.py b/pymongo/asynchronous/server.py
@@ -16,6 +16,8 @@
 from __future__ import annotations
 
 import logging
+import statistics
+import time
 from datetime import datetime
 from typing import (
     TYPE_CHECKING,
@@ -58,6 +60,12 @@
 _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,
@@ -204,8 +212,19 @@ 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:
diff --git a/pymongo/network_layer.py b/pymongo/network_layer.py
@@ -18,6 +18,7 @@
 import asyncio
 import errno
 import socket
+import statistics
 import struct
 import sys
 import time
@@ -68,6 +69,7 @@
 # Errors raised by sockets (and TLS sockets) when in non-blocking mode.
 BLOCKING_IO_ERRORS = (BlockingIOError, BLOCKING_IO_LOOKUP_ERROR, *ssl_support.BLOCKING_IO_ERRORS)
 
+
 async def async_sendall_stream(stream: asyncio.StreamWriter, buf: bytes) -> None:
     try:
         stream.write(buf)
@@ -77,161 +79,14 @@ async def async_sendall_stream(stream: asyncio.StreamWriter, buf: bytes) -> None
         raise socket.timeout("timed out") from exc
 
 
-if sys.platform != "win32":
-
-    async def _async_sendall_ssl(
-        sock: Union[socket.socket, _sslConn], buf: bytes, loop: AbstractEventLoop
-    ) -> None:
-        view = memoryview(buf)
-        sent = 0
-
-        def _is_ready(fut: Future) -> None:
-            if fut.done():
-                return
-            fut.set_result(None)
-
-        while sent < len(buf):
-            try:
-                sent += sock.send(view[sent:])
-            except BLOCKING_IO_ERRORS as exc:
-                fd = sock.fileno()
-                # Check for closed socket.
-                if fd == -1:
-                    raise SSLError("Underlying socket has been closed") from None
-                if isinstance(exc, BLOCKING_IO_READ_ERROR):
-                    fut = loop.create_future()
-                    loop.add_reader(fd, _is_ready, fut)
-                    try:
-                        await fut
-                    finally:
-                        loop.remove_reader(fd)
-                if isinstance(exc, BLOCKING_IO_WRITE_ERROR):
-                    fut = loop.create_future()
-                    loop.add_writer(fd, _is_ready, fut)
-                    try:
-                        await fut
-                    finally:
-                        loop.remove_writer(fd)
-                if _HAVE_PYOPENSSL and isinstance(exc, BLOCKING_IO_LOOKUP_ERROR):
-                    fut = loop.create_future()
-                    loop.add_reader(fd, _is_ready, fut)
-                    try:
-                        loop.add_writer(fd, _is_ready, fut)
-                        await fut
-                    finally:
-                        loop.remove_reader(fd)
-                        loop.remove_writer(fd)
-
-    async def _async_receive_ssl(
-        conn: _sslConn, length: int, loop: AbstractEventLoop, once: Optional[bool] = False
-    ) -> memoryview:
-        mv = memoryview(bytearray(length))
-        total_read = 0
-
-        def _is_ready(fut: Future) -> None:
-            if fut.done():
-                return
-            fut.set_result(None)
-
-        while total_read < length:
-            try:
-                read = conn.recv_into(mv[total_read:])
-                if read == 0:
-                    raise OSError("connection closed")
-                # KMS responses update their expected size after the first batch, stop reading after one loop
-                if once:
-                    return mv[:read]
-                total_read += read
-            except BLOCKING_IO_ERRORS as exc:
-                fd = conn.fileno()
-                # Check for closed socket.
-                if fd == -1:
-                    raise SSLError("Underlying socket has been closed") from None
-                if isinstance(exc, BLOCKING_IO_READ_ERROR):
-                    fut = loop.create_future()
-                    loop.add_reader(fd, _is_ready, fut)
-                    try:
-                        await fut
-                    finally:
-                        loop.remove_reader(fd)
-                if isinstance(exc, BLOCKING_IO_WRITE_ERROR):
-                    fut = loop.create_future()
-                    loop.add_writer(fd, _is_ready, fut)
-                    try:
-                        await fut
-                    finally:
-                        loop.remove_writer(fd)
-                if _HAVE_PYOPENSSL and isinstance(exc, BLOCKING_IO_LOOKUP_ERROR):
-                    fut = loop.create_future()
-                    loop.add_reader(fd, _is_ready, fut)
-                    try:
-                        loop.add_writer(fd, _is_ready, fut)
-                        await fut
-                    finally:
-                        loop.remove_reader(fd)
-                        loop.remove_writer(fd)
-        return mv
-
-else:
-    # The default Windows asyncio event loop does not support loop.add_reader/add_writer:
-    # https://docs.python.org/3/library/asyncio-platforms.html#asyncio-platform-support
-    # Note: In PYTHON-4493 we plan to replace this code with asyncio streams.
-    async def _async_sendall_ssl(
-        sock: Union[socket.socket, _sslConn], buf: bytes, dummy: AbstractEventLoop
-    ) -> None:
-        view = memoryview(buf)
-        total_length = len(buf)
-        total_sent = 0
-        # Backoff starts at 1ms, doubles on timeout up to 512ms, and halves on success
-        # down to 1ms.
-        backoff = 0.001
-        while total_sent < total_length:
-            try:
-                sent = sock.send(view[total_sent:])
-            except BLOCKING_IO_ERRORS:
-                await asyncio.sleep(backoff)
-                sent = 0
-            if sent > 0:
-                backoff = max(backoff / 2, 0.001)
-            else:
-                backoff = min(backoff * 2, 0.512)
-            total_sent += sent
-
-    async def _async_receive_ssl(
-        conn: _sslConn, length: int, dummy: AbstractEventLoop, once: Optional[bool] = False
-    ) -> memoryview:
-        mv = memoryview(bytearray(length))
-        total_read = 0
-        # Backoff starts at 1ms, doubles on timeout up to 512ms, and halves on success
-        # down to 1ms.
-        backoff = 0.001
-        while total_read < length:
-            try:
-                read = conn.recv_into(mv[total_read:])
-                if read == 0:
-                    raise OSError("connection closed")
-                # KMS responses update their expected size after the first batch, stop reading after one loop
-                if once:
-                    return mv[:read]
-            except BLOCKING_IO_ERRORS:
-                await asyncio.sleep(backoff)
-                read = 0
-            if read > 0:
-                backoff = max(backoff / 2, 0.001)
-            else:
-                backoff = min(backoff * 2, 0.512)
-            total_read += read
-        return mv
-
-
 def sendall(sock: Union[socket.socket, _sslConn], buf: bytes) -> None:
     sock.sendall(buf)
 
 
 async def _poll_cancellation(conn: AsyncConnection) -> None:
-    # while True:
-    #     if conn.cancel_context.cancelled:
-    #         return
+    while True:
+        if conn.cancel_context.cancelled:
+            return
 
         await asyncio.sleep(_POLL_TIMEOUT)
 
@@ -295,19 +150,11 @@ async def async_receive_data_socket(
         sock.settimeout(sock_timeout)
 
 
-
 async def _async_receive_stream(reader: asyncio.StreamReader, length: int) -> memoryview:
-    mv = bytearray(length)
-    total_read = 0
-
-    while total_read < length:
-        bytes = await reader.read(length)
-        chunk_length = len(bytes)
-        if chunk_length == 0:
-            raise OSError("connection closed")
-        mv[total_read:] = bytes
-        total_read += chunk_length
-    return memoryview(mv)
+    try:
+        return memoryview(await reader.readexactly(length))
+    except asyncio.IncompleteReadError:
+        raise OSError("connection closed")
 
 def receive_data(conn: Connection, length: int, deadline: Optional[float]) -> memoryview:
     buf = bytearray(length)
