Add broadcast to the threading implementation.

Fix #1466.

Author: aaugustin

docs/project/changelog.rst

@@ -56,6 +56,8 @@ New features
 
 * Validated compatibility with Python 3.15.
 
+* Added :func:`~sync.server.broadcast` to the :mod:`threading` implementation.
+
 * Added the ``--insecure`` option to the ``websockets`` CLI to disable TLS
   certificate validation.
 

docs/reference/features.rst

@@ -35,7 +35,7 @@ Both sides
     +------------------------------------+--------+--------+--------+--------+
     | Send a message                     | ✅     | ✅     | ✅     | ✅     |
     +------------------------------------+--------+--------+--------+--------+
-    | Broadcast a message                | ✅     | ❌     | —      | ✅     |
+    | Broadcast a message                | ✅     | ✅     | —      | ✅     |
     +------------------------------------+--------+--------+--------+--------+
     | Receive a message                  | ✅     | ✅     | ✅     | ✅     |
     +------------------------------------+--------+--------+--------+--------+

docs/reference/sync/server.rst

@@ -85,6 +85,11 @@ Using a connection
 
     .. autoproperty:: close_reason
 
+Broadcast
+---------
+
+.. autofunction:: broadcast
+
 HTTP Basic Authentication
 -------------------------
 

docs/topics/broadcast.rst

@@ -16,7 +16,8 @@ WebSocket servers often send the same message to all connected clients or to a
 subset of clients for which the message is relevant.
 
 Let's explore options for broadcasting a message, explain the design of
-:func:`~asyncio.server.broadcast`, and discuss alternatives.
+:func:`~asyncio.server.broadcast` in the :mod:`asyncio` implementation, and
+discuss alternatives.
 
 For each option, we'll provide a connection handler called ``handler()`` and a
 function or coroutine called ``broadcast()`` that sends a message to all

src/websockets/sync/connection.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+import concurrent.futures
 import contextlib
 import logging
 import random
 import socket
 import struct
 import threading
 import time
+import traceback
 import uuid
 from collections.abc import Iterable, Iterator, Mapping
 from types import TracebackType
@@ -1082,3 +1084,128 @@ def close_socket(self) -> None:
 
             # Acknowledge pings sent with the ack_on_close option.
             self.terminate_pending_pings()
+
+
+# broadcast() is defined in the connection module even though it's primarily
+# used by servers and documented in the server module because it works with
+# client connections too and because it's easier to test together with the
+# Connection class.
+
+
+def broadcast(
+    connections: Iterable[Connection],
+    message: DataLike,
+    raise_exceptions: bool = False,
+    *,
+    text: bool | None = None,
+    **kwargs: Any,
+) -> None:
+    """
+    Broadcast a message to several WebSocket connections.
+
+    A string (:class:`str`) is sent as a Text_ frame. A bytestring or bytes-like
+    object (:class:`bytes`, :class:`bytearray`, or :class:`memoryview`) is sent
+    as a Binary_ frame.
+
+    .. _Text: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
+    .. _Binary: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
+
+    You may override this behavior with the ``text`` argument:
+
+    * Set ``text=True`` to send an UTF-8 bytestring or bytes-like object
+      (:class:`bytes`, :class:`bytearray`, or :class:`memoryview`) in a
+      Text_ frame. This improves performance when the message is already
+      UTF-8 encoded, for example if the message contains JSON and you're
+      using a JSON library that produces a bytestring.
+    * Set ``text=False`` to send a string (:class:`str`) in a Binary_
+      frame. This may be useful for servers that expect binary frames
+      instead of text frames.
+
+    :func:`broadcast` relies on :class:`concurrent.futures.ThreadPoolExecutor`
+    to send the messages. Make sure the thread pool is large enough relative to
+    the number of clients, so that slow or stuck connections don't clog it. If
+    that's an issue, then you should be using an asynchronous implementation.
+    You can configure the thread pool by passing additional keyword arguments to
+    :func:`broadcast`, such as ``max_workers``.
+
+    Unlike :meth:`~websockets.asyncio.connection.Connection.send`,
+    :func:`broadcast` doesn't support sending fragmented messages. Indeed,
+    fragmentation is useful for sending large messages without buffering them in
+    memory, while :func:`broadcast` buffers one copy per connection as fast as
+    possible.
+
+    :func:`broadcast` skips connections that aren't open in order to avoid
+    errors on connections where the closing handshake is in progress.
+
+    :func:`broadcast` ignores failures to write the message on some connections.
+    It continues writing to other connections. You may set ``raise_exceptions``
+    to :obj:`True` to record failures and raise all exceptions in a :pep:`654`
+    :exc:`ExceptionGroup`.
+
+    While :func:`broadcast` makes more sense for servers, it works identically
+    with clients, if you have a use case for opening connections to many servers
+    and broadcasting a message to them.
+
+    Args:
+        websockets: WebSocket connections to which the message will be sent.
+        message: Message to send.
+        raise_exceptions: Whether to raise an exception in case of failures.
+        text: Force sending in Text_ or Binary_ frames.
+
+    Raises:
+        TypeError: If ``message`` doesn't have a supported type.
+
+    """
+    if isinstance(message, str):
+        send_method = "send_binary" if text is False else "send_text"
+        message = message.encode()
+    elif isinstance(message, BytesLike):
+        send_method = "send_text" if text is True else "send_binary"
+    else:
+        raise TypeError("data must be str or bytes")
+
+    if raise_exceptions:
+        exceptions: list[Exception] = []
+
+    def send_message(connection: Connection) -> None:
+        exception: Exception
+
+        with connection.protocol_mutex:
+            if connection.protocol.state is not OPEN:
+                return
+
+            if connection.send_in_progress:
+                if raise_exceptions:
+                    exception = ConcurrencyError("sending a fragmented message")
+                    exceptions.append(exception)
+                else:
+                    connection.logger.warning(
+                        "skipped broadcast: sending a fragmented message",
+                    )
+                return
+
+            try:
+                # Call connection.protocol.send_text or send_binary.
+                # Either way, message is already converted to bytes.
+                getattr(connection.protocol, send_method)(message)
+                connection.send_data()
+            except Exception as write_exception:
+                if raise_exceptions:
+                    exception = RuntimeError("failed to write message")
+                    exception.__cause__ = write_exception
+                    exceptions.append(exception)
+                else:
+                    connection.logger.warning(
+                        "skipped broadcast: failed to write message: %s",
+                        traceback.format_exception_only(write_exception)[0].strip(),
+                    )
+
+    with concurrent.futures.ThreadPoolExecutor(**kwargs) as executor:
+        executor.map(send_message, connections)
+
+    if raise_exceptions and exceptions:
+        raise ExceptionGroup("skipped broadcast", exceptions)
+
+
+# Pretend that broadcast is actually defined in the server module.
+broadcast.__module__ = "websockets.sync.server"

src/websockets/sync/server.py

@@ -27,11 +27,18 @@
 from ..protocol import CONNECTING, OPEN, Event
 from ..server import ServerProtocol
 from ..typing import LoggerLike, Origin, StatusLike, Subprotocol
-from .connection import Connection
+from .connection import Connection, broadcast
 from .utils import Deadline
 
 
-__all__ = ["serve", "unix_serve", "ServerConnection", "Server", "basic_auth"]
+__all__ = [
+    "broadcast",
+    "serve",
+    "unix_serve",
+    "ServerConnection",
+    "Server",
+    "basic_auth",
+]
 
 
 class ServerConnection(Connection):

tests/sync/test_connection.py

@@ -15,6 +15,7 @@
 from websockets.frames import CloseCode, Frame, Opcode
 from websockets.protocol import CLIENT, SERVER, Protocol, State
 from websockets.sync.connection import *
+from websockets.sync.connection import broadcast
 
 from ..protocol import RecordingProtocol
 from ..utils import MS
@@ -54,6 +55,11 @@ def assertFrameSent(self, frame):
         self.wait_for_remote_side()
         self.assertEqual(self.remote_connection.protocol.get_frames_rcvd(), [frame])
 
+    def assertFramesSent(self, frames):
+        """Check that several frames were sent."""
+        self.wait_for_remote_side()
+        self.assertEqual(self.remote_connection.protocol.get_frames_rcvd(), frames)
+
     def assertNoFrameSent(self):
         """Check that no frame was sent."""
         self.wait_for_remote_side()
@@ -992,6 +998,153 @@ def test_unexpected_failure_in_send_context(self, send_text):
             self.connection.send("😀")
         self.assertIsInstance(raised.exception.__cause__, AssertionError)
 
+    # Test broadcast.
+
+    def test_broadcast_text(self):
+        """broadcast broadcasts a text message."""
+        broadcast([self.connection], "😀")
+        self.assertFrameSent(Frame(Opcode.TEXT, "😀".encode()))
+
+    def test_broadcast_text_reports_no_errors(self):
+        """broadcast broadcasts a text message without raising exceptions."""
+        broadcast([self.connection], "😀", raise_exceptions=True)
+        self.assertFrameSent(Frame(Opcode.TEXT, "😀".encode()))
+
+    def test_broadcast_binary(self):
+        """broadcast broadcasts a binary message."""
+        broadcast([self.connection], b"\x01\x02\xfe\xff")
+        self.assertFrameSent(Frame(Opcode.BINARY, b"\x01\x02\xfe\xff"))
+
+    def test_broadcast_binary_reports_no_errors(self):
+        """broadcast broadcasts a binary message without raising exceptions."""
+        broadcast([self.connection], b"\x01\x02\xfe\xff", raise_exceptions=True)
+        self.assertFrameSent(Frame(Opcode.BINARY, b"\x01\x02\xfe\xff"))
+
+    def test_broadcast_text_from_bytes(self):
+        """broadcast broadcasts a text message from bytes."""
+        broadcast([self.connection], "😀".encode(), text=True)
+        self.assertFrameSent(Frame(Opcode.TEXT, "😀".encode()))
+
+    def test_broadcast_binary_from_str(self):
+        """broadcast broadcasts a binary message from a str."""
+        broadcast([self.connection], "😀", text=False)
+        self.assertFrameSent(Frame(Opcode.BINARY, "😀".encode()))
+
+    def test_broadcast_no_clients(self):
+        """broadcast does nothing when called with an empty list of clients."""
+        broadcast([], "😀")
+        self.assertNoFrameSent()
+
+    def test_broadcast_two_clients(self):
+        """broadcast broadcasts a message to several clients."""
+        broadcast([self.connection, self.connection], "😀")
+        self.assertFramesSent(
+            [
+                Frame(Opcode.TEXT, "😀".encode()),
+                Frame(Opcode.TEXT, "😀".encode()),
+            ]
+        )
+
+    def test_broadcast_skips_closed_connection(self):
+        """broadcast ignores closed connections."""
+        self.connection.close()
+        self.assertFrameSent(Frame(Opcode.CLOSE, b"\x03\xe8"))
+
+        with self.assertNoLogs("websockets", logging.WARNING):
+            broadcast([self.connection], "😀")
+        self.assertNoFrameSent()
+
+    def test_broadcast_skips_closing_connection(self):
+        """broadcast ignores closing connections."""
+
+        def closer():
+            with self.delay_frames_rcvd(MS):
+                self.connection.close()
+
+        with self.run_in_thread(closer):
+            self.assertFrameSent(Frame(Opcode.CLOSE, b"\x03\xe8"))
+
+            with self.assertNoLogs("websockets", logging.WARNING):
+                broadcast([self.connection], "😀")
+            self.assertNoFrameSent()
+
+    def test_broadcast_skips_connection_with_send_blocked(self):
+        """broadcast logs a warning when a connection is blocked in send."""
+        gate = threading.Event()
+
+        def fragments():
+            yield "⏳"
+            gate.wait()
+
+        with self.run_in_thread(self.connection.send, args=(fragments(),)):
+            self.assertFrameSent(Frame(Opcode.TEXT, "⏳".encode(), fin=False))
+
+            with self.assertLogs("websockets", logging.WARNING) as logs:
+                broadcast([self.connection], "😀")
+
+            self.assertEqual(
+                [record.getMessage() for record in logs.records],
+                ["skipped broadcast: sending a fragmented message"],
+            )
+
+            gate.set()
+
+    def test_broadcast_reports_connection_with_send_blocked(self):
+        """broadcast raises exceptions for connections blocked in send."""
+        gate = threading.Event()
+
+        def fragments():
+            yield "⏳"
+            gate.wait()
+
+        with self.run_in_thread(self.connection.send, args=(fragments(),)):
+            self.assertFrameSent(Frame(Opcode.TEXT, "⏳".encode(), fin=False))
+
+            with self.assertRaises(ExceptionGroup) as raised:
+                broadcast([self.connection], "😀", raise_exceptions=True)
+
+            self.assertEqual(
+                str(raised.exception), "skipped broadcast (1 sub-exception)"
+            )
+            exc = raised.exception.exceptions[0]
+            self.assertEqual(str(exc), "sending a fragmented message")
+            self.assertIsInstance(exc, ConcurrencyError)
+
+            gate.set()
+
+    @patch("socket.socket.sendall", side_effect=BrokenPipeError(32, "Broken pipe"))
+    def test_broadcast_skips_connection_failing_to_send(self, sendall):
+        """broadcast logs a warning when a connection fails to send."""
+        with self.assertLogs("websockets", logging.WARNING) as logs:
+            broadcast([self.connection], "😀")
+
+        self.assertEqual(
+            [record.getMessage() for record in logs.records],
+            [
+                "skipped broadcast: failed to write message: "
+                "BrokenPipeError: [Errno 32] Broken pipe"
+            ],
+        )
+
+    @patch("socket.socket.sendall", side_effect=BrokenPipeError(32, "Broken pipe"))
+    def test_broadcast_reports_connection_failing_to_send(self, sendall):
+        """broadcast raises exceptions for connections failing to send."""
+        with self.assertRaises(ExceptionGroup) as raised:
+            broadcast([self.connection], "😀", raise_exceptions=True)
+
+        self.assertEqual(str(raised.exception), "skipped broadcast (1 sub-exception)")
+        exc = raised.exception.exceptions[0]
+        self.assertEqual(str(exc), "failed to write message")
+        self.assertIsInstance(exc, RuntimeError)
+        cause = exc.__cause__
+        self.assertEqual(str(cause), "[Errno 32] Broken pipe")
+        self.assertIsInstance(cause, BrokenPipeError)
+
+    def test_broadcast_type_error(self):
+        """broadcast raises TypeError when called with an unsupported type."""
+        with self.assertRaises(TypeError):
+            broadcast([self.connection], ["⏳", "⌛️"])
+
 
 class ServerConnectionTests(ClientConnectionTests):
     LOCAL = SERVER