Add a close method to the Receiver interface

Also implement the method in all classes implementing the `Receiver`
interface.

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/frequenz/channels/_anycast.py

@@ -391,6 +391,9 @@ def __init__(self, channel: Anycast[_T], /) -> None:
         self._channel: Anycast[_T] = channel
         """The channel that this receiver belongs to."""
 
+        self._closed: bool = False
+        """Whether the receiver is closed."""
+
         self._next: _T | type[_Empty] = _Empty
 
     @override
@@ -409,6 +412,9 @@ async def ready(self) -> bool:
         if self._next is not _Empty:
             return True
 
+        if self._closed:
+            return False
+
         # pylint: disable=protected-access
         while len(self._channel._deque) == 0:
             if self._channel._closed:
@@ -436,6 +442,9 @@ def consume(self) -> _T:
         ):
             raise ReceiverStoppedError(self) from ChannelClosedError(self._channel)
 
+        if self._next is _Empty and self._closed:
+            raise ReceiverStoppedError(self)
+
         assert (
             self._next is not _Empty
         ), "`consume()` must be preceded by a call to `ready()`"
@@ -446,6 +455,14 @@ def consume(self) -> _T:
 
         return next_val
 
+    @override
+    def close(self) -> None:
+        """Close this receiver.
+
+        After closing, the receiver will not be able to receive any more messages.
+        """
+        self._closed = True
+
     def __str__(self) -> str:
         """Return a string representation of this receiver."""
         return f"{self._channel}:{type(self).__name__}"

src/frequenz/channels/_broadcast.py

@@ -417,6 +417,9 @@ def __init__(
         self._q: deque[_T] = deque(maxlen=limit)
         """The receiver's internal message queue."""
 
+        self._closed: bool = False
+        """Whether the receiver is closed."""
+
     def enqueue(self, message: _T, /) -> None:
         """Put a message into this receiver's queue.
 
@@ -466,7 +469,7 @@ async def ready(self) -> bool:
         # consumed, then we return immediately.
         # pylint: disable=protected-access
         while len(self._q) == 0:
-            if self._channel._closed:
+            if self._channel._closed or self._closed:
                 return False
             async with self._channel._recv_cv:
                 await self._channel._recv_cv.wait()
@@ -486,9 +489,25 @@ def consume(self) -> _T:
         if not self._q and self._channel._closed:  # pylint: disable=protected-access
             raise ReceiverStoppedError(self) from ChannelClosedError(self._channel)
 
+        if self._closed:
+            raise ReceiverStoppedError(self)
+
         assert self._q, "`consume()` must be preceded by a call to `ready()`"
         return self._q.popleft()
 
+    @override
+    def close(self) -> None:
+        """Close the receiver.
+
+        After calling this method, new messages will not be received.  Once the
+        receiver's buffer is drained, trying to receive a message will raise a
+        [`ReceiverStoppedError`][frequenz.channels.ReceiverStoppedError].
+        """
+        self._closed = True
+        self._channel._receivers.pop(  # pylint: disable=protected-access
+            hash(self), None
+        )
+
     def __str__(self) -> str:
         """Return a string representation of this receiver."""
         return f"{self._channel}:{type(self).__name__}"

src/frequenz/channels/_merge.py

@@ -191,6 +191,21 @@ def consume(self) -> ReceiverMessageT_co:
 
         return self._results.popleft()
 
+    @override
+    def close(self) -> None:
+        """Close the receiver.
+
+        After calling this method, new messages will not be received.  Once the
+        receiver's buffer is drained, trying to receive a message will raise a
+        [`ReceiverStoppedError`][frequenz.channels.ReceiverStoppedError].
+        """
+        for task in self._pending:
+            if not task.done() and task.get_loop().is_running():
+                task.cancel()
+        self._pending = set()
+        for recv in self._receivers.values():
+            recv.close()
+
     def __str__(self) -> str:
         """Return a string representation of this receiver."""
         if len(self._receivers) > 3:

src/frequenz/channels/_receiver.py

@@ -217,6 +217,15 @@ def consume(self) -> ReceiverMessageT_co:
             ReceiverError: If there is some problem with the receiver.
         """
 
+    def close(self) -> None:
+        """Close the receiver.
+
+        After calling this method, new messages will not be available from the receiver.
+        Once the receiver's buffer is drained, trying to receive a message will raise a
+        [`ReceiverStoppedError`][frequenz.channels.ReceiverStoppedError].
+        """
+        raise NotImplementedError("close() must be implemented by subclasses")
+
     def __aiter__(self) -> Self:
         """Get an async iterator over the received messages.
 
@@ -464,6 +473,16 @@ def consume(self) -> MappedMessageT_co:  # noqa: DOC502
         """
         return self._mapping_function(self._receiver.consume())
 
+    @override
+    def close(self) -> None:
+        """Close the receiver.
+
+        After calling this method, new messages will not be received.  Once the
+        receiver's buffer is drained, trying to receive a message will raise a
+        [`ReceiverStoppedError`][frequenz.channels.ReceiverStoppedError].
+        """
+        self._receiver.close()
+
     def __str__(self) -> str:
         """Return a string representation of the mapper."""
         return f"{type(self).__name__}:{self._receiver}:{self._mapping_function}"
@@ -553,6 +572,16 @@ def consume(self) -> ReceiverMessageT_co:
         self._next_message = _SENTINEL
         return message
 
+    @override
+    def close(self) -> None:
+        """Close the receiver.
+
+        After calling this method, new messages will not be received.  Once the
+        receiver's buffer is drained, trying to receive a message will raise a
+        [`ReceiverStoppedError`][frequenz.channels.ReceiverStoppedError].
+        """
+        self._receiver.close()
+
     def __str__(self) -> str:
         """Return a string representation of the filter."""
         return f"{type(self).__name__}:{self._receiver}:{self._filter_function}"

src/frequenz/channels/event.py

@@ -172,6 +172,11 @@ def consume(self) -> None:
         self._is_set = False
         self._event.clear()
 
+    @override
+    def close(self) -> None:
+        """Close this receiver."""
+        self.stop()
+
     def __str__(self) -> str:
         """Return a string representation of this event."""
         return f"{type(self).__name__}({self._name!r})"

src/frequenz/channels/file_watcher.py

@@ -232,6 +232,11 @@ def consume(self) -> Event:
         change, path_str = self._changes.pop()
         return Event(type=EventType(change), path=pathlib.Path(path_str))
 
+    @override
+    def close(self) -> None:
+        """Close this receiver."""
+        self._stop_event.set()
+
     def __str__(self) -> str:
         """Return a string representation of this receiver."""
         if len(self._paths) > 3:

src/frequenz/channels/timer.py

@@ -745,6 +745,11 @@ def consume(self) -> timedelta:
         self._current_drift = None
         return drift
 
+    @override
+    def close(self) -> None:
+        """Close the timer."""
+        self.stop()
+
     def _now(self) -> int:
         """Return the current monotonic clock time in microseconds.