Close connections when the threading server exits.

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
+  <sync.server.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.
 

docs/reference/sync/server.rst

@@ -28,6 +28,8 @@ Running a server
 
 .. autoclass:: Server
 
+    .. autoattribute:: connections
+
     .. automethod:: serve_forever
 
     .. automethod:: shutdown

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)

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,13 +66,15 @@ class ServerConnection(Connection):
     Args:
         socket: Socket connected to a WebSocket client.
         protocol: Sans-I/O connection.
+        server: Server that manages this connection.
 
     """
 
     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(

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(