Prevent deadlock when mocking large async HTTP responses.

Author: mindflayer

mocket/io.py

@@ -3,9 +3,6 @@
 from __future__ import annotations
 
 import io
-import os
-
-from mocket.mocket import Mocket
 
 
 class MocketSocketIO(io.BytesIO):
@@ -21,17 +18,12 @@ def __init__(self, address: tuple) -> None:
         super().__init__()
 
     def write(self, content: bytes) -> int:
-        """Write content to the buffer and the pipe if available.
+        """Write content to the in-memory buffer.
 
         Args:
             content: Bytes to write
 
         Returns:
             Number of bytes written
         """
-        super().write(content)
-
-        _, w_fd = Mocket.get_pair(self._address)
-        if w_fd:
-            os.write(w_fd, content)
-        return len(content)
+        return super().write(content)

mocket/mocket.py

@@ -22,7 +22,10 @@
 class Mocket:
     """Singleton class managing all mock socket operations and entries."""
 
+    _socket_ios: ClassVar[dict[Address, Any]] = {}
     _socket_pairs: ClassVar[dict[Address, tuple[int, int]]] = {}
+    _pipe_uses_data: ClassVar[dict[Address, bool]] = {}
+    _pending_readables: ClassVar[dict[Address, int]] = {}
     _address: ClassVar[Address | tuple[None, None]] = (None, None)
     _entries: ClassVar[dict[Address, list[MocketEntry]]] = collections.defaultdict(list)
     _requests: ClassVar[list] = []
@@ -90,6 +93,39 @@ def set_pair(cls, address: Address, pair: tuple[int, int]) -> None:
         """
         cls._socket_pairs[address] = pair
 
+    @classmethod
+    def get_io(cls, address: Address) -> Any:
+        """Get the shared socket I/O buffer for an address, if any."""
+        return cls._socket_ios.get(address)
+
+    @classmethod
+    def set_io(cls, address: Address, socket_io: Any) -> None:
+        """Store the shared socket I/O buffer for an address."""
+        cls._socket_ios[address] = socket_io
+
+    @classmethod
+    def pipe_uses_data(cls, address: Address) -> bool:
+        """Return whether the pipe for an address carries mirrored response bytes."""
+        return cls._pipe_uses_data.get(address, False)
+
+    @classmethod
+    def set_pipe_uses_data(cls, address: Address, uses_data: bool) -> None:
+        """Store whether the pipe for an address carries mirrored response bytes."""
+        cls._pipe_uses_data[address] = uses_data
+
+    @classmethod
+    def get_pending_readables(cls, address: Address) -> int:
+        """Get the number of readiness bytes currently queued for an address."""
+        return cls._pending_readables.get(address, 0)
+
+    @classmethod
+    def set_pending_readables(cls, address: Address, count: int) -> None:
+        """Store the number of readiness bytes queued for an address."""
+        if count > 0:
+            cls._pending_readables[address] = count
+        else:
+            cls._pending_readables.pop(address, None)
+
     @classmethod
     def register(cls, *entries: MocketEntry) -> None:
         """Register mock entries with Mocket.
@@ -132,10 +168,15 @@ def collect(cls, data: Any) -> None:
     @classmethod
     def reset(cls) -> None:
         """Reset all Mocket state and clean up file descriptors."""
+        for socket_io in cls._socket_ios.values():
+            socket_io.close()
         for r_fd, w_fd in cls._socket_pairs.values():
             os.close(r_fd)
             os.close(w_fd)
+        cls._socket_ios = {}
         cls._socket_pairs = {}
+        cls._pipe_uses_data = {}
+        cls._pending_readables = {}
         cls._entries = collections.defaultdict(list)
         cls._requests = []
         cls._record_storage = None

mocket/socket.py

@@ -200,8 +200,14 @@ def proto(self) -> int:
     @property
     def io(self) -> MocketSocketIO:
         """Get or create the socket I/O object."""
-        if self._io is None:
-            self._io = MocketSocketIO((self._host, self._port))
+        if self._io is None or getattr(self._io, "closed", False):
+            address = self._address_key()
+            self._io = Mocket.get_io(address)
+            if self._io is not None and getattr(self._io, "closed", False):
+                self._io = None
+            if self._io is None:
+                self._io = MocketSocketIO(address)
+                Mocket.set_io(address, self._io)
         return self._io
 
     def fileno(self) -> int:
@@ -214,9 +220,76 @@ def fileno(self) -> int:
         r_fd, _ = Mocket.get_pair(address)
         if not r_fd:
             r_fd, w_fd = os.pipe()
+            os.set_blocking(r_fd, False)
+            os.set_blocking(w_fd, False)
             Mocket.set_pair(address, (r_fd, w_fd))
+        if self._io is not None and self._buffered_bytes():
+            if Mocket.pipe_uses_data(address):
+                self._mirror_buffer_to_pipe()
+            else:
+                self._sync_readable_pipe()
         return r_fd
 
+    def _address_key(self) -> Address:
+        """Return the current socket address tuple."""
+        return self._host, self._port
+
+    def _buffered_bytes(self) -> int:
+        """Return the number of unread bytes buffered in the socket I/O."""
+        return len(self.io.getbuffer()) - self.io.tell()
+
+    def _clear_readable_pipe(self) -> None:
+        """Drain any stale readiness bytes from the pipe for this socket."""
+        address = self._address_key()
+        r_fd, _ = Mocket.get_pair(address)
+        if not r_fd:
+            return
+
+        while True:
+            try:
+                if not os.read(r_fd, self._buflen):
+                    break
+            except BlockingIOError:
+                break
+
+        Mocket.set_pending_readables(address, 0)
+
+    def _mirror_buffer_to_pipe(self) -> None:
+        """Mirror unread response bytes into the pipe for small payloads."""
+        address = self._address_key()
+        _, w_fd = Mocket.get_pair(address)
+        if not w_fd:
+            return
+
+        unread = self.io.getbuffer()[self.io.tell() :]
+        if unread:
+            os.write(w_fd, unread)
+
+    def _sync_readable_pipe(self) -> None:
+        """Keep the readiness pipe in sync with the unread buffer size.
+
+        The pipe is used only to wake selector-based async clients. Response bytes
+        remain in the in-memory buffer so large payloads do not block on OS pipe
+        capacity.
+        """
+        address = self._address_key()
+        _, w_fd = Mocket.get_pair(address)
+        if not w_fd:
+            return
+
+        pending = Mocket.get_pending_readables(address)
+        desired = min(self._buffered_bytes(), self._buflen)
+        if desired <= pending:
+            return
+
+        try:
+            written = os.write(w_fd, b"\0" * (desired - pending))
+        except BlockingIOError:
+            written = 0
+
+        if written:
+            Mocket.set_pending_readables(address, pending + written)
+
     def gettimeout(self) -> float | None:
         """Get the socket timeout.
 
@@ -375,10 +448,17 @@ def sendall(
             response = self.true_sendall(data, *args, **kwargs)
 
         if response is not None:
+            address = self._address_key()
             self.io.seek(0)
+            self._clear_readable_pipe()
             self.io.write(response)
             self.io.truncate()
             self.io.seek(0)
+            Mocket.set_pipe_uses_data(address, len(response) <= self._buflen)
+            if Mocket.pipe_uses_data(address):
+                self._mirror_buffer_to_pipe()
+            else:
+                self._sync_readable_pipe()
 
     def sendmsg(
         self,
@@ -537,11 +617,25 @@ def recv(self, buffersize: int, flags: int | None = None) -> bytes:
         Raises:
             BlockingIOError: If socket is non-blocking and no data available
         """
-        r_fd, _ = Mocket.get_pair((self._host, self._port))
-        if r_fd:
-            return os.read(r_fd, buffersize)
+        address = self._address_key()
+        r_fd, _ = Mocket.get_pair(address)
+        if r_fd and Mocket.pipe_uses_data(address):
+            try:
+                return os.read(r_fd, buffersize)
+            except BlockingIOError:
+                pass
+
+        pending = Mocket.get_pending_readables(address)
+        if r_fd and self._buffered_bytes() and pending == 0:
+            self._sync_readable_pipe()
+            pending = Mocket.get_pending_readables(address)
+
         data = self.io.read(buffersize)
         if data:
+            if r_fd and pending:
+                drained = os.read(r_fd, min(len(data), pending))
+                Mocket.set_pending_readables(address, pending - len(drained))
+                self._sync_readable_pipe()
             return data
         # used by Redis mock
         exc = BlockingIOError()