Close all *pool's receivers when stopping the pool (#1150)

Author: shsms

RELEASE_NOTES.md

@@ -63,3 +63,5 @@ This release includes a new `ConfigManager` class to simplify managing the confi
 - Fix a bug where if a string was passed to the `ConfigManagingActor` it would be interpreted as a sequence of 1 character strings.
 
 - Remove a confusing log message in the power distributing actor.
+
+- Close all receivers owned by a *pool when stopping the pool.

src/frequenz/sdk/timeseries/battery_pool/_battery_pool_reference_store.py

@@ -90,9 +90,12 @@ def __init__(  # pylint: disable=too-many-arguments
         self._working_batteries: set[int] = set()
 
         self._update_battery_status_task: asyncio.Task[None] | None = None
+        self._batteries_status_receiver: Receiver[ComponentPoolStatus] = (
+            batteries_status_receiver
+        )
         if self._batteries:
             self._update_battery_status_task = asyncio.create_task(
-                self._update_battery_status(batteries_status_receiver)
+                self._update_battery_status(self._batteries_status_receiver)
             )
 
         self._min_update_interval: timedelta = min_update_interval
@@ -128,6 +131,7 @@ async def stop(self) -> None:
         if self._update_battery_status_task:
             tasks_to_stop.append(cancel_and_await(self._update_battery_status_task))
         await asyncio.gather(*tasks_to_stop)
+        self._batteries_status_receiver.close()
 
     def _get_all_batteries(self) -> frozenset[int]:
         """Get all batteries from the microgrid.

src/frequenz/sdk/timeseries/battery_pool/_component_metric_fetcher.py

@@ -21,6 +21,7 @@
     ComponentMetricId,
     InverterData,
 )
+from typing_extensions import override
 
 from ..._internal._asyncio import AsyncConstructible
 from ..._internal._constants import MAX_BATTERY_DATA_AGE_SEC
@@ -68,6 +69,10 @@ async def async_new(
     async def fetch_next(self) -> ComponentMetricsData | None:
         """Fetch metrics for this component."""
 
+    @abstractmethod
+    def stop(self) -> None:
+        """Stop the metric fetcher."""
+
 
 class LatestMetricsFetcher(ComponentMetricFetcher, Generic[T], ABC):
     """Subscribe for the latest component data and extract the needed metrics."""
@@ -143,6 +148,11 @@ async def fetch_next(self) -> ComponentMetricsData | None:
 
         return ComponentMetricsData(self._component_id, data.timestamp, metrics)
 
+    @override
+    def stop(self) -> None:
+        """Stop the metric fetcher."""
+        self._receiver.close()
+
     @abstractmethod
     def _extract_metric(self, data: T, mid: ComponentMetricId) -> float: ...
 

src/frequenz/sdk/timeseries/battery_pool/_methods.py

@@ -112,6 +112,8 @@ def __init__(
             set()
         )
 
+        self._fetchers: dict[int, ComponentMetricFetcher] = {}
+
     @classmethod
     def name(cls) -> str:
         """Get name of the method.
@@ -166,28 +168,29 @@ async def stop(self) -> None:
         await asyncio.gather(
             *[cancel_and_await(self._send_task), cancel_and_await(self._update_task)]
         )
+        for fetcher in self._fetchers.values():
+            fetcher.stop()
 
-    async def _create_data_fetchers(self) -> dict[int, ComponentMetricFetcher]:
+    async def _create_data_fetchers(self) -> None:
         fetchers: dict[int, ComponentMetricFetcher] = {
             cid: await LatestBatteryMetricsFetcher.async_new(cid, metrics)
             for cid, metrics in self._metric_calculator.battery_metrics.items()
         }
+        self._fetchers.update(fetchers)
         inverter_fetchers = {
             cid: await LatestInverterMetricsFetcher.async_new(cid, metrics)
             for cid, metrics in self._metric_calculator.inverter_metrics.items()
         }
-        fetchers.update(inverter_fetchers)
-        return fetchers
+        self._fetchers.update(inverter_fetchers)
 
-    def _remove_metric_fetcher(
-        self, fetchers: dict[int, ComponentMetricFetcher], component_id: int
-    ) -> None:
+    def _remove_metric_fetcher(self, component_id: int) -> None:
         _logger.error(
             "Removing component %d from the %s formula.",
             component_id,
             self._result_channel._name,  # pylint: disable=protected-access
         )
-        fetchers.pop(component_id)
+        fetcher = self._fetchers.pop(component_id)
+        fetcher.stop()
 
     def _metric_updated(self, new_metrics: ComponentMetricsData) -> bool:
         cid = new_metrics.component_id
@@ -197,11 +200,11 @@ def _metric_updated(self, new_metrics: ComponentMetricsData) -> bool:
 
     async def _update_and_notify(self) -> None:
         """Receive component metrics and send notification when they change."""
-        fetchers = await self._create_data_fetchers()
+        await self._create_data_fetchers()
 
         self._pending_data_fetchers = {
             asyncio.create_task(fetcher.fetch_next(), name=str(cid))
-            for cid, fetcher in fetchers.items()
+            for cid, fetcher in self._fetchers.items()
         }
         while len(self._pending_data_fetchers) > 0:
             done, self._pending_data_fetchers = await asyncio.wait(
@@ -210,7 +213,7 @@ async def _update_and_notify(self) -> None:
             for item in done:
                 metrics = item.result()
                 if metrics is None:
-                    self._remove_metric_fetcher(fetchers, int(item.get_name()))
+                    self._remove_metric_fetcher(int(item.get_name()))
                     continue
                 if self._metric_updated(metrics):
                     self._update_event.set()
@@ -220,7 +223,7 @@ async def _update_and_notify(self) -> None:
                 self._cached_metrics[cid] = metrics
                 # Add fetcher back to the processing list.
                 self._pending_data_fetchers.add(
-                    asyncio.create_task(fetchers[cid].fetch_next(), name=str(cid))
+                    asyncio.create_task(self._fetchers[cid].fetch_next(), name=str(cid))
                 )
 
     async def _send_on_update(self, min_update_interval: timedelta) -> None:

src/frequenz/sdk/timeseries/ev_charger_pool/_ev_charger_pool_reference_store.py

@@ -107,3 +107,4 @@ async def stop(self) -> None:
         """Stop all tasks and channels owned by the EVChargerPool."""
         await self.formula_pool.stop()
         await self.bounds_tracker.stop()
+        self.status_receiver.close()

src/frequenz/sdk/timeseries/formula_engine/_formula_engine.py

@@ -117,6 +117,9 @@ async def _stop(self) -> None:
         if self._task is None:
             return
         await cancel_and_await(self._task)
+        _, fetchers = self._builder.finalize()
+        for fetcher in fetchers.values():
+            fetcher.stream.close()
 
     @classmethod
     def from_receiver(

src/frequenz/sdk/timeseries/pv_pool/_pv_pool_reference_store.py

@@ -108,3 +108,4 @@ async def stop(self) -> None:
         """Stop all tasks and channels owned by the PVInverterPool."""
         await self.formula_pool.stop()
         await self.bounds_tracker.stop()
+        self.status_receiver.close()