diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index f297359a7..b3acd7ece 100644 --- a/docs/project/changelog.rst +++ b/docs/project/changelog.rst @@ -78,6 +78,13 @@ New features * Added :func:`~sync.server.broadcast` to the :mod:`threading` implementation. +* Made the set of active connections available in the :attr:`Server.connections + ` property in the :mod:`threading` + implementation. + +* Closed connections when shutting down the server in the :mod:`threading` + implementation. See :meth:`~sync.server.Server.shutdown` for details. + * Added the ``--insecure`` option to the ``websockets`` CLI to disable TLS certificate validation. diff --git a/docs/reference/sync/server.rst b/docs/reference/sync/server.rst index 6df014f0d..38458255b 100644 --- a/docs/reference/sync/server.rst +++ b/docs/reference/sync/server.rst @@ -28,6 +28,8 @@ Running a server .. autoclass:: Server + .. autoattribute:: connections + .. automethod:: serve_forever .. automethod:: shutdown diff --git a/src/websockets/asyncio/server.py b/src/websockets/asyncio/server.py index bfe2af7f1..3eb6a7d4c 100644 --- a/src/websockets/asyncio/server.py +++ b/src/websockets/asyncio/server.py @@ -292,8 +292,9 @@ def __init__( logger = logging.getLogger("websockets.server") self.logger = logger - # Keep track of active connections. - self.handlers: dict[ServerConnection, asyncio.Task[None]] = {} + # Keep track of active connections and connection handler tasks. + self.all_connections: set[ServerConnection] = set() + self.handler_tasks: set[asyncio.Task[None]] = set() # Task responsible for closing the server and terminating connections. self.close_task: asyncio.Task[None] | None = None @@ -311,7 +312,11 @@ def connections(self) -> set[ServerConnection]: It can be useful in combination with :func:`~broadcast`. """ - return {connection for connection in self.handlers if connection.state is OPEN} + return { + connection + for connection in self.all_connections + if connection.state is OPEN + } def wrap(self, server: asyncio.Server) -> None: """ @@ -392,18 +397,23 @@ async def conn_handler(self, connection: ServerConnection) -> None: # Registration is tied to the lifecycle of conn_handler() because # the server waits for connection handlers to terminate, even if # all connections are already closed. - del self.handlers[connection] + self.all_connections.discard(connection) + task = asyncio.current_task() + assert task is not None # help mypy + self.handler_tasks.discard(task) def start_connection_handler(self, connection: ServerConnection) -> None: """ Register a connection with this server. """ - # The connection must be registered in self.handlers immediately. + # The connection must be registered in self.all_connections now. # If it was registered in conn_handler(), a race condition could # happen when closing the server after scheduling conn_handler() # but before it starts executing. - self.handlers[connection] = self.loop.create_task(self.conn_handler(connection)) + self.all_connections.add(connection) + handler_task = self.loop.create_task(self.conn_handler(connection)) + self.handler_tasks.add(handler_task) def close( self, @@ -458,10 +468,10 @@ async def _close( # HTTP 503 error. if close_connections: - # Close OPEN connections with code 1001 by default. + # Close OPEN connections. close_tasks = [ asyncio.create_task(connection.close(code, reason)) - for connection in self.handlers + for connection in self.all_connections if connection.protocol.state is not CONNECTING ] # asyncio.wait doesn't accept an empty first argument. @@ -473,8 +483,8 @@ async def _close( # Wait until all connection handlers terminate. # asyncio.wait doesn't accept an empty first argument. - if self.handlers: - await asyncio.wait(self.handlers.values()) + if self.handler_tasks: + await asyncio.wait(self.handler_tasks) # Tell wait_closed() to return. self.closed_waiter.set_result(None) diff --git a/src/websockets/sync/server.py b/src/websockets/sync/server.py index 8b532833d..5967c2225 100644 --- a/src/websockets/sync/server.py +++ b/src/websockets/sync/server.py @@ -1,5 +1,6 @@ from __future__ import annotations +import concurrent.futures import hmac import http import logging @@ -9,6 +10,7 @@ import ssl as ssl_module import sys import threading +import time import warnings from collections.abc import Iterable, Sequence from types import TracebackType @@ -64,6 +66,7 @@ class ServerConnection(Connection): Args: socket: Socket connected to a WebSocket client. protocol: Sans-I/O connection. + server: Server that manages this connection. """ @@ -71,6 +74,7 @@ def __init__( self, sock: socket.socket, protocol: ServerProtocol, + server: Server, *, ping_interval: float | None = 20, ping_timeout: float | None = 20, @@ -87,6 +91,7 @@ def __init__( close_timeout=close_timeout, max_queue=max_queue, ) + self.server = server self.username: str # see basic_auth() self.handler: Callable[[ServerConnection], None] # see route() self.handler_kwargs: Mapping[str, Any] # see route() @@ -157,7 +162,13 @@ def handshake( ) if response is None: - self.response = self.protocol.accept(self.request) + if self.server.socket_closed.is_set(): + self.response = self.protocol.reject( + http.HTTPStatus.SERVICE_UNAVAILABLE, + "Server is shutting down.\n", + ) + else: + self.response = self.protocol.accept(self.request) else: self.response = response @@ -231,6 +242,9 @@ class Server: :meth:`~socketserver.BaseServer.shutdown` methods, as well as the context manager protocol. + It keeps track of WebSocket connections in order to close them properly + when shutting down. + Args: socket: Server socket listening for new connections. handler: Handler for one connection. Receives the socket and address @@ -241,6 +255,8 @@ class Server: """ + SHUTDOWN_POLLING_INTERVAL = 0.1 # seconds + def __init__( self, sock: socket.socket, @@ -252,11 +268,40 @@ def __init__( if logger is None: logger = logging.getLogger("websockets.server") self.logger = logger + + # Synchronize access to closing, all_connections, and handler_threads. + self.lock = threading.Lock() + + # Keep track of active connections and connection handler threads. + self.all_connections: set[ServerConnection] = set() + self.handler_threads: set[threading.Thread] = set() + # On Windows, closing the socket wakes up the poller in serve_forever(), # making the notification mechanism unnecessary. if sys.platform != "win32": self.shutdown_watcher, self.shutdown_notifier = socket.socketpair() + # Set when serve_forever() no longer accepts new connections and starts + # threads to handle them. + self.socket_closed = threading.Event() + + @property + def connections(self) -> set[ServerConnection]: + """ + Set of active connections. + + This property contains all connections that completed the opening + handshake successfully and didn't start the closing handshake yet. + It can be useful in combination with :func:`~broadcast`. + + """ + with self.lock: + return { + connection + for connection in self.all_connections + if connection.protocol.state is OPEN + } + def serve_forever(self) -> None: """ See :meth:`socketserver.BaseServer.serve_forever`. @@ -293,20 +338,50 @@ def serve_forever(self) -> None: sock, addr = self.socket.accept() except OSError: break - # Since we don't track connections, we cannot wait for handlers - # to terminate, so we cannot use daemon threads. If we did, all - # connections would be closed brutally when closing the server. + # shutdown() can let existing connections terminate on their own + # or close them. Either way, it waits for connection handlers to + # terminate, so there's no point using daemon threads. thread = threading.Thread(target=self.handler, args=(sock, addr)) + # The thread must be registered in self.handler_threads now, + # before it's started. Otherwise, a race condition could + # happen when closing the server after starting the thread + # but before it starts executing. + with self.lock: + self.handler_threads.add(thread) thread.start() finally: + self.socket_closed.set() if sys.platform != "win32": self.shutdown_watcher.close() - def shutdown(self) -> None: + def shutdown( + self, + close_connections: bool = True, + code: CloseCode | int = CloseCode.GOING_AWAY, + reason: str = "", + ) -> None: """ - See :meth:`socketserver.BaseServer.shutdown`. + Close the server. + + * Close the listening socket to stop accepting new connections. + * When ``close_connections`` is :obj:`True`, which is the default, close + existing connections. Specifically: + + * Reject opening WebSocket connections with an HTTP 503 (service + unavailable) error. This happens when the server accepted the TCP + connection but didn't complete the opening handshake before closing. + * Close open WebSocket connections with code 1001 (going away). + ``code`` and ``reason`` can be customized, for example to use code + 1012 (service restart). + + * Wait until all connection handlers terminate. + + :meth:`shutdown` is idempotent. """ + self.logger.info("server closing") + + # Stop accepting new connections. self.socket.close() if sys.platform != "win32": try: @@ -316,6 +391,54 @@ def shutdown(self) -> None: finally: self.shutdown_notifier.close() + # Wait until serve_forever() no longer accepts new connections nor + # starts threads to handle them, meaning that self.handler_threads + # won't get new entries. + self.socket_closed.wait() + + if close_connections: + # At this point, all threads are started, but some may still be in + # the opening handshake. Close open connections until no thread is + # executing anymore. Some threads may be cleaning up; in that case + # they're expected to terminate quickly, so waiting is fine. + while True: + with self.lock: + # Inline self.connections because it acquires self.lock, + # which isn't reentrant. + connections = [ + connection + for connection in self.all_connections + if connection.protocol.state is OPEN + ] + threads = list(self.handler_threads) + # No threads are executing anymore. Server is fully closed. + if not threads: + break + # Some threads are still executing, but no connections are OPEN. + # Wait for connections to complete the opening handshake, or for + # handler threads to terminate. + if not connections: + time.sleep(self.SHUTDOWN_POLLING_INTERVAL) + continue + # Close OPEN connections and wait until they're closed. + with concurrent.futures.ThreadPoolExecutor() as executor: + for connection in connections: + executor.submit(connection.close, code, reason) + + else: + # At this point, all threads are started. + with self.lock: + threads = list(self.handler_threads) + # Wait until all connection handlers terminate. + for thread in threads: + # This raises RuntimeError if shutdown() is called from a + # connection handler. It's documented to return after all + # connection handlers terminate, which is impossible when + # it's called from a connection handler. + thread.join() + + self.logger.info("server closed") + def fileno(self) -> int: """ See :meth:`socketserver.BaseServer.fileno`. @@ -600,14 +723,22 @@ def protocol_select_subprotocol( connection = create_connection( sock, protocol, + server, ping_interval=ping_interval, ping_timeout=ping_timeout, close_timeout=close_timeout, max_queue=max_queue, ) except Exception: - sock.close() - return + try: + sock.close() + return + finally: + with server.lock: + server.handler_threads.discard(threading.current_thread()) + + with server.lock: + server.all_connections.add(connection) try: try: @@ -645,9 +776,19 @@ def protocol_select_subprotocol( # Don't leak sockets on unexpected errors. sock.close() + finally: + # Registration is tied to the lifecycle of conn_handler() because + # the server waits for connection handlers to terminate, even if + # all connections are already closed. + with server.lock: + server.all_connections.discard(connection) + server.handler_threads.discard(threading.current_thread()) + # Initialize server - return Server(sock, conn_handler, logger) + # The `server` variable is captured by the closure of conn_handler(). + server = Server(sock, conn_handler, logger) + return server def unix_serve( diff --git a/tests/asyncio/test_client.py b/tests/asyncio/test_client.py index 9fb1bff2f..dde37fb3a 100644 --- a/tests/asyncio/test_client.py +++ b/tests/asyncio/test_client.py @@ -96,19 +96,19 @@ async def test_additional_headers(self): ) as client: self.assertEqual(client.request.headers["Authorization"], "Bearer ...") - async def test_override_user_agent(self): + async def test_override_user_agent_header(self): """Client can override User-Agent header with user_agent_header.""" async with serve(*args) as server: async with connect(get_uri(server), user_agent_header="Smith") as client: self.assertEqual(client.request.headers["User-Agent"], "Smith") - async def test_remove_user_agent(self): + async def test_remove_user_agent_header(self): """Client can remove User-Agent header with user_agent_header.""" async with serve(*args) as server: async with connect(get_uri(server), user_agent_header=None) as client: self.assertNotIn("User-Agent", client.request.headers) - async def test_legacy_user_agent(self): + async def test_legacy_user_agent_header(self): """Client can override User-Agent header with additional_headers.""" async with serve(*args) as server: async with connect( diff --git a/tests/asyncio/test_server.py b/tests/asyncio/test_server.py index 475920d49..1722dcf16 100644 --- a/tests/asyncio/test_server.py +++ b/tests/asyncio/test_server.py @@ -359,13 +359,13 @@ async def process_response(ws, request, response): ["BOOM"], ) - async def test_override_server(self): + async def test_override_server_header(self): """Server can override Server header with server_header.""" async with serve(*args, server_header="Neo") as server: async with connect(get_uri(server)) as client: await self.assertEval(client, "ws.response.headers['Server']", "Neo") - async def test_remove_server(self): + async def test_remove_server_header(self): """Server can remove Server header with server_header.""" async with serve(*args, server_header=None) as server: async with connect(get_uri(server)) as client: @@ -416,10 +416,16 @@ async def test_connections(self): """Server provides a connections property.""" async with serve(*args) as server: self.assertEqual(server.connections, set()) - async with connect(get_uri(server)) as client: - self.assertEqual(len(server.connections), 1) - ws_id = str(next(iter(server.connections)).id) - await self.assertEval(client, "ws.id", ws_id) + async with ( + connect(get_uri(server)) as client1, + connect(get_uri(server)) as client2, + ): + await client1.send("ws.id") + await client2.send("ws.id") + self.assertEqual( + {str(connection.id) for connection in server.connections}, + {await client1.recv(), await client2.recv()}, + ) self.assertEqual(server.connections, set()) async def test_handshake_fails(self): @@ -481,8 +487,8 @@ async def test_junk_handshake(self): "invalid HTTP request line: HELO relay.invalid", ) - async def test_close_server_rejects_connecting_connections(self): - """Server rejects connecting connections with HTTP 503 when closing.""" + async def test_close_rejects_connecting_connections(self): + """Server rejects connecting connections with HTTP 503.""" async def process_request(ws, _request): while ws.server.is_serving(): @@ -498,8 +504,8 @@ async def process_request(ws, _request): "server rejected WebSocket connection: HTTP 503", ) - async def test_close_server_closes_open_connections(self): - """Server closes open connections with close code 1001 when closing.""" + async def test_close_closes_open_connections(self): + """Server closes open connections with close code 1001.""" async with serve(*args) as server: async with connect(get_uri(server)) as client: server.close() @@ -510,8 +516,8 @@ async def test_close_server_closes_open_connections(self): "received 1001 (going away); then sent 1001 (going away)", ) - async def test_close_server_closes_open_connections_with_code_and_reason(self): - """Server closes open connections with custom code and reason when closing.""" + async def test_close_closes_open_connections_with_code_and_reason(self): + """Server closes open connections with custom code and reason.""" async with serve(*args) as server: async with connect(get_uri(server)) as client: server.close(code=1012, reason="restarting") @@ -523,8 +529,8 @@ async def test_close_server_closes_open_connections_with_code_and_reason(self): "then sent 1012 (service restart) restarting", ) - async def test_close_server_keeps_connections_open(self): - """Server waits for client to close open connections when closing.""" + async def test_close_keeps_connections_open(self): + """Server waits for the client to close open connections.""" async with serve(*args) as server: async with connect(get_uri(server)) as client: server.close(close_connections=False) @@ -543,7 +549,7 @@ async def test_close_server_keeps_connections_open(self): async with asyncio.timeout(MS): await server.wait_closed() - async def test_close_server_keeps_handlers_running(self): + async def test_close_keeps_handlers_running(self): """Server waits for connection handlers to terminate.""" async with serve(*args) as server: async with connect(get_uri(server) + "/delay") as client: @@ -557,10 +563,16 @@ async def test_close_server_keeps_handlers_running(self): async with asyncio.timeout(2 * MS): await server.wait_closed() - # Set a large timeout here, else the test becomes flaky. + # Set a longer timeout here, else the test becomes flaky. async with asyncio.timeout(5 * MS): await server.wait_closed() + async def test_sockets(self): + """Server provides a sockets property.""" + async with serve(*args) as server: + sock = server.sockets[0] + self.assertIsInstance(sock.fileno(), int) + SSL_OBJECT = "ws.transport.get_extra_info('ssl_object')" diff --git a/tests/sync/server.py b/tests/sync/server.py index cadaa267e..d242e8bfe 100644 --- a/tests/sync/server.py +++ b/tests/sync/server.py @@ -1,6 +1,7 @@ import contextlib import ssl import threading +import time import urllib.parse from websockets.sync.router import * @@ -26,6 +27,10 @@ def handler(ws): raise RuntimeError elif path == "/no-op": pass + elif path == "/delay": + delay = float(ws.recv()) + ws.close() + time.sleep(delay) else: raise AssertionError(f"unexpected path: {path}") @@ -47,30 +52,12 @@ def run_server_or_router( with serve_or_route(handler_or_url_map, host, port, **kwargs) as server: thread = threading.Thread(target=server.serve_forever) thread.start() - - # HACK: since the sync server doesn't track connections (yet), we record - # a reference to the thread handling the most recent connection, then we - # can wait for that thread to terminate when exiting the context. - handler_thread = None - original_handler = server.handler - - def handler(sock, addr): - nonlocal handler_thread - handler_thread = threading.current_thread() - original_handler(sock, addr) - - server.handler = handler - try: yield server finally: server.shutdown() thread.join() - # HACK: wait for the thread handling the most recent connection. - if handler_thread is not None: - handler_thread.join() - def run_server(handler=handler, **kwargs): return run_server_or_router(serve, handler, **kwargs) diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index fb1261c84..492991383 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -71,19 +71,19 @@ def test_additional_headers(self): ) as client: self.assertEqual(client.request.headers["Authorization"], "Bearer ...") - def test_override_user_agent(self): + def test_override_user_agent_header(self): """Client can override User-Agent header with user_agent_header.""" with run_server() as server: with connect(get_uri(server), user_agent_header="Smith") as client: self.assertEqual(client.request.headers["User-Agent"], "Smith") - def test_remove_user_agent(self): + def test_remove_user_agent_header(self): """Client can remove User-Agent header with user_agent_header.""" with run_server() as server: with connect(get_uri(server), user_agent_header=None) as client: self.assertNotIn("User-Agent", client.request.headers) - def test_legacy_user_agent(self): + def test_legacy_user_agent_header(self): """Client can override User-Agent header with additional_headers.""" with run_server() as server: with connect( diff --git a/tests/sync/test_server.py b/tests/sync/test_server.py index 20e57f7b1..8eaf80f5a 100644 --- a/tests/sync/test_server.py +++ b/tests/sync/test_server.py @@ -3,8 +3,10 @@ import http import logging import socket +import threading import time import unittest +from unittest.mock import patch from websockets.exceptions import ( ConnectionClosedError, @@ -243,13 +245,13 @@ def process_response(ws, request, response): ["BOOM"], ) - def test_override_server(self): + def test_override_server_header(self): """Server can override Server header with server_header.""" with run_server(server_header="Neo") as server: with connect(get_uri(server)) as client: self.assertEval(client, "ws.response.headers['Server']", "Neo") - def test_remove_server(self): + def test_remove_server_header(self): """Server can remove Server header with server_header.""" with run_server(server_header=None) as server: with connect(get_uri(server)) as client: @@ -294,18 +296,21 @@ def create_connection(*args, **kwargs): with connect(get_uri(server)) as client: self.assertEval(client, "ws.create_connection_ran", "True") - def test_fileno(self): - """Server provides a fileno attribute.""" - with run_server() as server: - self.assertIsInstance(server.fileno(), int) - - def test_shutdown(self): - """Server provides a shutdown method.""" + def test_connections(self): + """Server provides a connections property.""" with run_server() as server: - server.shutdown() - # Check that the server socket is closed. - with self.assertRaises(OSError): - server.socket.accept() + self.assertEqual(server.connections, set()) + with ( + connect(get_uri(server)) as client1, + connect(get_uri(server)) as client2, + ): + client1.send("ws.id") + client2.send("ws.id") + self.assertEqual( + {str(connection.id) for connection in server.connections}, + {client1.recv(), client2.recv()}, + ) + self.assertEqual(server.connections, set()) def test_handshake_fails(self): """Server receives connection from client but the handshake fails.""" @@ -361,6 +366,96 @@ def test_junk_handshake(self): "invalid HTTP request line: HELO relay.invalid", ) + def test_shutdown_rejects_connecting_connections(self): + """Server rejects connecting connections with HTTP 503.""" + + def process_request(ws, _request): + ws.server.socket_closed.wait() + + def shutdown_server(server): + time.sleep(5 * MS) + server.shutdown() + + with run_server(process_request=process_request) as server: + shutdown_thread = threading.Thread(target=shutdown_server, args=(server,)) + shutdown_thread.start() + with self.assertRaises(InvalidStatus) as raised: + with connect(get_uri(server)): + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "server rejected WebSocket connection: HTTP 503", + ) + shutdown_thread.join(MS) + + def test_shutdown_closes_open_connections(self): + """Server closes open connections with code 1001.""" + with run_server() as server: + with connect(get_uri(server)) as client: + server.shutdown() + with self.assertRaises(ConnectionClosedOK) as raised: + client.recv() + self.assertEqual( + str(raised.exception), + "received 1001 (going away); then sent 1001 (going away)", + ) + + def test_shutdown_closes_open_connections_with_code_and_reason(self): + """Server closes open connections with a custom code and reason.""" + with run_server() as server: + with connect(get_uri(server)) as client: + server.shutdown(code=1012, reason="restarting") + with self.assertRaises(ConnectionClosedError) as raised: + client.recv() + self.assertEqual( + str(raised.exception), + "received 1012 (service restart) restarting; " + "then sent 1012 (service restart) restarting", + ) + + def test_shutdown_keeps_connections_open(self): + """Server waits for the client to close open connections.""" + with run_server() as server: + with connect(get_uri(server)) as client: + shutdown_thread = threading.Thread( + target=server.shutdown, + kwargs={"close_connections": False}, + ) + shutdown_thread.start() + + # The server waits for the client to close the connection. + shutdown_thread.join(MS) + self.assertTrue(shutdown_thread.is_alive()) + + # Once the client closes the connection, the server terminates. + client.close() + shutdown_thread.join(MS) + self.assertFalse(shutdown_thread.is_alive()) + + @patch("websockets.sync.server.Server.SHUTDOWN_POLLING_INTERVAL", MS) + def test_shutdown_keeps_handlers_running(self): + """Server waits for connection handlers to terminate.""" + with run_server() as server: + with connect(get_uri(server) + "/delay") as client: + # Delay termination of the connection handler. + client.send(str(3 * MS)) + + shutdown_thread = threading.Thread(target=server.shutdown) + shutdown_thread.start() + + # The server waits for the connection handler to terminate. + shutdown_thread.join(2 * MS) + self.assertTrue(shutdown_thread.is_alive()) + + # Set a longer timeout here, else the test becomes flaky. + shutdown_thread.join(5 * MS) + self.assertFalse(shutdown_thread.is_alive()) + + def test_fileno(self): + """Server provides a fileno method.""" + with run_server() as server: + self.assertIsInstance(server.fileno(), int) + class SecureServerTests(EvalShellMixin, unittest.TestCase): def test_connection(self):