Add & use close channel method in the registry (#225)

Add all required code to automatically close channels when we're done
sending
values in the mock_microgrid.

Required for the datasourcing actor benchmark.

- Add private method to allow closing of channels
- Update mock_microgrid to close channels when done sending.
- Update datasourcing actor to close channels.

Author: mathias-baumann-frequenz

RELEASE_NOTES.md

@@ -12,6 +12,7 @@
 
 * A new class `SerializableRingbuffer` is now available, extending the `OrderedRingBuffer` class with the ability to load & dump the data to disk.
 * Add the `run(*actors)` function for running and synchronizing the execution of actors. This new function simplifies the way actors are managed on the client side, allowing for a cleaner and more streamlined approach. Users/apps can now run actors simply by calling run(actor1, actor2, actor3...) without the need to manually call join() and deal with linting errors.
+* The datasourcing actor now automatically closes all sending channels when the input channel closes.
 
 ## Bug Fixes
 

src/frequenz/sdk/actor/_channel_registry.py

@@ -49,3 +49,15 @@ def new_receiver(self, key: str) -> Receiver[Any]:
         if key not in self._channels:
             self._channels[key] = Broadcast(f"{self._name}-{key}")
         return self._channels[key].new_receiver()
+
+    async def _close_channel(self, key: str) -> None:
+        """Close a channel with the given key.
+
+        This method is private and should only be used in special cases.
+
+        Args:
+            key: A key to identify the channel.
+        """
+        if key in self._channels:
+            if channel := self._channels.pop(key, None):
+                await channel.close()

src/frequenz/sdk/actor/_data_sourcing/microgrid_api_source.py

@@ -343,18 +343,39 @@ async def _handle_data_stream(
             stream_senders = self._get_metric_senders(
                 category, self._req_streaming_metrics[comp_id]
             )
-        api_data_receiver = self.comp_data_receivers[comp_id]
+        api_data_receiver: Receiver[Any] = self.comp_data_receivers[comp_id]
+
+        senders_done: asyncio.Event = asyncio.Event()
+        pending_messages = 0
 
         def process_msg(data: Any) -> None:
             tasks = []
             for extractor, senders in stream_senders:
                 for sender in senders:
                     tasks.append(sender.send(Sample(data.timestamp, extractor(data))))
             asyncio.gather(*tasks)
+            nonlocal pending_messages
+            pending_messages -= 1
+            if pending_messages == 0:
+                senders_done.set()
 
         async for data in api_data_receiver:
+            pending_messages += 1
+            senders_done.clear()
             process_msg(data)
 
+        while pending_messages > 0:
+            await senders_done.wait()
+
+        await asyncio.gather(
+            *[
+                # pylint: disable=protected-access
+                self._registry._close_channel(r.get_channel_name())
+                for requests in self._req_streaming_metrics[comp_id].values()
+                for r in requests
+            ]
+        )
+
     async def _update_streams(
         self,
         comp_id: int,

tests/timeseries/mock_microgrid.py

@@ -148,6 +148,8 @@ async def _comp_data_send_task(
                 await self._microgrid.send(make_comp_data(val_to_send, timestamp))
             await asyncio.sleep(self._sample_rate_s)
 
+        await self._microgrid.close_channel(comp_id)
+
     def _start_meter_streaming(self, meter_id: int) -> None:
         self._streaming_coros.append(
             self._comp_data_send_task(

tests/utils/mock_microgrid.py

@@ -49,6 +49,13 @@ def __init__(self, components: Set[Component], connections: Set[Connection]):
         meter_channels = self._create_meter_channels()
         ev_charger_channels = self._create_ev_charger_channels()
 
+        self._all_channels: Dict[int, Broadcast[Any]] = {
+            **bat_channels,
+            **inv_channels,
+            **meter_channels,
+            **ev_charger_channels,
+        }
+
         mock_api = self._create_mock_api(
             bat_channels, inv_channels, meter_channels, ev_charger_channels
         )
@@ -128,6 +135,15 @@ async def send(self, data: ComponentData) -> bool:
 
         raise RuntimeError(f"{type(data)} is not supported in MockMicrogridClient.")
 
+    async def close_channel(self, cid: int) -> None:
+        """Close channel for given component id.
+
+        Args:
+            cid: Component id
+        """
+        if cid in self._all_channels:
+            await self._all_channels[cid].close()
+
     def _create_battery_channels(self) -> Dict[int, Broadcast[BatteryData]]:
         """Create channels for the batteries.