Split __anext__ in receivers into two methods: _ready and _get

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

Author: shsms

src/frequenz/channels/anycast.py

@@ -7,7 +7,7 @@
 
 from asyncio import Condition
 from collections import deque
-from typing import Deque, Generic
+from typing import Deque, Generic, Optional
 
 from frequenz.channels.base_classes import Receiver as BaseReceiver
 from frequenz.channels.base_classes import Sender as BaseSender
@@ -162,26 +162,38 @@ def __init__(self, chan: Anycast[T]) -> None:
             chan: A reference to the channel that this receiver belongs to.
         """
         self._chan = chan
+        self._next: Optional[T] = None
 
-    async def __anext__(self) -> T:
-        """Receive a message from the channel.
-
-        Waits for an message to become available, and returns that message.
-        When there are multiple receivers for the channel, only one receiver
-        will receive each message.
+    async def _ready(self) -> None:
+        """Wait until the receiver is ready with a value.
 
         Raises:
-            StopAsyncIteration: When the channel is closed.
-
-        Returns:
-            The received message.
+            StopAsyncIteration: if the underlying channel is closed.
         """
+        # if a message is already ready, then return immediately.
+        if self._next is not None:
+            return
+
         while len(self._chan.deque) == 0:
             if self._chan.closed:
                 raise StopAsyncIteration()
             async with self._chan.recv_cv:
                 await self._chan.recv_cv.wait()
-        ret = self._chan.deque.popleft()
+        self._next = self._chan.deque.popleft()
         async with self._chan.send_cv:
             self._chan.send_cv.notify(1)
-        return ret
+
+    def _get(self) -> T:
+        """Return the latest value once `_ready()` is complete.
+
+        Raises:
+            EOFError: When called before a call to `_ready()` finishes.
+
+        Returns:
+            The next value that was received.
+        """
+        if self._next is None:
+            raise EOFError("_get was called before a call to _ready finished.")
+        next_val = self._next
+        self._next = None
+        return next_val

src/frequenz/channels/base_classes.py

@@ -60,7 +60,6 @@ async def send(self, msg: T) -> bool:
 class Receiver(ABC, Generic[T]):
     """A channel Receiver."""
 
-    @abstractmethod
     async def __anext__(self) -> T:
         """Await the next value in the async iteration over received values.
 
@@ -70,6 +69,32 @@ async def __anext__(self) -> T:
         Raises:
             StopAsyncIteration: if the underlying channel is closed.
         """
+        await self._ready()
+        return self._get()
+
+    @abstractmethod
+    async def _ready(self) -> None:
+        """Wait until the receiver is ready with a value.
+
+        Once a call to `_ready` has finished, the value should be read with a call to
+        `_get()`.
+
+        Raises:
+            StopAsyncIteration: if the underlying channel is closed.
+        """
+
+    @abstractmethod
+    def _get(self) -> T:
+        """Return the latest value once `_ready` is complete.
+
+        `_ready()` must be called before each call to `_get()`.
+
+        Returns:
+            The next value received.
+
+        Raises:
+            StopAsyncIteration: if the underlying channel is closed.
+        """
 
     def __aiter__(self) -> Receiver[T]:
         """Initialize the async iterator over received values.
@@ -168,11 +193,14 @@ def __init__(self, recv: Receiver[T], transform: Callable[[T], U]) -> None:
         self._recv = recv
         self._transform = transform
 
-    async def __anext__(self) -> U:
-        """Return a transformed message received from the input channel.
+    async def _ready(self) -> None:
+        """Wait until the receiver is ready with a value."""
+        await self._recv._ready()  # pylint: disable=protected-access
+
+    def _get(self) -> U:
+        """Return a transformed value once `_ready()` is complete.
 
         Returns:
-            A transformed message.
+            The next value that was received.
         """
-        msg = await self._recv.__anext__()
-        return self._transform(msg)
+        return self._transform(self._recv._get())  # pylint: disable=protected-access

src/frequenz/channels/bidirectional.py

@@ -82,10 +82,14 @@ async def send(self, msg: T) -> bool:
         """
         return await self._sender.send(msg)
 
-    async def __anext__(self) -> U:
-        """Receive a value from the other side.
+    async def _ready(self) -> None:
+        """Wait until the receiver is ready with a value."""
+        await self._receiver._ready()  # pylint: disable=protected-access
+
+    def _get(self) -> U:
+        """Return the latest value once `_ready` is complete.
 
         Returns:
-            Received value, or `None` if the channels are closed.
+            The next value that was received.
         """
-        return await self._receiver.__anext__()
+        return self._receiver._get()  # pylint: disable=protected-access

src/frequenz/channels/broadcast.py

@@ -249,32 +249,32 @@ def __len__(self) -> int:
         """
         return len(self._q)
 
-    async def __anext__(self) -> T:
-        """Receive a message from the Broadcast channel.
-
-        Waits until there are messages available in the channel and returns
-        them.  If there are no remaining messages in the buffer and the channel
-        is closed, returns `None` immediately.
-
-        If [into_peekable()][frequenz.channels.Receiver.into_peekable] is called
-        on a broadcast `Receiver`, further calls to `receive`, will raise an
-        `EOFError`.
+    async def _ready(self) -> None:
+        """Wait until the receiver is ready with a value.
 
         Raises:
-            StopAsyncIteration: When the channel is closed.
-            EOFError: when the receiver has been converted into a `Peekable`.
-
-        Returns:
-            `None`, if the channel is closed, a message otherwise.
+            EOFError: When called before a call to `_ready()` finishes.
+            StopAsyncIteration: if the underlying channel is closed.
         """
         if not self._active:
             raise EOFError("This receiver is no longer active.")
 
+        # Use a while loop here, to handle spurious wakeups of condition variables.
+        #
+        # The condition also makes sure that if there are already messages ready to be
+        # consumed, then we return immediately.
         while len(self._q) == 0:
             if self._chan.closed:
                 raise StopAsyncIteration()
             async with self._chan.recv_cv:
                 await self._chan.recv_cv.wait()
+
+    def _get(self) -> T:
+        """Return the latest value once `_ready` is complete.
+
+        Returns:
+            The next value that was received.
+        """
         ret = self._q.popleft()
         return ret
 

src/frequenz/channels/merge.py

@@ -44,21 +44,18 @@ def __del__(self) -> None:
         for task in self._pending:
             task.cancel()
 
-    async def __anext__(self) -> T:
-        """Wait until there's a message in any of the channels.
+    async def _ready(self) -> None:
+        """Wait until the receiver is ready with a value.
 
         Raises:
-            StopAsyncIteration: When the channel is closed.
-
-        Returns:
-            The next message that was received, or `None`, if all channels have
-                closed.
+            StopAsyncIteration: if the underlying channel is closed.
         """
         # we use a while loop to continue to wait for new data, in case the
         # previous `wait` completed because a channel was closed.
         while True:
+            # if there are messages waiting to be consumed, return immediately.
             if len(self._results) > 0:
-                return self._results.popleft()
+                return
 
             if len(self._pending) == 0:
                 raise StopAsyncIteration()
@@ -73,5 +70,20 @@ async def __anext__(self) -> T:
                 result = item.result()
                 self._results.append(result)
                 self._pending.add(
+                    # pylint: disable=unnecessary-dunder-call
                     asyncio.create_task(self._receivers[name].__anext__(), name=name)
                 )
+
+    def _get(self) -> T:
+        """Return the latest value once `_ready` is complete.
+
+        Raises:
+            EOFError: When called before a call to `_ready()` finishes.
+
+        Returns:
+            The next value that was received.
+        """
+        if not self._results:
+            raise EOFError("_get called before _ready finished.")
+
+        return self._results.popleft()

src/frequenz/channels/merge_named.py

@@ -31,21 +31,18 @@ def __del__(self) -> None:
         for task in self._pending:
             task.cancel()
 
-    async def __anext__(self) -> Tuple[str, T]:
+    async def _ready(self) -> None:
         """Wait until there's a message in any of the channels.
 
         Raises:
             StopAsyncIteration: When the channel is closed.
-
-        Returns:
-            The next message that was received, or `None`, if all channels have
-                closed.
         """
         # we use a while loop to continue to wait for new data, in case the
         # previous `wait` completed because a channel was closed.
         while True:
+            # if there are messages waiting to be consumed, return immediately.
             if len(self._results) > 0:
-                return self._results.popleft()
+                return
 
             if len(self._pending) == 0:
                 raise StopAsyncIteration()
@@ -60,5 +57,20 @@ async def __anext__(self) -> Tuple[str, T]:
                 result = item.result()
                 self._results.append((name, result))
                 self._pending.add(
+                    # pylint: disable=unnecessary-dunder-call
                     asyncio.create_task(self._receivers[name].__anext__(), name=name)
                 )
+
+    def _get(self) -> Tuple[str, T]:
+        """Return the latest value once `_ready` is complete.
+
+        Raises:
+            EOFError: When called before a call to `_ready()` finishes.
+
+        Returns:
+            The next value that was received, along with its name.
+        """
+        if not self._results:
+            raise EOFError("_get called before _ready finished.")
+
+        return self._results.popleft()

src/frequenz/channels/utils/file_watcher.py

@@ -8,6 +8,7 @@
 from typing import List, Optional, Set, Union
 
 from watchfiles import Change, awatch
+from watchfiles.main import FileChange
 
 from frequenz.channels.base_classes import Receiver
 
@@ -52,6 +53,7 @@ def __init__(
                 and pathlib.Path(path_str).is_file()
             ),
         )
+        self._changes: Set[FileChange] = set()
 
     def __del__(self) -> None:
         """Cleanup registered watches.
@@ -62,7 +64,7 @@ def __del__(self) -> None:
         """
         self._stop_event.set()
 
-    async def __anext__(self) -> pathlib.Path:
+    async def _ready(self) -> None:
         """Wait for the next file event and return its path.
 
         Raises:
@@ -71,13 +73,17 @@ async def __anext__(self) -> pathlib.Path:
         Returns:
             Path of next file.
         """
-        while True:
-            changes = await self._awatch.__anext__()
-            for change in changes:
-                # Tuple of (Change, path) returned by watchfiles
-                if change is None or len(change) != 2:
-                    raise StopAsyncIteration()
-
-                _, path_str = change
-                path = pathlib.Path(path_str)
-                return path
+        # if there are messages waiting to be consumed, return immediately.
+        if self._changes:
+            return
+
+        self._changes = await self._awatch.__anext__()
+
+    def _get(self) -> pathlib.Path:
+        change = self._changes.pop()
+        # Tuple of (Change, path) returned by watchfiles
+        if change is None or len(change) != 2:
+            raise StopAsyncIteration()
+        _, path_str = change
+        path = pathlib.Path(path_str)
+        return path

src/frequenz/channels/utils/timer.py

@@ -5,6 +5,7 @@
 
 import asyncio
 from datetime import datetime, timedelta, timezone
+from typing import Optional
 
 from frequenz.channels.base_classes import Receiver
 
@@ -61,6 +62,7 @@ def __init__(self, interval: float) -> None:
         self._stopped = False
         self._interval = timedelta(seconds=interval)
         self._next_msg_time = datetime.now(timezone.utc) + self._interval
+        self._now: Optional[datetime] = None
 
     def reset(self) -> None:
         """Reset the timer to start timing from `now`."""
@@ -75,11 +77,12 @@ def stop(self) -> None:
         """
         self._stopped = True
 
-    async def __anext__(self) -> datetime:
+    async def _ready(self) -> None:
         """Return the current time (in UTC) once the next tick is due.
 
         Raises:
-            StopAsyncIteration: When the channel is closed.
+            StopAsyncIteration: if [stop()][frequenz.channels.Timer.stop] has been
+                called on the timer.
 
         Returns:
             The time of the next tick in UTC or `None` if
@@ -90,6 +93,10 @@ async def __anext__(self) -> datetime:
             * **v0.11.0:** Returns a timezone-aware datetime with UTC timezone
               instead of a native datetime object.
         """
+        # if there are messages waiting to be consumed, return immediately.
+        if self._now is not None:
+            return
+
         if self._stopped:
             raise StopAsyncIteration()
         now = datetime.now(timezone.utc)
@@ -98,7 +105,21 @@ async def __anext__(self) -> datetime:
             await asyncio.sleep(diff.total_seconds())
             now = datetime.now(timezone.utc)
             diff = self._next_msg_time - now
+        self._now = now
+
+        self._next_msg_time = self._now + self._interval
+
+    def _get(self) -> datetime:
+        """Return the latest value once `_ready` is complete.
 
-        self._next_msg_time = now + self._interval
+        Raises:
+            EOFError: When called before a call to `_ready()` finishes.
 
+        Returns:
+            The timestamp for the next tick.
+        """
+        if self._now is None:
+            raise EOFError("_get called before _ready finished")
+        now = self._now
+        self._now = None
         return now