Implement auto-close support for Broadcast channels

Signed-off-by: Sahas Subramanian <sahas.subramanian@proton.me>

Author: shsms

src/frequenz/channels/_broadcast.py

@@ -21,7 +21,9 @@
 _logger = logging.getLogger(__name__)
 
 
-class Broadcast(Generic[ChannelMessageT]):
+class Broadcast(  # pylint: disable=too-many-instance-attributes
+    Generic[ChannelMessageT]
+):
     """A channel that deliver all messages to all receivers.
 
     # Description
@@ -184,7 +186,13 @@ async def main() -> None:
         ```
     """
 
-    def __init__(self, *, name: str, resend_latest: bool = False) -> None:
+    def __init__(
+        self,
+        *,
+        name: str,
+        resend_latest: bool = False,
+        auto_close: bool = False,
+    ) -> None:
         """Initialize this channel.
 
         Args:
@@ -197,6 +205,8 @@ def __init__(self, *, name: str, resend_latest: bool = False) -> None:
                 wait for the next message on the channel to arrive.  It is safe to be
                 set in data/reporting channels, but is not recommended for use in
                 channels that stream control instructions.
+            auto_close: If True, the channel will be closed when all senders or all
+                receivers are closed.
         """
         self._name: str = name
         """The name of the broadcast channel.
@@ -221,6 +231,9 @@ def __init__(self, *, name: str, resend_latest: bool = False) -> None:
         self._latest: ChannelMessageT | None = None
         """The latest message sent to the channel."""
 
+        self._auto_close_enabled: bool = auto_close
+        """Whether to close the channel when all senders or all receivers are closed."""
+
         self.resend_latest: bool = resend_latest
         """Whether to resend the latest message to new receivers.
 
@@ -367,6 +380,10 @@ async def send(self, message: _T, /) -> None:
             raise SenderError("The channel was closed", self) from ChannelClosedError(
                 self._channel
             )
+        if self._channel._auto_close_enabled and len(self._channel._receivers) == 0:
+            raise SenderError("The channel was closed", self) from ChannelClosedError(
+                self._channel
+            )
         self._channel._latest = message
         stale_refs = []
         for _hash, recv_ref in self._channel._receivers.items():
@@ -394,6 +411,12 @@ async def aclose(self) -> None:
         self._closed = True
         self._channel._sender_count -= 1
 
+        if (
+            self._channel._sender_count == 0  # pylint: disable=protected-access
+            and self._channel._auto_close_enabled  # pylint: disable=protected-access
+        ):
+            await self._channel.aclose()
+
     def __del__(self) -> None:
         """Clean up this sender."""
         if not self._closed:
@@ -527,6 +550,11 @@ async def ready(self) -> bool:
         while len(self._q) == 0:
             if self._channel._closed or self._closed:
                 return False
+            if self._channel._auto_close_enabled and (
+                self._channel._sender_count == 0 or len(self._channel._receivers) == 0
+            ):
+                await self._channel.aclose()
+                return False
             async with self._channel._recv_cv:
                 await self._channel._recv_cv.wait()
         return True