Support closing of Senders

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

Author: shsms

src/frequenz/channels/__init__.py

@@ -100,7 +100,7 @@
     select,
     selected_from,
 )
-from ._sender import Sender, SenderError
+from ._sender import Sender, SenderClosedError, SenderError
 
 __all__ = [
     "Anycast",
@@ -120,6 +120,7 @@
     "SelectError",
     "Selected",
     "Sender",
+    "SenderClosedError",
     "SenderError",
     "SenderMessageT_co",
     "SenderMessageT_contra",

src/frequenz/channels/_anycast.py

@@ -15,7 +15,7 @@
 from ._exceptions import ChannelClosedError
 from ._generic import ChannelMessageT
 from ._receiver import Receiver, ReceiverStoppedError
-from ._sender import Sender, SenderError
+from ._sender import Sender, SenderClosedError, SenderError
 
 _logger = logging.getLogger(__name__)
 
@@ -327,6 +327,9 @@ def __init__(self, channel: Anycast[_T], /) -> None:
         self._channel: Anycast[_T] = channel
         """The channel that this sender belongs to."""
 
+        self._closed: bool = False
+        """Whether the sender is closed."""
+
     @override
     async def send(self, message: _T, /) -> None:
         """Send a message across the channel.
@@ -343,7 +346,11 @@ async def send(self, message: _T, /) -> None:
             SenderError: If the underlying channel was closed.
                 A [ChannelClosedError][frequenz.channels.ChannelClosedError] is
                 set as the cause.
+            SenderClosedError: If this sender was closed.
         """
+        if self._closed:
+            raise SenderClosedError(self)
+
         # pylint: disable=protected-access
         if self._channel._closed:
             raise SenderError("The channel was closed", self) from ChannelClosedError(
@@ -367,6 +374,16 @@ async def send(self, message: _T, /) -> None:
             self._channel._recv_cv.notify(1)
         # pylint: enable=protected-access
 
+    @override
+    async def aclose(self) -> None:
+        """Close this sender.
+
+        After closing, the sender will not be able to send any more messages. Any
+        attempt to send a message through a closed sender will raise a
+        [SenderError][frequenz.channels.SenderError].
+        """
+        self._closed = True
+
     def __str__(self) -> str:
         """Return a string representation of this sender."""
         return f"{self._channel}:{type(self).__name__}"

src/frequenz/channels/_broadcast.py

@@ -16,7 +16,7 @@
 from ._exceptions import ChannelClosedError
 from ._generic import ChannelMessageT
 from ._receiver import Receiver, ReceiverStoppedError
-from ._sender import Sender, SenderError
+from ._sender import Sender, SenderClosedError, SenderError
 
 _logger = logging.getLogger(__name__)
 
@@ -334,6 +334,9 @@ def __init__(self, channel: Broadcast[_T], /) -> None:
         self._channel: Broadcast[_T] = channel
         """The broadcast channel this sender belongs to."""
 
+        self._closed: bool = False
+        """Whether this sender is closed."""
+
     @override
     async def send(self, message: _T, /) -> None:
         """Send a message to all broadcast receivers.
@@ -345,7 +348,10 @@ async def send(self, message: _T, /) -> None:
             SenderError: If the underlying channel was closed.
                 A [ChannelClosedError][frequenz.channels.ChannelClosedError] is
                 set as the cause.
+            SenderClosedError: If this sender was closed.
         """
+        if self._closed:
+            raise SenderClosedError(self)
         # pylint: disable=protected-access
         if self._channel._closed:
             raise SenderError("The channel was closed", self) from ChannelClosedError(
@@ -365,6 +371,16 @@ async def send(self, message: _T, /) -> None:
             self._channel._recv_cv.notify_all()
         # pylint: enable=protected-access
 
+    @override
+    async def aclose(self) -> None:
+        """Close this sender.
+
+        After a sender is closed, it can no longer be used to send messages. Any
+        attempt to send a message through a closed sender will raise a
+        [SenderError][frequenz.channels.SenderError].
+        """
+        self._closed = True
+
     def __str__(self) -> str:
         """Return a string representation of this sender."""
         return f"{self._channel}:{type(self).__name__}"

src/frequenz/channels/_sender.py

@@ -70,6 +70,15 @@ async def send(self, message: SenderMessageT_contra, /) -> None:
             SenderError: If there was an error sending the message.
         """
 
+    @abstractmethod
+    async def aclose(self) -> None:
+        """Close this sender.
+
+        After a sender is closed, it can no longer be used to send messages. Any
+        attempt to send a message through a closed sender will raise a
+        [SenderClosedError][frequenz.channels.SenderClosedError].
+        """
+
 
 class SenderError(Error, Generic[SenderMessageT_co]):
     """An error that originated in a [Sender][frequenz.channels.Sender].
@@ -88,3 +97,15 @@ def __init__(self, message: str, sender: Sender[SenderMessageT_co]):
         super().__init__(message)
         self.sender: Sender[SenderMessageT_co] = sender
         """The sender where the error happened."""
+
+
+class SenderClosedError(SenderError[SenderMessageT_co]):
+    """An error indicating that a send operation was attempted a closed sender."""
+
+    def __init__(self, sender: Sender[SenderMessageT_co]):
+        """Initialize this error.
+
+        Args:
+            sender: The [Sender][frequenz.channels.Sender] that was closed.
+        """
+        super().__init__("Sender is closed", sender)

src/frequenz/channels/experimental/_relay_sender.py

@@ -7,15 +7,15 @@
 to the senders it was created with.
 """
 
-import typing
+import asyncio
 
 from typing_extensions import override
 
 from .._generic import SenderMessageT_contra
 from .._sender import Sender
 
 
-class RelaySender(typing.Generic[SenderMessageT_contra], Sender[SenderMessageT_contra]):
+class RelaySender(Sender[SenderMessageT_contra]):
     """A Sender for sending messages to multiple senders.
 
     The `RelaySender` class takes multiple senders and forwards all the messages sent to
@@ -57,3 +57,8 @@ async def send(self, message: SenderMessageT_contra, /) -> None:
         """
         for sender in self._senders:
             await sender.send(message)
+
+    @override
+    async def aclose(self) -> None:
+        """Close this sender."""
+        await asyncio.gather(*(sender.aclose() for sender in self._senders))