Streamer: Write out events on the channels as well

Signed-off-by: Mathias L. Baumann <[email protected]>

Author: Marenz

RELEASE_NOTES.md

@@ -10,7 +10,7 @@
 
 ## New Features
 
-<!-- Here goes the main new features and examples or instructions on how to use them -->
+* The streaming client now also sends state change events out.
 
 ## Bug Fixes
 

src/frequenz/client/base/streaming.py

@@ -6,6 +6,8 @@
 import asyncio
 import logging
 from collections.abc import Callable
+from dataclasses import dataclass
+from enum import StrEnum
 from typing import AsyncIterable, Generic, TypeVar
 
 import grpc.aio
@@ -24,6 +26,33 @@
 """The output type of the stream."""
 
 
+class State(StrEnum):
+    """State of the GrpcStreamBroadcaster."""
+
+    CONNECTING = "connecting"
+    """The broadcaster is connecting to the stream."""
+
+    CONNECTED = "connected"
+    """The broadcaster is connected to the stream."""
+
+    DISCONNECTED = "disconnected"
+    """The broadcaster is disconnected from the stream."""
+
+
+@dataclass(frozen=True, kw_only=True)
+class Message:
+    """Message sent by the broadcaster."""
+
+    state: State
+    """State change of the broadcaster."""
+
+    error: Exception | None = None
+    """An error that occurred during streaming, if any. Defaults to None."""
+
+
+MessageT = TypeVar("MessageT", bound=Message)
+
+
 class GrpcStreamBroadcaster(Generic[InputT, OutputT]):
     """Helper class to handle grpc streaming methods."""
 
@@ -55,14 +84,14 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-positional-argument
         )
         self._retry_on_exhausted_stream = retry_on_exhausted_stream
 
-        self._channel: channels.Broadcast[OutputT] = channels.Broadcast(
+        self._channel: channels.Broadcast[Message | OutputT] = channels.Broadcast(
             name=f"GrpcStreamBroadcaster-{stream_name}"
         )
         self._task = asyncio.create_task(self._run())
 
     def new_receiver(
         self, maxsize: int = 50, warn_on_overflow: bool = True
-    ) -> channels.Receiver[OutputT]:
+    ) -> channels.Receiver[Message | OutputT]:
         """Create a new receiver for the stream.
 
         Args:
@@ -107,20 +136,16 @@ async def _run(self) -> None:
             _logger.info("%s: starting to stream", self._stream_name)
             try:
                 call = self._stream_method()
+                await sender.send(Message(state=State.CONNECTED, error=None))
                 async for msg in call:
                     await sender.send(self._transform(msg))
             except grpc.aio.AioRpcError as err:
                 error = err
-            except Exception as err:  # pylint: disable=broad-except
-                _logger.exception(
-                    "%s: raise an unexpected exception",
-                    self._stream_name,
-                )
-                error = err
             if error is None and not self._retry_on_exhausted_stream:
                 _logger.info(
                     "%s: connection closed, stream exhausted", self._stream_name
                 )
+                await sender.send(Message(state=State.DISCONNECTED, error=None))
                 await self._channel.close()
                 break
             error_str = f"Error: {error}" if error else "Stream exhausted"
@@ -132,6 +157,7 @@ async def _run(self) -> None:
                     self._retry_strategy.get_progress(),
                     error_str,
                 )
+                await sender.send(Message(state=State.DISCONNECTED, error=error))
                 await self._channel.close()
                 break
             _logger.warning(
@@ -141,4 +167,6 @@ async def _run(self) -> None:
                 interval,
                 error_str,
             )
+            await sender.send(Message(state=State.DISCONNECTED, error=error))
+            await sender.send(Message(state=State.CONNECTING, error=None))
             await asyncio.sleep(interval)

tests/streaming/test_grpc_stream_broadcaster.py

@@ -11,8 +11,10 @@
 
 import grpc.aio
 import pytest
+from frequenz.channels import Receiver
 
 from frequenz.client.base import retry, streaming
+from frequenz.client.base.streaming import Message, State
 
 
 def _transformer(x: int) -> str:
@@ -62,6 +64,28 @@ async def asynciter(ready_event: asyncio.Event) -> AsyncIterator[int]:
     await helper.stop()
 
 
+async def _split_message(
+    receiver: Receiver[Message | str],
+) -> tuple[list[str], list[Message]]:
+    """Split the items received from the receiver into items and messages.
+
+    Args:
+        receiver: The receiver to process.
+
+    Returns:
+        A tuple containing a list of transformed items and a list of messages.
+    """
+    items: list[str] = []
+    events: list[Message] = []
+    async for item in receiver:
+        match item:
+            case Message():
+                events.append(item)
+            case str():
+                items.append(item)
+    return items, events
+
+
 class _ErroringAsyncIter(AsyncIterator[int]):
     """Async iterator that raises an error after a certain number of successes."""
 
@@ -93,11 +117,12 @@ async def test_streaming_success_retry_on_exhausted(
     """Test streaming success."""
     caplog.set_level(logging.INFO)
     items: list[str] = []
+    events: list[Message] = []
     async with asyncio.timeout(1):
         receiver = ok_helper.new_receiver()
         receiver_ready_event.set()
-        async for item in receiver:
-            items.append(item)
+        items, events = await _split_message(receiver)
+
     no_retry.next_interval.assert_called_once_with()
     assert items == [
         "transformed_0",
@@ -106,6 +131,10 @@ async def test_streaming_success_retry_on_exhausted(
         "transformed_3",
         "transformed_4",
     ]
+    assert events == [
+        Message(state=State.DISCONNECTED, error=None),
+    ]
+
     assert caplog.record_tuples == [
         (
             "frequenz.client.base.streaming",
@@ -128,11 +157,13 @@ async def test_streaming_success(
     """Test streaming success."""
     caplog.set_level(logging.INFO)
     items: list[str] = []
+    events: list[Message] = []
+
     async with asyncio.timeout(1):
         receiver = ok_helper.new_receiver()
         receiver_ready_event.set()
-        async for item in receiver:
-            items.append(item)
+        items, events = await _split_message(receiver)
+
     assert (
         no_retry.next_interval.call_count == 0
     ), "next_interval should not be called when streaming is successful"
@@ -144,6 +175,9 @@ async def test_streaming_success(
         "transformed_3",
         "transformed_4",
     ]
+    assert events == [
+        Message(state=State.DISCONNECTED, error=None),
+    ]
     assert caplog.record_tuples == [
         (
             "frequenz.client.base.streaming",
@@ -191,13 +225,13 @@ async def test_streaming_error(  # pylint: disable=too-many-arguments
     )
 
     items: list[str] = []
+    events: list[Message] = []
     async with AsyncExitStack() as stack:
         stack.push_async_callback(helper.stop)
 
         receiver = helper.new_receiver()
         receiver_ready_event.set()
-        async for item in receiver:
-            items.append(item)
+        items, events = await _split_message(receiver)
 
     no_retry.next_interval.assert_called_once_with()
     assert items == [f"transformed_{i}" for i in range(successes)]
@@ -251,13 +285,13 @@ async def test_retry_next_interval_zero(  # pylint: disable=too-many-arguments
     )
 
     items: list[str] = []
+    events: list[Message] = []
     async with AsyncExitStack() as stack:
         stack.push_async_callback(helper.stop)
 
         receiver = helper.new_receiver()
         receiver_ready_event.set()
-        async for item in receiver:
-            items.append(item)
+        items, events = await _split_message(receiver)
 
     assert not items
     assert mock_retry.next_interval.mock_calls == [mock.call(), mock.call()]
@@ -282,3 +316,46 @@ async def test_retry_next_interval_zero(  # pylint: disable=too-many-arguments
             f"giving up. Error: {expected_error_str}.",
         ),
     ]
+
+
+async def test_messages_on_retry(
+    receiver_ready_event: asyncio.Event,  # pylint: disable=redefined-outer-name
+) -> None:
+    """Test that messages are sent on retry."""
+    helper = streaming.GrpcStreamBroadcaster(
+        stream_name="test_helper",
+        stream_method=lambda: _ErroringAsyncIter(
+            grpc.aio.AioRpcError(
+                code=_NamedMagicMock(name="mock grpc code"),
+                initial_metadata=mock.MagicMock(),
+                trailing_metadata=mock.MagicMock(),
+                details="mock details",
+                debug_error_string="mock debug_error_string",
+            ),
+            receiver_ready_event,
+        ),
+        transform=_transformer,
+        retry_strategy=retry.LinearBackoff(
+            limit=1,
+            interval=0.01,
+        ),
+        retry_on_exhausted_stream=True,
+    )
+
+    items: list[str] = []
+    events: list[Message] = []
+    async with AsyncExitStack() as stack:
+        stack.push_async_callback(helper.stop)
+
+        receiver = helper.new_receiver()
+        receiver_ready_event.set()
+        items, events = await _split_message(receiver)
+
+    assert items == []
+    assert [e.state for e in events] == [
+        State.CONNECTED,
+        State.DISCONNECTED,
+        State.CONNECTING,
+        State.CONNECTED,
+        State.DISCONNECTED,
+    ]