ActorDispatcher: Support dispatches referring to different actors

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

Author: Marenz

src/frequenz/dispatch/_actor_dispatcher.py

@@ -41,7 +41,7 @@ class ActorDispatcher(BackgroundService):
     import os
     import asyncio
     from typing import override
-    from frequenz.dispatch import Dispatcher, DispatchManagingActor, DispatchInfo
+    from frequenz.dispatch import Dispatcher, ActorDispatcher, DispatchInfo
     from frequenz.client.dispatch.types import TargetComponents
     from frequenz.client.common.microgrid.components import ComponentCategory
     from frequenz.channels import Receiver, Broadcast, select, selected_from
@@ -125,9 +125,10 @@ async def main():
 
         status_receiver = dispatcher.new_running_state_event_receiver("EXAMPLE_TYPE")
 
-        managing_actor = DispatchManagingActor(
+        managing_actor = ActorDispatcher(
             actor_factory=MyActor.new_with_dispatch,
             running_status_receiver=status_receiver,
+            map_dispatch=lambda dispatch: dispatch.id,
         )
 
         await run(managing_actor)
@@ -138,18 +139,21 @@ def __init__(
         self,
         actor_factory: Callable[[DispatchInfo, Receiver[DispatchInfo]], Actor],
         running_status_receiver: Receiver[Dispatch],
+        map_dispatch: Callable[[Dispatch], int],
     ) -> 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.
         """
         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[DispatchInfo](
             name="dispatch_updates_channel", resend_latest=True
         )
@@ -167,7 +171,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"
@@ -179,10 +185,12 @@ async def _start_actor(self, dispatch: Dispatch) -> None:
             )
         else:
             _logger.info("Starting actor for dispatch type %r", dispatch.type)
-            self._actor = self._actor_factory(
+            actor = self._actor_factory(
                 dispatch_update, self._updates_channel.new_receiver()
             )
-            self._actor.start()
+            self._actors[self._map_dispatch(dispatch)] = actor
+
+            actor.start()
 
     async def _stop_actor(self, stopping_dispatch: Dispatch, msg: str) -> None:
         """Stop all actors.
@@ -191,13 +199,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."""

tests/test_mananging_actor.py

@@ -69,21 +69,13 @@ class TestEnv:
     running_status_sender: Sender[Dispatch]
     generator: DispatchGenerator = DispatchGenerator()
 
-    @property
-    def actor(self) -> MockActor | None:
+    def actor(self, identity: int) -> MockActor:
         """Return the actor."""
         # pylint: disable=protected-access
-        if self.actors_service._actor is None:
-            return None
-        return cast(MockActor, self.actors_service._actor)
+        assert identity in self.actors_service._actors
+        return cast(MockActor, self.actors_service._actors[identity])
         # pylint: enable=protected-access
 
-    @property
-    def updates_receiver(self) -> Receiver[DispatchInfo]:
-        """Return the updates receiver."""
-        assert self.actor is not None
-        return self.actor.receiver
-
 
 @fixture
 async def test_env() -> AsyncIterator[TestEnv]:
@@ -93,6 +85,7 @@ async def test_env() -> AsyncIterator[TestEnv]:
     actors_service = ActorDispatcher(
         actor_factory=MockActor,
         running_status_receiver=channel.new_receiver(),
+        map_dispatch=lambda dispatch: dispatch.id,
     )
 
     actors_service.start()
@@ -116,6 +109,7 @@ async def test_simple_start_stop(
     dispatch = test_env.generator.generate_dispatch()
     dispatch = replace(
         dispatch,
+        id=1,
         active=True,
         dry_run=False,
         duration=duration,
@@ -135,24 +129,24 @@ async def test_simple_start_stop(
     await asyncio.sleep(1)
     logging.info("Sent dispatch")
 
-    assert test_env.actor is not None
-    event = test_env.actor.initial_dispatch
+    event = test_env.actor(1).initial_dispatch
     assert event.options == {"test": True}
     assert event.components == dispatch.target
     assert event.dry_run is False
 
     logging.info("Received dispatch")
 
-    assert test_env.actor is not None
-    assert test_env.actor.is_running is True
+    assert test_env.actor(1).is_running is True
 
     fake_time.shift(duration)
     await test_env.running_status_sender.send(Dispatch(dispatch))
 
     # Give await actor.stop a chance to run
     await asyncio.sleep(1)
 
-    assert test_env.actor is None
+    # pylint: disable=protected-access
+    assert 1 not in test_env.actors_service._actors
+    # pylint: enable=protected-access
 
 
 def test_heapq_dispatch_compare(test_env: TestEnv) -> None:
@@ -209,6 +203,7 @@ async def test_dry_run(test_env: TestEnv, fake_time: time_machine.Coordinates) -
     dispatch = test_env.generator.generate_dispatch()
     dispatch = replace(
         dispatch,
+        id=1,
         dry_run=True,
         active=True,
         start_time=_now(),
@@ -224,20 +219,16 @@ async def test_dry_run(test_env: TestEnv, fake_time: time_machine.Coordinates) -
     fake_time.shift(timedelta(seconds=1))
     await asyncio.sleep(1)
 
-    assert test_env.actor is not None
-    event = test_env.actor.initial_dispatch
+    event = test_env.actor(1).initial_dispatch
 
     assert event.dry_run is dispatch.dry_run
     assert event.components == dispatch.target
     assert event.options == dispatch.payload
-    assert test_env.actor is not None
-    assert test_env.actor.is_running is True
+    assert test_env.actor(1).is_running is True
 
     assert dispatch.duration is not None
     fake_time.shift(dispatch.duration)
     await test_env.running_status_sender.send(Dispatch(dispatch))
 
     # Give await actor.stop a chance to run
     await asyncio.sleep(1)
-
-    assert test_env.actor is None