Change tracking batteries status feature to use channel communication

BatteryStatusTracker receives messages from components, parse it.
If status change BatteryStatusTracker sends new status using channel.
BatteriesStatus is now only class that listen for change of
status from any batteries and updates set of working batteries.

Signed-off-by: ela-kotulska-frequenz <[email protected]>

Author: ela-kotulska-frequenz

src/frequenz/sdk/actor/power_distributing/_batteries_status.py

@@ -6,33 +6,75 @@
 
 import asyncio
 import logging
+from dataclasses import dataclass
 from typing import Dict, Set
 
-from ..._internal.asyncio import AsyncConstructible
-from ...microgrid._battery import BatteryStatus, BatteryStatusTracker
+from frequenz.channels import Broadcast, Receiver
+from frequenz.channels.util import MergeNamed
+
+from ..._internal.asyncio import cancel_and_await
+from ...microgrid._battery import BatteryStatus, BatteryStatusTracker, RequestResult
 
 _logger = logging.getLogger(__name__)
 
 
-class BatteriesStatus(AsyncConstructible):
-    """Return status of batteries in the pool.
+@dataclass
+class Status:
+    """Status of the batteries."""
+
+    working: Set[int]
+    """Set of working battery ids."""
+
+    uncertain: Set[int]
+    """Set of batteries that should be used only if there are no working batteries."""
+
+    def get_working_batteries(self, batteries: Set[int]) -> Set[int]:
+        """From the given set of batteries return working batteries.
+
+        Args:
+            batteries: Set of batteries
 
-    To create an instance of this class you should use `async_new` class method.
-    Standard constructor (__init__) is not supported and using it will raise
-    `NotSyncConstructible` error.
+        Returns:
+            Subset with working batteries.
+        """
+        working = self.working.intersection(batteries)
+        if len(working) > 0:
+            return working
+        return self.uncertain.intersection(batteries)
+
+
+@dataclass
+class _BatteryStatusChannelHelper:
+    """Helper class to create battery status channel.
+
+    Channel has only one receiver.
+    Receiver has size 1, because we need only latest status.
     """
 
-    # This is instance attribute.
-    # Don't assign default value, because then it becomes class attribute.
-    _batteries: Dict[int, BatteryStatusTracker]
+    battery_id: int
+    """Id of the battery for which we should create channel."""
+
+    def __post_init__(self):
+        self.name: str = f"battery-{self.battery_id}-status"
+        channel = Broadcast[BatteryStatus](self.name)
+
+        receiver_name = f"{self.name}-receiver"
+        self.receiver = channel.new_receiver(name=receiver_name, maxsize=1)
+        self.sender = channel.new_sender()
+
+
+class BatteriesStatus:
+    """Track status of the batteries.
+
+    Send set of working and uncertain batteries, when the any battery change status.
+    """
 
-    @classmethod
-    async def async_new(
-        cls,
+    def __init__(
+        self,
         battery_ids: Set[int],
         max_data_age_sec: float,
         max_blocking_duration_sec: float,
-    ) -> BatteriesStatus:
+    ) -> None:
         """Create BatteriesStatus instance.
 
         Args:
@@ -46,75 +88,103 @@ async def async_new(
 
         Raises:
             RuntimeError: If any battery has no adjacent inverter.
-
-        Returns:
-            New instance of this class.
         """
-        self: BatteriesStatus = BatteriesStatus.__new__(cls)
-
-        tasks = [
-            BatteryStatusTracker.async_new(
-                id, max_data_age_sec, max_blocking_duration_sec
+        # At first no battery is working, we will get notification when they start
+        # working.
+        self._current_status = Status(working=set(), uncertain=set())
+
+        # Channel for sending results of requests to the batteries
+        request_result_channel = Broadcast[RequestResult]("battery_request_status")
+        self._request_result_sender = request_result_channel.new_sender()
+
+        self._batteries: Dict[str, BatteryStatusTracker] = {}
+
+        # Receivers for individual battery statuses are needed to create a `MergeNamed`
+        # object.
+        receivers: Dict[str, Receiver[BatteryStatus]] = {}
+
+        for battery_id in battery_ids:
+            channel = _BatteryStatusChannelHelper(battery_id)
+            receivers[channel.name] = channel.receiver
+
+            self._batteries[channel.name] = BatteryStatusTracker(
+                battery_id=battery_id,
+                max_data_age_sec=max_data_age_sec,
+                max_blocking_duration_sec=max_blocking_duration_sec,
+                status_sender=channel.sender,
+                request_result_receiver=request_result_channel.new_receiver(
+                    f"battery_{battery_id}_request_status"
+                ),
             )
-            for id in battery_ids
-        ]
 
-        trackers = await asyncio.gather(*tasks)
-        self._batteries = {tracker.battery_id: tracker for tracker in trackers}
+        self._battery_status_channel = MergeNamed[BatteryStatus](
+            **receivers,
+        )
 
-        return self
+        self._task = asyncio.create_task(self._run())
 
-    def get_working_batteries(self, battery_ids: Set[int]) -> Set[int]:
-        """Get subset of battery_ids with working batteries.
+    async def stop(self) -> None:
+        """Stop tracking batteries status."""
+        await asyncio.gather(
+            *[
+                tracker.stop()  # pylint: disable=protected-access
+                for tracker in self._batteries.values()
+            ]
+        )
+        await cancel_and_await(self._task)
 
-        Args:
-            battery_ids: batteries ids
+    async def _run(self) -> None:
+        """Start tracking batteries status."""
+        while True:
+            try:
+                await self._update_status(self._battery_status_channel)
+            except Exception as err:  # pylint: disable=broad-except
+                _logger.error("BatteriesStatus failed with error: %s. Restarting.", err)
 
-        Raises:
-            KeyError: If any battery in the given batteries is not in the pool.
+    async def _update_status(self, status_channel: MergeNamed[BatteryStatus]) -> None:
+        """Wait for any battery to change status and update status.
 
-        Returns:
-            Subset of given batteries with working batteries.
+        Args:
+            status_channel: Receivers packed in Select object.
         """
-        working: Set[int] = set()
-        uncertain: Set[int] = set()
-        for bat_id in battery_ids:
-            if bat_id not in battery_ids:
-                ids = str(self._batteries.keys())
-                raise KeyError(f"No battery {bat_id} in pool. All batteries: {ids}")
-            battery_status = self._batteries[bat_id].get_status()
-            if battery_status == BatteryStatus.WORKING:
-                working.add(bat_id)
-            elif battery_status == BatteryStatus.UNCERTAIN:
-                uncertain.add(bat_id)
+        async for channel_name, status in status_channel:
+            battery_id = self._batteries[channel_name].battery_id
+            if status == BatteryStatus.WORKING:
+                self._current_status.working.add(battery_id)
+                self._current_status.uncertain.discard(battery_id)
+            elif status == BatteryStatus.UNCERTAIN:
+                self._current_status.working.discard(battery_id)
+                self._current_status.uncertain.add(battery_id)
+            elif status == BatteryStatus.NOT_WORKING:
+                self._current_status.working.discard(battery_id)
+                self._current_status.uncertain.discard(battery_id)
+
+            # In the future here we should send status to the subscribed actors
+
+    async def update_status(
+        self, succeed_batteries: Set[int], failed_batteries: Set[int]
+    ) -> None:
+        """Notify which batteries succeed and failed in the request.
+
+        Batteries that failed will be considered as broken and will be blocked for
+        some time.
+        Batteries that succeed will be unblocked.
 
-        if len(working) > 0:
-            return working
-
-        _logger.warning(
-            "There are no working batteries in %s. Falling back to using uncertain batteries %s.",
-            str(battery_ids),
-            str(uncertain),
+        Args:
+            succeed_batteries: Batteries that succeed request
+            failed_batteries: Batteries that failed request
+        """
+        await self._request_result_sender.send(
+            RequestResult(succeed_batteries, failed_batteries)
         )
-        return uncertain
 
-    def update_status(self, succeed_batteries: Set[int], failed_batteries: Set[int]):
-        """Update batteries in pool based on the last result from the request.
+    def get_working_batteries(self, batteries: Set[int]) -> Set[int]:
+        """From the given set of batteries get working.
 
         Args:
-            succeed_batteries: Batteries that succeed request
-            failed_batteries: Batteries that failed request.
+            batteries: Set of batteries
+
+        Returns:
+            Subset with working batteries.
         """
-        for battery_id in succeed_batteries:
-            self._batteries[battery_id].unblock()
-
-        for battery_id in failed_batteries:
-            duration = self._batteries[battery_id].block()
-            if duration > 0:
-                _logger.warning(
-                    "Battery %d failed last response. Block it for %f sec",
-                    battery_id,
-                    duration,
-                )
-
-            self._batteries[battery_id].block()
+        return self._current_status.get_working_batteries(batteries)

src/frequenz/sdk/actor/power_distributing/power_distributing.py

@@ -183,6 +183,12 @@ def __init__(
         self._users_tasks = self._create_users_tasks()
         self._started = asyncio.Event()
 
+        self._batteries_status = BatteriesStatus(
+            battery_ids=set(self._bat_inv_map.keys()),
+            max_blocking_duration_sec=30.0,
+            max_data_age_sec=10.0,
+        )
+
     def _create_users_tasks(self) -> List[asyncio.Task[Empty]]:
         """For each user create a task to wait for request.
 
@@ -246,11 +252,6 @@ async def run(self) -> None:
         await self._create_channels()
 
         api = microgrid.get().api_client
-        battery_pool = await BatteriesStatus.async_new(
-            battery_ids=set(self._bat_inv_map.keys()),
-            max_blocking_duration_sec=30.0,
-            max_data_age_sec=10.0,
-        )
 
         # Wait few seconds to get data from the channels created above.
         await asyncio.sleep(self._wait_for_data_sec)
@@ -261,7 +262,7 @@ async def run(self) -> None:
 
             try:
                 pairs_data: List[InvBatPair] = self._get_components_data(
-                    battery_pool.get_working_batteries(request.batteries)
+                    self._batteries_status.get_working_batteries(request.batteries)
                 )
             except KeyError as err:
                 await user.channel.send(Error(request, str(err)))
@@ -316,8 +317,14 @@ async def run(self) -> None:
                     excess_power=distribution.remaining_power,
                 )
 
-            battery_pool.update_status(succeed_batteries, failed_batteries)
-            await user.channel.send(response)
+            asyncio.gather(
+                *[
+                    self._batteries_status.update_status(
+                        succeed_batteries, failed_batteries
+                    ),
+                    user.channel.send(response),
+                ]
+            )
 
     async def _set_distributed_power(
         self,
@@ -645,4 +652,5 @@ async def _cancel_tasks(self, tasks: Iterable[asyncio.Task[Any]]) -> None:
     async def _stop_actor(self) -> None:
         """Stop all running async tasks."""
         await asyncio.gather(*[cancel_and_await(t) for t in self._users_tasks])
+        await self._batteries_status.stop()
         await self._stop()  # type: ignore # pylint: disable=no-member

src/frequenz/sdk/microgrid/_battery/__init__.py

@@ -6,9 +6,10 @@
 Stores features for the batteries.
 """
 
-from ._status import BatteryStatus, BatteryStatusTracker
+from ._status import BatteryStatus, BatteryStatusTracker, RequestResult
 
 __all__ = [
     "BatteryStatusTracker",
     "BatteryStatus",
+    "RequestResult",
 ]