Add broadcast to the trio implementation.

docs/reference/features.rst

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

docs/reference/trio/server.rst

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

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"

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 {

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

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):