Reuse GrpcStreamBroadcasters for aggregated data streams

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

Author: shsms

src/frequenz/client/reporting/_client.py

@@ -3,7 +3,7 @@
 
 """Client for requests to the Reporting API."""
 
-from collections.abc import AsyncIterable, AsyncIterator
+from collections.abc import AsyncIterable
 from datetime import datetime, timedelta
 from typing import cast
 
@@ -117,6 +117,17 @@ def __init__(
                 PBReceiveMicrogridSensorsDataStreamResponse, SensorsDataBatch
             ],
         ] = {}
+        self._aggregated_data_streams: dict[
+            tuple[
+                int,  # microgrid_id
+                str,  # metric name
+                str,  # aggregation_formula
+                float | None,  # start_time timestamp
+                float | None,  # end_time timestamp
+                int | None,  # resampling period in seconds
+            ],
+            GrpcStreamBroadcaster[PBAggregatedStreamResponse, MetricSample],
+        ] = {}
 
         self._metadata = (("key", key),) if key else ()
 
@@ -509,7 +520,7 @@ async def receive_aggregated_data(
         start_time: datetime | None,
         end_time: datetime | None,
         resampling_period: timedelta,
-    ) -> AsyncIterator[MetricSample]:
+    ) -> Receiver[MetricSample]:
         """Iterate over aggregated data for a single metric using GrpcStreamBroadcaster.
 
         For now this only supports a single metric and aggregation formula.
@@ -521,63 +532,78 @@ async def receive_aggregated_data(
             end_time: end datetime, if None starts streaming indefinitely from start_time
             resampling_period: The period for resampling the data.
 
-        Yields:
-            An iterator over the aggregated metric samples.
+        Returns:
+            A receiver of `MetricSample`s.
 
         Raises:
             ValueError: If the resampling_period is not provided.
         """
-        if not resampling_period:
-            raise ValueError("resampling_period must be provided")
-
-        aggregation_config = PBAggregationConfig(
-            microgrid_id=microgrid_id,
-            metric=metric.to_proto(),
-            aggregation_formula=aggregation_formula,
+        stream_key = (
+            microgrid_id,
+            metric.name,
+            aggregation_formula,
+            start_time.timestamp() if start_time else None,
+            end_time.timestamp() if end_time else None,
+            round(resampling_period.total_seconds()) if resampling_period else None,
         )
+        if (
+            stream_key not in self._aggregated_data_streams
+            or not self._aggregated_data_streams[stream_key].is_running
+        ):
 
-        def dt2ts(dt: datetime) -> PBTimestamp:
-            ts = PBTimestamp()
-            ts.FromDatetime(dt)
-            return ts
+            if not resampling_period:
+                raise ValueError("resampling_period must be provided")
 
-        time_filter = PBTimeFilter(
-            start=dt2ts(start_time) if start_time else None,
-            end=dt2ts(end_time) if end_time else None,
-        )
+            aggregation_config = PBAggregationConfig(
+                microgrid_id=microgrid_id,
+                metric=metric.to_proto(),
+                aggregation_formula=aggregation_formula,
+            )
 
-        stream_filter = PBAggregatedStreamRequest.AggregationStreamFilter(
-            time_filter=time_filter,
-            resampling_options=PBResamplingOptions(
-                resolution=round(resampling_period.total_seconds())
-            ),
-        )
+            def dt2ts(dt: datetime) -> PBTimestamp:
+                ts = PBTimestamp()
+                ts.FromDatetime(dt)
+                return ts
 
-        request = PBAggregatedStreamRequest(
-            aggregation_configs=[aggregation_config],
-            filter=stream_filter,
-        )
+            time_filter = PBTimeFilter(
+                start=dt2ts(start_time) if start_time else None,
+                end=dt2ts(end_time) if end_time else None,
+            )
 
-        def transform_response(response: PBAggregatedStreamResponse) -> MetricSample:
-            return AggregatedMetric(response).sample()
+            stream_filter = PBAggregatedStreamRequest.AggregationStreamFilter(
+                time_filter=time_filter,
+                resampling_options=PBResamplingOptions(
+                    resolution=round(resampling_period.total_seconds())
+                ),
+            )
 
-        async def stream_method() -> AsyncIterable[PBAggregatedStreamResponse]:
-            call_iterator = self.stub.ReceiveAggregatedMicrogridComponentsDataStream(
-                request, metadata=self._metadata
+            request = PBAggregatedStreamRequest(
+                aggregation_configs=[aggregation_config],
+                filter=stream_filter,
             )
 
-            async for response in cast(
-                AsyncIterable[PBAggregatedStreamResponse], call_iterator
-            ):
-                yield response
+            def transform_response(
+                response: PBAggregatedStreamResponse,
+            ) -> MetricSample:
+                return AggregatedMetric(response).sample()
+
+            async def stream_method() -> AsyncIterable[PBAggregatedStreamResponse]:
+                call_iterator = (
+                    self.stub.ReceiveAggregatedMicrogridComponentsDataStream(
+                        request, metadata=self._metadata
+                    )
+                )
 
-        broadcaster = GrpcStreamBroadcaster(
-            stream_name="aggregated-microgrid-data-stream",
-            stream_method=stream_method,
-            transform=transform_response,
-            retry_strategy=None,
-        )
+                async for response in cast(
+                    AsyncIterable[PBAggregatedStreamResponse], call_iterator
+                ):
+                    yield response
+
+            self._aggregated_data_streams[stream_key] = GrpcStreamBroadcaster(
+                stream_name="aggregated-microgrid-data-stream",
+                stream_method=stream_method,
+                transform=transform_response,
+                retry_strategy=None,
+            )
 
-        receiver = broadcaster.new_receiver()
-        async for data in receiver:
-            yield data
+        return self._aggregated_data_streams[stream_key].new_receiver()

src/frequenz/client/reporting/cli/__main__.py

@@ -171,7 +171,7 @@ async def data_iter() -> AsyncIterator[MetricSample]:
         for formula in formulas:
             assert resampling_period is not None
             for metric in metrics:
-                async for sample in client.receive_aggregated_data(
+                async for sample in await client.receive_aggregated_data(
                     microgrid_id=microgrid_id,
                     metric=metric,
                     aggregation_formula=formula,