Raise ChannelClosedError to signal that a channel is closed

... instead of using `None`

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

Author: shsms

benchmarks/benchmark_anycast.py

@@ -51,10 +51,7 @@ async def benchmark_anycast(
     recv_trackers = [0]
 
     async def update_tracker_on_receive(chan: Receiver[int]) -> None:
-        while True:
-            msg = await chan.receive()
-            if msg is None:
-                return
+        async for _ in chan:
             recv_trackers[0] += 1
 
     receivers = []

benchmarks/benchmark_broadcast.py

@@ -69,10 +69,7 @@ async def benchmark_broadcast(
     recv_trackers = [0]
 
     async def update_tracker_on_receive(chan: Receiver[int]) -> None:
-        while True:
-            msg = await chan.receive()
-            if msg is None:
-                return
+        async for _ in chan:
             recv_trackers[0] += 1
 
     receivers = []

src/frequenz/channels/base_classes.py

@@ -6,12 +6,41 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Callable, Generic, Optional, TypeVar
+from typing import Any, Callable, Generic, Optional, TypeVar
 
 T = TypeVar("T")
 U = TypeVar("U")
 
 
+class ChannelError(RuntimeError):
+    """Base channel error.
+
+    All exceptions generated by channels inherit from this exception.
+    """
+
+    def __init__(self, message: Any, channel: Any = None):
+        """Create a ChannelError instance.
+
+        Args:
+            message: An error message.
+            channel: A reference to the channel that encountered the error.
+        """
+        super().__init__(message)
+        self.channel: Any = channel
+
+
+class ChannelClosedError(ChannelError):
+    """Error raised when trying to operate on a closed channel."""
+
+    def __init__(self, channel: Any = None):
+        """Create a `ChannelClosedError` instance.
+
+        Args:
+            channel: A reference to the channel that was closed.
+        """
+        super().__init__(f"Channel {channel} was closed", channel)
+
+
 class Sender(ABC, Generic[T]):
     """A channel Sender."""
 
@@ -50,16 +79,19 @@ def __aiter__(self) -> Receiver[T]:
         """
         return self
 
-    async def receive(self) -> Optional[T]:
+    async def receive(self) -> T:
         """Receive a message from the channel.
 
+        Raises:
+            ChannelClosedError: if the underlying channel is closed.
+
         Returns:
             The received message.
         """
         try:
             received = await self.__anext__()  # pylint: disable=unnecessary-dunder-call
-        except StopAsyncIteration:
-            return None
+        except StopAsyncIteration as exc:
+            raise ChannelClosedError() from exc
         return received
 
     def map(self, call: Callable[[T], U]) -> Receiver[U]:

tests/test_anycast.py

@@ -5,7 +5,10 @@
 
 import asyncio
 
+import pytest
+
 from frequenz.channels import Anycast, Receiver, Sender
+from frequenz.channels.base_classes import ChannelClosedError
 
 
 async def test_anycast() -> None:
@@ -29,8 +32,9 @@ async def send_msg(chan: Sender[int]) -> None:
 
     async def update_tracker_on_receive(receiver_id: int, chan: Receiver[int]) -> None:
         while True:
-            msg = await chan.receive()
-            if msg is None:
+            try:
+                msg = await chan.receive()
+            except ChannelClosedError:
                 return
             recv_trackers[receiver_id] += msg
             # without the sleep, decomissioning receivers temporarily, all
@@ -56,7 +60,8 @@ async def update_tracker_on_receive(receiver_id: int, chan: Receiver[int]) -> No
     await receivers_runs
 
     assert await after_close_sender.send(5) is False
-    assert await after_close_receiver.receive() is None
+    with pytest.raises(ChannelClosedError):
+        await after_close_receiver.receive()
 
     actual_sum = 0
     for ctr in recv_trackers:
@@ -79,7 +84,8 @@ async def test_anycast_after_close() -> None:
 
     assert await sender.send(5) is False
     assert await receiver.receive() == 2
-    assert await receiver.receive() is None
+    with pytest.raises(ChannelClosedError):
+        await receiver.receive()
 
 
 async def test_anycast_full() -> None:

tests/test_broadcast.py

@@ -9,6 +9,7 @@
 import pytest
 
 from frequenz.channels import Broadcast, Receiver, Sender
+from frequenz.channels.base_classes import ChannelClosedError
 
 
 async def test_broadcast() -> None:
@@ -32,8 +33,9 @@ async def send_msg(chan: Sender[int]) -> None:
 
     async def update_tracker_on_receive(receiver_id: int, chan: Receiver[int]) -> None:
         while True:
-            msg = await chan.receive()
-            if msg is None:
+            try:
+                msg = await chan.receive()
+            except ChannelClosedError:
                 return
             recv_trackers[receiver_id] += msg
 
@@ -68,7 +70,8 @@ async def test_broadcast_after_close() -> None:
     await bcast.close()
 
     assert await sender.send(5) is False
-    assert await receiver.receive() is None
+    with pytest.raises(ChannelClosedError):
+        await receiver.receive()
 
 
 async def test_broadcast_overflow() -> None:

tests/test_merge.py

@@ -26,7 +26,7 @@ async def send(ch1: Sender[int], ch2: Sender[int]) -> None:
 
     merge = Merge(chan1.get_receiver(), chan2.get_receiver())
     results: List[int] = []
-    while item := await merge.receive():
+    async for item in merge:
         results.append(item)
     await senders
     for ctr in range(5):

tests/test_mergenamed.py

@@ -27,7 +27,7 @@ async def send(ch1: Sender[int], ch2: Sender[int]) -> None:
 
     merge = MergeNamed(**recvs)
     results: List[Tuple[str, int]] = []
-    while item := await merge.receive():
+    async for item in merge:
         results.append(item)
     await senders
     for ctr in range(5):