PYTHON-4636 - Avoid blocking I/O calls in async code paths

mongodb · Sep 19, 2024 · 7fa5d47 · 7fa5d47
1 parent 9df635f
commit 7fa5d47
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/pymongo/asynchronous/network.py b/pymongo/asynchronous/network.py
@@ -355,6 +355,7 @@ async def wait_for_read(conn: AsyncConnection, deadline: Optional[float]) -> Non
     """Block until at least one byte is read, or a timeout, or a cancel."""
     sock = conn.conn
     timed_out = False
+    timeout = _POLL_TIMEOUT
     # Check if the connection's socket has been manually closed
     if sock.fileno() == -1:
         return
@@ -373,16 +374,18 @@ async def wait_for_read(conn: AsyncConnection, deadline: Optional[float]) -> Non
                 if remaining <= 0:
                     timed_out = True
                 timeout = max(min(remaining, _POLL_TIMEOUT), 0)
+            if _IS_SYNC:
+                readable = conn.socket_checker.select(sock, read=True, timeout=timeout)
             else:
-                timeout = _POLL_TIMEOUT
-            readable = conn.socket_checker.select(sock, read=True, timeout=timeout)
+                readable = conn.socket_checker.select(sock, read=True, timeout=0.0001)
         if conn.cancel_context.cancelled:
             raise _OperationCancelled("operation cancelled")
         if readable:
             return
         if timed_out:
             raise socket.timeout("timed out")
-        await asyncio.sleep(0)
+        if not _IS_SYNC:
+            await asyncio.sleep(timeout / 100.0)
 
 
 async def _receive_data_on_socket(

diff --git a/pymongo/synchronous/network.py b/pymongo/synchronous/network.py
@@ -352,6 +352,7 @@ def wait_for_read(conn: Connection, deadline: Optional[float]) -> None:
     """Block until at least one byte is read, or a timeout, or a cancel."""
     sock = conn.conn
     timed_out = False
+    timeout = _POLL_TIMEOUT
     # Check if the connection's socket has been manually closed
     if sock.fileno() == -1:
         return
@@ -370,15 +371,18 @@ def wait_for_read(conn: Connection, deadline: Optional[float]) -> None:
                 if remaining <= 0:
                     timed_out = True
                 timeout = max(min(remaining, _POLL_TIMEOUT), 0)
+            if _IS_SYNC:
+                readable = conn.socket_checker.select(sock, read=True, timeout=timeout)
             else:
-                timeout = _POLL_TIMEOUT
-            readable = conn.socket_checker.select(sock, read=True, timeout=timeout)
+                readable = conn.socket_checker.select(sock, read=True, timeout=0.0001)
         if conn.cancel_context.cancelled:
             raise _OperationCancelled("operation cancelled")
         if readable:
             return
         if timed_out:
             raise socket.timeout("timed out")
+        if not _IS_SYNC:
+            time.sleep(timeout / 100.0)
 
 
 def _receive_data_on_socket(conn: Connection, length: int, deadline: Optional[float]) -> memoryview: