Avoid blocking reads in the sync implementation.

docs/project/changelog.rst

@@ -39,6 +39,9 @@ Improvements
 
 .. _uvloop: https://uvloop.readthedocs.io/
 
+* Used non-blocking reads in the :mod:`threading` implementation. Depending on
+  platforms, this can make connection termination more reliable.
+
 Bug fixes
 .........
 

src/websockets/sync/connection.py

@@ -4,6 +4,7 @@
 import contextlib
 import logging
 import random
+import selectors
 import socket
 import struct
 import threading
@@ -23,7 +24,7 @@
 )
 from ..frames import DATA_OPCODES, PONG, CloseCode, Frame
 from ..http11 import Request, Response
-from ..protocol import CLOSED, CONNECTING, OPEN, Event, Protocol, State
+from ..protocol import CLOSED, OPEN, Event, Protocol, State
 from ..typing import BytesLike, Data, DataLike, LoggerLike, Subprotocol
 from .messages import Assembler
 from .utils import Deadline
@@ -129,6 +130,9 @@ def __init__(
         # ConnectionClosed in order to show why the TCP connection dropped.
         self.recv_exc: BaseException | None = None
 
+        # Notification to wake up recv_events() when the connection is closed.
+        self.shutdown_watcher, self.shutdown_notifier = socket.socketpair()
+
         # Receiving events from the socket. This thread is marked as daemon to
         # allow creating a connection in a non-daemon thread and using it in a
         # daemon thread. This mustn't prevent the interpreter from exiting.
@@ -847,34 +851,41 @@ def recv_events(self) -> None:
         ``recv_events()`` exits immediately when ``self.socket`` is closed.
 
         """
-        # When the opening handshake fails, we cannot trust rules in RFC 6455
-        # for closing TCP connections will be followed. The HTTP server could
-        # keep the connection alive after sending the response. We attempt to
-        # close the connection immediately. Unfortunately, this is unreliable
-        # on macOS; recv() may block until close_timeout elapses:
-        # https://github.com/python-websockets/websockets/issues/1596
-        # https://github.com/python/cpython/issues/154224
-        # In that case, break out of the loop to prevent recv() from blocking
-        # until close_timeout elapses.
-        close_expected_while_connecting = False
-
+        poller = selectors.DefaultSelector()
         try:
-            while not close_expected_while_connecting:
+            poller.register(self.shutdown_watcher, selectors.EVENT_READ)
+            try:
+                poller.register(self.socket, selectors.EVENT_READ, data="socket")
+            except (OSError, ValueError):
+                # close_socket() ran before poller.register().
+                # * OSError: [Errno 9] Bad file descriptor
+                #   (only seen on free-threaded Python for poller.register())
+                # * ValueError: Invalid file descriptor: -1
+                return
+
+            while True:
                 try:
                     # If the assembler buffer is full, block until it drains.
                     with self.recv_flow_control:
                         pass
-                    if self.close_deadline is not None:
-                        self.socket.settimeout(self.close_deadline.timeout())
-                    data = self.socket.recv(self.recv_bufsize)
+                    # Wait for the socket to be ready for reading.
+                    for key, _ in poller.select():
+                        if key.data == "socket":
+                            data = self.socket.recv(self.recv_bufsize)
+                            break
+                    else:
+                        # If the socket wasn't ready for reading, then it's
+                        # shutdown_watcher that was notified, meaning we're
+                        # closing the socket.
+                        data = b""
                 except Exception as exc:
                     if self.debug:
                         self.logger.debug(
                             "! error while receiving data",
                             exc_info=True,
                         )
                     # When the closing handshake is initiated by our side,
-                    # recv() may block until send_context() closes the socket.
+                    # select() may block until send_context() closes the socket.
                     # In that case, send_context() already set recv_exc.
                     # Calling set_recv_exc() avoids overwriting it.
                     with self.protocol_mutex:
@@ -903,7 +914,7 @@ def recv_events(self) -> None:
                             )
                         # Similarly to the above, avoid overriding an exception
                         # set by send_context(), in case of a race condition
-                        # i.e. send_context() closes the socket after recv()
+                        # i.e. send_context() closes the socket after select()
                         # returns above but before send_data() calls send().
                         self.set_recv_exc(exc)
                         break
@@ -912,8 +923,6 @@ def recv_events(self) -> None:
                     if self.protocol.close_expected():
                         if self.close_deadline is None:
                             self.close_deadline = Deadline(self.close_timeout)
-                        if self.protocol.state is CONNECTING:
-                            close_expected_while_connecting = True
 
                 # Unlock conn_mutex before processing events. Else, the
                 # application can't send messages in response to events.
@@ -926,8 +935,8 @@ def recv_events(self) -> None:
                     # This isn't expected to raise an exception.
                     self.process_event(event)
 
-            # Breaking out of the while not close_expected_while_connecting: ...
-            # loop means that we believe that the socket doesn't work anymore.
+            # Breaking out of the while True: ... loop means that we believe
+            # that the socket doesn't work anymore.
 
             with self.protocol_mutex:
                 # Feed the end of the data stream to the protocol.
@@ -957,6 +966,8 @@ def recv_events(self) -> None:
         finally:
             # This isn't expected to raise an exception.
             self.close_socket()
+            poller.close()
+            self.shutdown_watcher.close()
 
     @contextlib.contextmanager
     def send_context(
@@ -1103,14 +1114,16 @@ def close_socket(self) -> None:
         Shutdown and close socket. Close message assembler.
 
         Calling close_socket() guarantees that recv_events() terminates. Indeed,
-        recv_events() may block only on socket.recv() or on recv_messages.put().
+        recv_events() may block only on poller.select() or recv_messages.put().
 
         """
-        # shutdown() is required to interrupt recv() on Linux.
         try:
-            self.socket.shutdown(socket.SHUT_RDWR)
-        except OSError:  # socket already closed
+            self.shutdown_notifier.send(b"x")
+        except OSError:  # close_socket() was already called
             pass
+        finally:
+            self.shutdown_notifier.close()
+
         try:
             self.socket.close()
         except OSError:  # socket already closed  # pragma: no cover

src/websockets/trio/connection.py

@@ -886,6 +886,8 @@ async def recv_events(self) -> None:
                     if self.recv_flow_control is not None:
                         await self.recv_flow_control.wait()
                     data = await self.stream.receive_some()
+                except trio.ClosedResourceError:
+                    break
                 except Exception as exc:
                     if self.debug:
                         self.logger.debug(

tests/asyncio/test_client.py

@@ -516,12 +516,41 @@ def http_response(connection, request):
         self.assertEqual(raised.exception.response.status_code, 200)
         self.assertEqual(raised.exception.response.body.decode(), "👌")
 
+    async def test_http_response_with_keep_alive(self):
+        """Client reads HTTP response and server keeps connection alive."""
+
+        async def keep_alive(reader, writer):
+            # Wait for the client to send the handshake request.
+            await reader.read(4096)
+            writer.write(
+                b"HTTP/1.1 410 Gone\r\n"
+                b"Connection: keep-alive\r\n"  # default in HTTP/1.1
+                b"Content-Length: 0\r\n"  # avoids "read body until EOF"
+                b"\r\n"
+            )
+            # Wait for the client to close the connection.
+            await reader.read(4096)
+            writer.close()
+
+        server = await asyncio.start_server(keep_alive, "localhost", 0)
+        host, port = get_host_port(server)
+        async with server:
+            with self.assertRaises(InvalidStatus) as raised:
+                async with asyncio.timeout(10 * MS):
+                    async with connect(f"ws://{host}:{port}"):
+                        self.fail("did not raise")
+
+        self.assertEqual(
+            str(raised.exception),
+            "server rejected WebSocket connection: HTTP 410",
+        )
+
     async def test_junk_handshake_response(self):
         """Client closes the connection when receiving non-HTTP response from server."""
 
         async def junk(reader, writer):
             # Wait for the client to send the handshake request.
-            await asyncio.sleep(MS)
+            await reader.read(4096)
             writer.write(b"220 smtp.invalid ESMTP Postfix\r\n")
             # Wait for the client to close the connection.
             await reader.read(4096)

tests/asyncio/test_connection.py

@@ -1051,14 +1051,14 @@ async def test_pong_unsupported_type(self):
     async def test_keepalive(self, getrandbits):
         """keepalive sends pings at ping_interval and measures latency."""
         getrandbits.side_effect = itertools.count(1918987876)
-        self.connection.ping_interval = 3 * MS
+        self.connection.ping_interval = 4 * MS
         self.connection.start_keepalive()
         self.assertIsNotNone(self.connection.keepalive_task)
         self.assertEqual(self.connection.latency, 0)
-        # 3 ms: keepalive() sends a ping frame.
-        # 3.x ms: a pong frame is received.
-        await asyncio.sleep(5 * MS)
-        # 5 ms: check that the ping frame was sent.
+        # 4 ms: keepalive() sends a ping frame.
+        # 4.x ms: a pong frame is received.
+        await asyncio.sleep(6 * MS)
+        # 6 ms: check that the ping frame was sent.
         await self.assertFrameSent(Frame(PING, b"rand"))
         self.assertGreater(self.connection.latency, 0)
         self.assertLess(self.connection.latency, MS)

tests/sync/test_client.py

@@ -530,8 +530,8 @@ def test_http_response_with_keep_alive(self):
 
         class KeepAliveHandler(socketserver.BaseRequestHandler):
             def handle(self):
-                time.sleep(MS)  # wait for the client to send the handshake request
-                self.request.recv(4096)  # assume we read the full handshake request
+                # Wait for the client to send the handshake request.
+                self.request.recv(4096)
                 self.request.send(
                     b"HTTP/1.1 410 Gone\r\n"
                     b"Connection: keep-alive\r\n"  # default in HTTP/1.1
@@ -543,7 +543,7 @@ def handle(self):
                 self.request.close()
 
         server = socketserver.TCPServer(("localhost", 0), KeepAliveHandler)
-        host, port = server.server_address
+        host, port = get_host_port(server)
         with server:
             thread = threading.Thread(target=server.serve_forever, args=(MS,))
             thread.start()
@@ -561,7 +561,7 @@ def handle(self):
             str(raised.exception),
             "server rejected WebSocket connection: HTTP 410",
         )
-        self.assertLess(t1 - t0, 5 * MS)
+        self.assertLess(t1 - t0, 10 * MS)
 
     def test_junk_handshake_response(self):
         """Client closes the connection when receiving non-HTTP response from server."""

tests/sync/test_connection.py

@@ -16,7 +16,7 @@
     ConnectionClosedOK,
 )
 from websockets.frames import BINARY, CLOSE, PING, PONG, TEXT, CloseCode, Frame
-from websockets.protocol import CLIENT, CLOSED, CONNECTING, OPEN, SERVER, Protocol
+from websockets.protocol import CLIENT, CLOSED, OPEN, SERVER, Protocol
 from websockets.sync.connection import *
 from websockets.sync.connection import broadcast
 
@@ -787,14 +787,14 @@ def test_pong_unsupported_type(self):
     def test_keepalive(self, getrandbits):
         """keepalive sends pings at ping_interval and measures latency."""
         getrandbits.side_effect = itertools.count(1918987876)
-        self.connection.ping_interval = 3 * MS
+        self.connection.ping_interval = 4 * MS
         self.connection.start_keepalive()
         self.assertIsNotNone(self.connection.keepalive_thread)
         self.assertEqual(self.connection.latency, 0)
-        # 3 ms: keepalive() sends a ping frame.
-        # 3.x ms: a pong frame is received.
-        time.sleep(5 * MS)
-        # 5 ms: check that the ping frame was sent.
+        # 4 ms: keepalive() sends a ping frame.
+        # 4.x ms: a pong frame is received.
+        time.sleep(6 * MS)
+        # 6 ms: check that the ping frame was sent.
         self.assertFrameSent(Frame(PING, b"rand"))
         self.assertGreater(self.connection.latency, 0)
         self.assertLess(self.connection.latency, MS)
@@ -1033,6 +1033,15 @@ def test_writing_in_send_context_fails(self):
 
     # Test safety nets — catching all exceptions in case of bugs.
 
+    @patch("socket.socket.recv", side_effect=AssertionError)
+    def test_unexpected_failure_in_recv(self, recv):
+        """Unexpected internal error in recv() is correctly reported."""
+        self.remote_connection.send("😀")
+        # Reading the message will trigger the injected fault.
+        with self.assertRaises(ConnectionClosedError) as raised:
+            self.connection.recv()
+        self.assertIsInstance(raised.exception.__cause__, AssertionError)
+
     @patch("websockets.protocol.Protocol.events_received", side_effect=AssertionError)
     def test_unexpected_failure_in_recv_events(self, events_received):
         """Unexpected internal error in recv_events() is correctly reported."""
@@ -1050,30 +1059,6 @@ def test_unexpected_failure_in_send_context(self, send_text):
             self.connection.send("😀")
         self.assertIsInstance(raised.exception.__cause__, AssertionError)
 
-    # Test recv_events().
-
-    def test_close_expected_during_connecting(self):
-        """recv_events() terminates without waiting for close_timeout."""
-        # See https://github.com/python-websockets/websockets/issues/1596.
-        socket_, remote_socket = socket.socketpair()
-        protocol = Protocol(self.LOCAL, state=CONNECTING)
-        # Simulate ClientProtocol / ServerProtocol deciding that the TCP
-        # connection must close because the opening handshake failed.
-        protocol.send_eof()
-        connection = Connection(socket_, protocol, close_timeout=5 * MS)
-        try:
-            t0 = time.time()
-            remote_socket.sendall(b"\x00")
-            connection.recv_events_thread.join(3 * MS)
-            t1 = time.time()
-
-            self.assertFalse(connection.recv_events_thread.is_alive())
-            self.assertLess(t1 - t0, 3 * MS)
-        finally:
-            remote_socket.close()
-            connection.close_socket()
-            connection.recv_events_thread.join()
-
     # Test garbage collection.
 
     @unittest.skipUnless(

tests/trio/test_client.py

@@ -508,12 +508,48 @@ def http_response(connection, request):
         self.assertEqual(raised.exception.response.status_code, 200)
         self.assertEqual(raised.exception.response.body.decode(), "👌")
 
+    async def test_http_response_with_keep_alive(self):
+        """Client reads HTTP response and server keeps connection alive."""
+
+        async def keep_alive(stream):
+            # Wait for the client to send the handshake request.
+            await stream.receive_some()
+            await stream.send_all(
+                b"HTTP/1.1 410 Gone\r\n"
+                b"Connection: keep-alive\r\n"  # default in HTTP/1.1
+                b"Content-Length: 0\r\n"  # avoids "read body until EOF"
+                b"\r\n"
+            )
+            # Wait for the client to close the connection.
+            await stream.receive_some()
+            await stream.aclose()
+
+        async with trio.open_nursery() as nursery:
+            listeners = await nursery.start(trio.serve_tcp, keep_alive, 0)
+            host, port = next(
+                listener
+                for listener in listeners
+                if listener.socket.family == socket.AF_INET
+            ).socket.getsockname()
+            try:
+                with self.assertRaises(InvalidStatus) as raised:
+                    with trio.fail_after(10 * MS):
+                        async with connect(f"ws://{host}:{port}"):
+                            self.fail("did not raise")
+            finally:
+                nursery.cancel_scope.cancel()
+
+        self.assertEqual(
+            str(raised.exception),
+            "server rejected WebSocket connection: HTTP 410",
+        )
+
     async def test_junk_handshake_response(self):
         """Client closes the connection when receiving non-HTTP response from server."""
 
         async def junk(stream):
             # Wait for the client to send the handshake request.
-            await trio.testing.wait_all_tasks_blocked()
+            await stream.receive_some()
             await stream.send_all(b"220 smtp.invalid ESMTP Postfix\r\n")
             # Wait for the client to close the connection.
             await stream.receive_some()