Use select() to retry blocked socket writes instead of spinning or failing

When raw.write() returns None on a blocked non-blocking socket, use
select() to wait up to SOCK_WRITE_TIMEOUT seconds for the socket to
become writable before retrying. This avoids both the original silent
buffer truncation bug and the infinite spin loop that the previous fix
could produce. If the socket stays blocked past the timeout, raise
BlockingIOError with the write buffer preserved.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: 2 people

cheroot/makefile.py

@@ -2,12 +2,16 @@
 
 # prefer slower Python-based io module
 import _pyio as io
+import select
 import socket
 
 
 # Write only 16K at a time to sockets
 SOCK_WRITE_BLOCKSIZE = 16384
 
+# Seconds to wait for a blocked socket to become writable
+SOCK_WRITE_TIMEOUT = 10
+
 
 class BufferedWriter(io.BufferedWriter):
     """Faux file object attached to a socket object."""
@@ -35,8 +39,20 @@ def _flush_unlocked(self):
                 )
             except io.BlockingIOError as e:
                 n = e.characters_written
-            if n:
-                del self._write_buf[:n]
+            if n is None:
+                _, writable, _ = select.select(
+                    [],
+                    [self.raw],
+                    [],
+                    SOCK_WRITE_TIMEOUT,
+                )
+                if not writable:
+                    raise io.BlockingIOError(
+                        0,
+                        'raw stream blocked; no bytes written',
+                    )
+                continue
+            del self._write_buf[:n]
 
 
 class StreamReader(io.BufferedReader):

cheroot/makefile.pyi

@@ -1,6 +1,7 @@
 import io
 
 SOCK_WRITE_BLOCKSIZE: int
+SOCK_WRITE_TIMEOUT: int
 
 class BufferedWriter(io.BufferedWriter):
     def write(self, b): ...

cheroot/test/test_makefile.py

@@ -54,43 +54,97 @@ def test_bytes_written():
 
 
 class _RawWriteBlockOnce:
-    """Mock raw.write() that returns None on the first call, then writes normally."""
+    """Mock raw.write() returning None once, then writing normally."""
 
     def __init__(self):
         """Initialize _RawWriteBlockOnce."""
         self.call_count = 0
         self.written = bytearray()
 
     def __call__(self, chunk):
-        """Return None on first call to simulate a blocked socket write."""
+        """Return None on first call to simulate a blocked write."""
         self.call_count += 1
         if self.call_count == 1:
-            return (
-                None  # simulates socket returning None on first blocked write
-            )
+            return None
         self.written.extend(chunk)
         return len(chunk)
 
+    def fileno(self):
+        """Return a fake fd for select()."""
+        return -1
 
-def test_flush_when_raw_write_returns_none():
-    """_flush_unlocked() must not treat None from raw.write() as a byte count.
 
-    io.RawIOBase.write() returns None when a non-blocking socket cannot accept
-    data. del self._write_buf[:None] is equivalent to del self._write_buf[:]
-    which silently clears the entire buffer, truncating the response without
-    raising an exception.
+class _RawWriteBlockAlways:
+    """Mock raw.write() that always returns None."""
+
+    def __init__(self):
+        """Initialize _RawWriteBlockAlways."""
+        self.call_count = 0
+
+    def __call__(self, chunk):
+        """Return None to simulate a permanently blocked socket."""
+        self.call_count += 1
+
+    def fileno(self):
+        """Return a fake fd for select()."""
+        return -1
+
+
+def test_flush_recovers_from_temporary_block(monkeypatch):
+    """_flush_unlocked() retries after select when raw.write() returns None.
+
+    A temporarily blocked socket should recover once select() reports
+    the socket is writable again, delivering all buffered data.
     """
-    data = b'x' * (makefile.SOCK_WRITE_BLOCKSIZE * 2)  # stress the write loop
+    data = b'x' * (makefile.SOCK_WRITE_BLOCKSIZE * 2)
 
     sock = MockSocket()
     wfile = makefile.MakeFile(sock, 'w')
     wfile._write_buf.extend(data)
 
     mock = _RawWriteBlockOnce()
     wfile.raw.write = mock
+
+    # select() reports writable immediately
+    monkeypatch.setattr(
+        'cheroot.makefile.select.select',
+        lambda _rlist, wlist, _xlist, _timeout: ([], wlist, []),
+    )
     wfile._flush_unlocked()
 
     assert bytes(mock.written) == data, (
-        f'Expected {len(data)} bytes but only {len(mock.written)} reached raw.write(): '
-        'buffer was silently discarded when raw.write() returned None'
+        'all buffered data should be written after select retry'
+    )
+
+
+def test_flush_raises_on_sustained_block(monkeypatch):
+    """_flush_unlocked() raises BlockingIOError after select timeout.
+
+    If the socket stays blocked past SOCK_WRITE_TIMEOUT, the write
+    buffer must be preserved and BlockingIOError raised.
+    """
+    import io
+
+    import pytest
+
+    data = b'x' * makefile.SOCK_WRITE_BLOCKSIZE
+
+    sock = MockSocket()
+    wfile = makefile.MakeFile(sock, 'w')
+    wfile._write_buf.extend(data)
+
+    mock = _RawWriteBlockAlways()
+    wfile.raw.write = mock
+
+    # select() reports not writable (timeout)
+    monkeypatch.setattr(
+        'cheroot.makefile.select.select',
+        lambda _rlist, _wlist, _xlist, _timeout: ([], [], []),
+    )
+
+    with pytest.raises(io.BlockingIOError):
+        wfile._flush_unlocked()
+
+    assert len(wfile._write_buf) == len(data), (
+        'write buffer must be preserved when socket stays blocked'
     )