diff --git a/docs/reference/features.rst b/docs/reference/features.rst index 8cf18842..cfc6d0cb 100644 --- a/docs/reference/features.rst +++ b/docs/reference/features.rst @@ -36,7 +36,7 @@ Both sides +------------------------------------+--------+--------+--------+--------+--------+ | Send a message | ✅ | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ - | Broadcast a message | ✅ | ✅ | ❌ | — | ✅ | + | Broadcast a message | ✅ | ✅ | ✅ | — | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ | Receive a message | ✅ | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ diff --git a/docs/reference/trio/server.rst b/docs/reference/trio/server.rst index d9b92360..ef04becd 100644 --- a/docs/reference/trio/server.rst +++ b/docs/reference/trio/server.rst @@ -88,6 +88,11 @@ Using a connection .. autoproperty:: close_reason +Broadcast +--------- + +.. autofunction:: broadcast + HTTP Basic Authentication ------------------------- diff --git a/src/websockets/trio/connection.py b/src/websockets/trio/connection.py index b55058eb..58678f6c 100644 --- a/src/websockets/trio/connection.py +++ b/src/websockets/trio/connection.py @@ -4,6 +4,7 @@ import logging import random import struct +import traceback import uuid from collections.abc import AsyncIterable, AsyncIterator, Iterable, Mapping from types import TracebackType @@ -1118,3 +1119,124 @@ async def close_stream(self) -> None: # Unblock coroutines waiting on self.stream_closed. self.stream_closed.set() + + +# 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. + + +async def broadcast( + connections: Iterable[Connection], + message: DataLike, + raise_exceptions: bool = False, + *, + text: bool | None = None, +) -> 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` is equivalent to calling + :meth:`~websockets.trio.connection.Connection.send` for each connection. + It returns when all messages have been sent. + + Unlike :meth:`~websockets.trio.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] = [] + + async def send_message(connection: Connection) -> None: + exception: Exception + + if connection.protocol.state is not OPEN: + return + + if connection.send_in_progress is not None: + 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) + await 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(), + ) + + async with trio.open_nursery() as nursery: + for connection in connections: + nursery.start_soon(send_message, connection) + + if raise_exceptions and exceptions: + raise ExceptionGroup("skipped broadcast", exceptions) + + +# Pretend that broadcast is actually defined in the server module. +broadcast.__module__ = "websockets.asyncio.server" diff --git a/src/websockets/trio/server.py b/src/websockets/trio/server.py index 8d5b104b..82f8221b 100644 --- a/src/websockets/trio/server.py +++ b/src/websockets/trio/server.py @@ -21,11 +21,12 @@ 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 race_events __all__ = [ + "broadcast", "serve", "ServerConnection", "Server", @@ -248,6 +249,7 @@ def connections(self) -> set[ServerConnection]: 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`. """ return { diff --git a/tests/trio/test_connection.py b/tests/trio/test_connection.py index bc767c67..847b6cea 100644 --- a/tests/trio/test_connection.py +++ b/tests/trio/test_connection.py @@ -14,6 +14,7 @@ from websockets.frames import CloseCode, Frame, Opcode from websockets.protocol import CLIENT, SERVER, Protocol, State from websockets.trio.connection import * +from websockets.trio.connection import broadcast from ..protocol import RecordingProtocol from ..utils import MS, alist @@ -1220,7 +1221,7 @@ async def test_close_reason(self): async def test_writing_in_recv_events_fails(self): """Error when responding to incoming frames is correctly reported.""" - # Inject a fault by shutting down the stream for writing — but not the + # Inject a fault by closing the stream for writing. Don't close the # stream for reading because that would terminate the connection. self.connection.stream.send_stream.close() # Receive a ping. Responding with a pong will fail. @@ -1231,7 +1232,7 @@ async def test_writing_in_recv_events_fails(self): async def test_writing_in_send_context_fails(self): """Error when sending outgoing frame is correctly reported.""" - # Inject a fault by shutting down the stream for writing — but not the + # Inject a fault by closing the stream for writing. Don't close the # stream for reading because that would terminate the connection. self.connection.stream.send_stream.close() # Sending a pong will fail. @@ -1258,6 +1259,156 @@ async def test_unexpected_failure_in_send_context(self, send_text): await self.connection.send("😀") self.assertIsInstance(raised.exception.__cause__, AssertionError) + # Test broadcast. + + async def test_broadcast_text(self): + """broadcast broadcasts a text message.""" + await broadcast([self.connection], "😀") + await self.assertFrameSent(Frame(Opcode.TEXT, "😀".encode())) + + async def test_broadcast_text_reports_no_errors(self): + """broadcast broadcasts a text message without raising exceptions.""" + await broadcast([self.connection], "😀", raise_exceptions=True) + await self.assertFrameSent(Frame(Opcode.TEXT, "😀".encode())) + + async def test_broadcast_binary(self): + """broadcast broadcasts a binary message.""" + await broadcast([self.connection], b"\x01\x02\xfe\xff") + await self.assertFrameSent(Frame(Opcode.BINARY, b"\x01\x02\xfe\xff")) + + async def test_broadcast_binary_reports_no_errors(self): + """broadcast broadcasts a binary message without raising exceptions.""" + await broadcast([self.connection], b"\x01\x02\xfe\xff", raise_exceptions=True) + await self.assertFrameSent(Frame(Opcode.BINARY, b"\x01\x02\xfe\xff")) + + async def test_broadcast_text_from_bytes(self): + """broadcast broadcasts a text message from bytes.""" + await broadcast([self.connection], "😀".encode(), text=True) + await self.assertFrameSent(Frame(Opcode.TEXT, "😀".encode())) + + async def test_broadcast_binary_from_str(self): + """broadcast broadcasts a binary message from a str.""" + await broadcast([self.connection], "😀", text=False) + await self.assertFrameSent(Frame(Opcode.BINARY, "😀".encode())) + + async def test_broadcast_no_clients(self): + """broadcast does nothing when called with an empty list of clients.""" + await broadcast([], "😀") + await self.assertNoFrameSent() + + async def test_broadcast_two_clients(self): + """broadcast broadcasts a message to several clients.""" + await broadcast([self.connection, self.connection], "😀") + await self.assertFramesSent( + [ + Frame(Opcode.TEXT, "😀".encode()), + Frame(Opcode.TEXT, "😀".encode()), + ] + ) + + async def test_broadcast_skips_closed_connection(self): + """broadcast ignores closed connections.""" + await self.connection.aclose() + await self.assertFrameSent(Frame(Opcode.CLOSE, b"\x03\xe8")) + + with self.assertNoLogs("websockets", logging.WARNING): + await broadcast([self.connection], "😀") + await self.assertNoFrameSent() + + async def test_broadcast_skips_closing_connection(self): + """broadcast ignores closing connections.""" + async with self.delay_frames_rcvd(MS): + self.nursery.start_soon(self.connection.aclose) + await trio.testing.wait_all_tasks_blocked() + await self.assertFrameSent(Frame(Opcode.CLOSE, b"\x03\xe8")) + + with self.assertNoLogs("websockets", logging.WARNING): + await broadcast([self.connection], "😀") + await self.assertNoFrameSent() + + async def test_broadcast_skips_connection_with_send_blocked(self): + """broadcast logs a warning when a connection is blocked in send.""" + gate = trio.Event() + + async def fragments(): + yield "⏳" + await gate.wait() + + self.nursery.start_soon(self.connection.send, fragments()) + await trio.testing.wait_all_tasks_blocked() + await self.assertFrameSent(Frame(Opcode.TEXT, "⏳".encode(), fin=False)) + + with self.assertLogs("websockets", logging.WARNING) as logs: + await broadcast([self.connection], "😀") + + self.assertEqual( + [record.getMessage() for record in logs.records], + ["skipped broadcast: sending a fragmented message"], + ) + + gate.set() + + async def test_broadcast_reports_connection_with_send_blocked(self): + """broadcast raises exceptions for connections blocked in send.""" + gate = trio.Event() + + async def fragments(): + yield "⏳" + await gate.wait() + + self.nursery.start_soon(self.connection.send, fragments()) + await trio.testing.wait_all_tasks_blocked() + await self.assertFrameSent(Frame(Opcode.TEXT, "⏳".encode(), fin=False)) + + with self.assertRaises(ExceptionGroup) as raised: + await 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() + + async def test_broadcast_skips_connection_failing_to_send(self): + """broadcast logs a warning when a connection fails to send.""" + # Inject a fault by closing the stream for writing. Don't close the + # stream for reading because that would terminate the connection. + self.connection.stream.send_stream.close() + + with self.assertLogs("websockets", logging.WARNING) as logs: + await broadcast([self.connection], "😀") + + self.assertEqual( + [record.getMessage() for record in logs.records], + [ + "skipped broadcast: failed to write message: " + "trio.ClosedResourceError: virtual connection closed" + ], + ) + + async def test_broadcast_reports_connection_failing_to_send(self): + """broadcast raises exceptions for connections failing to send.""" + # Inject a fault by closing the stream for writing. Don't close the + # stream for reading because that would terminate the connection. + self.connection.stream.send_stream.close() + + with self.assertRaises(ExceptionGroup) as raised: + await 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), "virtual connection closed") + self.assertIsInstance(cause, trio.ClosedResourceError) + + async def test_broadcast_type_error(self): + """broadcast raises TypeError when called with an unsupported type.""" + with self.assertRaises(TypeError): + await broadcast([self.connection], ["⏳", "⌛️"]) + class ServerConnectionTests(ClientConnectionTests): LOCAL = SERVER diff --git a/tests/trio/test_server.py b/tests/trio/test_server.py index f4e54849..935d15b8 100644 --- a/tests/trio/test_server.py +++ b/tests/trio/test_server.py @@ -503,7 +503,7 @@ async def close_server(server): uri = get_uri(server) async with connect(uri): self.nursery.start_soon(close_server, server) - await trio.sleep(0) + await trio.testing.wait_all_tasks_blocked() # Server cannot receive new connections. with self.assertRaises(OSError):