diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index ddd21319..5f80bd19 100644 --- a/docs/project/changelog.rst +++ b/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. diff --git a/docs/reference/features.rst b/docs/reference/features.rst index e5f6e0de..4e4ca613 100644 --- a/docs/reference/features.rst +++ b/docs/reference/features.rst @@ -35,7 +35,7 @@ Both sides +------------------------------------+--------+--------+--------+--------+ | Send a message | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+ - | Broadcast a message | ✅ | ❌ | — | ✅ | + | Broadcast a message | ✅ | ✅ | — | ✅ | +------------------------------------+--------+--------+--------+--------+ | Receive a message | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+ diff --git a/docs/reference/sync/server.rst b/docs/reference/sync/server.rst index 32f14653..6df014f0 100644 --- a/docs/reference/sync/server.rst +++ b/docs/reference/sync/server.rst @@ -85,6 +85,11 @@ Using a connection .. autoproperty:: close_reason +Broadcast +--------- + +.. autofunction:: broadcast + HTTP Basic Authentication ------------------------- diff --git a/docs/topics/broadcast.rst b/docs/topics/broadcast.rst index 66b0819b..21369b0e 100644 --- a/docs/topics/broadcast.rst +++ b/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 diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index a89d080c..3b20e03d 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -1,5 +1,6 @@ from __future__ import annotations +import concurrent.futures import contextlib import logging import random @@ -7,6 +8,7 @@ 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" diff --git a/src/websockets/sync/server.py b/src/websockets/sync/server.py index aea09449..6f730a93 100644 --- a/src/websockets/sync/server.py +++ b/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): diff --git a/tests/sync/test_connection.py b/tests/sync/test_connection.py index 6c45bc74..b44646ef 100644 --- a/tests/sync/test_connection.py +++ b/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