WIP: Support dipatches refering to different actors

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

Author: Marenz

src/frequenz/dispatch/_actors_service.py

@@ -7,11 +7,13 @@
 import logging
 from collections.abc import Callable
 from dataclasses import dataclass
+from re import L
 from typing import Any
 
 from frequenz.channels import Broadcast, Receiver
 from frequenz.client.dispatch.types import TargetComponents
 from frequenz.sdk.actor import Actor, BackgroundService
+from grpc.aio import Call
 
 from ._dispatch import Dispatch
 
@@ -41,7 +43,7 @@ class DispatchActorsService(BackgroundService):
     import os
     import asyncio
     from typing import override
-    from frequenz.dispatch import Dispatcher, DispatchManagingActor, DispatchUpdate
+    from frequenz.dispatch import Dispatcher, DispatchActorsService, DispatchUpdate
     from frequenz.client.dispatch.types import TargetComponents
     from frequenz.client.common.microgrid.components import ComponentCategory
     from frequenz.channels import Receiver, Broadcast, select, selected_from
@@ -128,7 +130,7 @@ async def main():
         def actor_factory(initial_dispatch, dispatch_receiver):
             return MyActor.new_with_dispatch(initial_dispatch, dispatch_receiver)
 
-        managing_actor = DispatchManagingActor(
+        managing_actor = DispatchActorsService(
             actor_factory=actor_factory,
             running_status_receiver=status_receiver,
         )
@@ -141,18 +143,22 @@ def __init__(
         self,
         actor_factory: Callable[[DispatchUpdate, Receiver[DispatchUpdate]], Actor],
         running_status_receiver: Receiver[Dispatch],
+        map_dispatch: Callable[[Dispatch], int] = lambda dispatch: hash(dispatch.type),
     ) -> None:
         """Initialize the dispatch handler.
 
         Args:
             actor_factory: A callable that creates an actor with some initial dispatch
                 information.
             running_status_receiver: The receiver for dispatch running status changes.
+            map_dispatch: A function to identify to which actor a dispatch
+                          refers. By default, it uses the hash of the dispatch type.
         """
         super().__init__()
+        self._map_dispatch = map_dispatch
         self._dispatch_rx = running_status_receiver
         self._actor_factory = actor_factory
-        self._actor: Actor | None = None
+        self._actors: dict[int, Actor] = {}
         self._updates_channel = Broadcast[DispatchUpdate](
             name="dispatch_updates_channel", resend_latest=True
         )
@@ -178,7 +184,9 @@ async def _start_actor(self, dispatch: Dispatch) -> None:
             options=dispatch.payload,
         )
 
-        if self._actor:
+        actor: Actor | None = self._actors.get(self._map_dispatch(dispatch))
+
+        if actor:
             sent_str = ""
             if self._updates_sender is not None:
                 sent_str = ", sent a dispatch update instead of creating a new actor"
@@ -190,8 +198,10 @@ async def _start_actor(self, dispatch: Dispatch) -> None:
             )
         else:
             _logger.info("Starting actor for dispatch type %r", dispatch.type)
-            self._actor = self._actor_factory(dispatch_update, self.new_receiver())
-            self._actor.start()
+            actor = self._actor_factory(dispatch_update, self.new_receiver())
+            self._actors[self._map_dispatch(dispatch)] = actor
+
+            actor.start()
 
     async def _stop_actor(self, stopping_dispatch: Dispatch, msg: str) -> None:
         """Stop all actors.
@@ -200,13 +210,12 @@ async def _stop_actor(self, stopping_dispatch: Dispatch, msg: str) -> None:
             stopping_dispatch: The dispatch that is stopping the actor.
             msg: The message to be passed to the actors being stopped.
         """
-        if self._actor is None:
+        if actor := self._actors.pop(self._map_dispatch(stopping_dispatch), None):
+            await actor.stop(msg)
+        else:
             _logger.warning(
                 "Actor for dispatch type %r is not running", stopping_dispatch.type
             )
-        else:
-            await self._actor.stop(msg)
-            self._actor = None
 
     async def _run(self) -> None:
         """Wait for dispatches and handle them."""

src/frequenz/dispatch/_bg_service.py

@@ -53,18 +53,18 @@ def filter(self) -> Callable[[Dispatch], bool]:
         """
 
 
-class MergeByType(MergeStrategy):
-    """Merge running intervals based on the dispatch type."""
+class MergeByX(MergeStrategy):
+    """Merge running intervals based on a dispatch configuration."""
 
     @property
     @override
     def filter(self) -> Callable[[Dispatch], bool]:
         """Get a filter function for dispatches."""
         return self._filter_func
 
-    def criteria(self, previous_dispatch: Dispatch, new_dispatch: Dispatch) -> bool:
-        """Define the criteria for checking other running dispatches."""
-        return previous_dispatch.type == new_dispatch.type
+    @abstractmethod
+    def identity(self, dispatch: Dispatch) -> int:
+        """Identity function for the merge criteria."""
 
     def _filter_func(self, new_dispatch: Dispatch) -> bool:
         """Filter dispatches based on the merge strategy.
@@ -76,27 +76,40 @@ def _filter_func(self, new_dispatch: Dispatch) -> bool:
         if new_dispatch.started:
             return True
 
-        # pylint: disable=protected-access
         other_dispatches_running = any(
             dispatch.started
             for dispatch in self._scheduler.dispatches
-            if self.criteria(dispatch, new_dispatch)
+            if self.identity(dispatch) == self.identity(new_dispatch)
         )
-        # pylint: enable=protected-access
 
         return not other_dispatches_running
 
 
+class MergeByType(MergeByX):
+    """Merge running intervals based on the dispatch type."""
+
+    @override
+    def identity(self, dispatch: Dispatch) -> int:
+        """Identity function for the merge criteria."""
+        return hash(dispatch.type)
+
+
 class MergeByTypeTarget(MergeByType):
     """Merge running intervals based on the dispatch type and target."""
 
     @override
-    def criteria(self, previous_dispatch: Dispatch, new_dispatch: Dispatch) -> bool:
-        """Define the criteria for checking other running dispatches."""
-        return (
-            previous_dispatch.type == new_dispatch.type
-            and previous_dispatch.target == new_dispatch.target
-        )
+    def identity(self, dispatch: Dispatch) -> int:
+        """Identity function for the merge criteria."""
+        return hash((dispatch.type, dispatch.target))
+
+
+class MergeByTypeID(MergeByX):
+    """Merge running intervals based on the dispatch type and ID."""
+
+    @override
+    def identity(self, dispatch: Dispatch) -> int:
+        """Identity function for the merge criteria."""
+        return hash((dispatch.type, dispatch.id))
 
 
 # pylint: disable=too-many-instance-attributes