Re-create actors on dispatch updates

This commit modifies the `DispatchManagingActor` to re-create the actor
when a new dispatch is received and the actor should start running
instead of just start/stop the actor. To create the actor a factory
function is used, which passes the initial dispatch information to the
actor, so it can be properly initialized instead of having the
initialization in 2 steps, the creation and the receiving of the dispatch
update.

Author: llucax

src/frequenz/dispatch/_managing_actor.py

@@ -3,13 +3,15 @@
 
 """Helper class to manage actors based on dispatches."""
 
+import asyncio
 import logging
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Set
+from typing import Any
 
-from frequenz.channels import Receiver, Sender
+from frequenz.channels import Broadcast, Receiver
 from frequenz.client.dispatch.types import TargetComponents
-from frequenz.sdk.actor import Actor
+from frequenz.sdk.actor import Actor, BackgroundService
 
 from ._dispatch import Dispatch
 
@@ -30,38 +32,71 @@ class DispatchUpdate:
     """Additional options."""
 
 
-class DispatchManagingActor(Actor):
+class DispatchManagingActor(BackgroundService):
     """Helper class to manage actors based on dispatches.
 
     Example usage:
 
     ```python
     import os
     import asyncio
+    from typing import override
     from frequenz.dispatch import Dispatcher, DispatchManagingActor, DispatchUpdate
     from frequenz.client.dispatch.types import TargetComponents
     from frequenz.client.common.microgrid.components import ComponentCategory
-
-    from frequenz.channels import Receiver, Broadcast
+    from frequenz.channels import Receiver, Broadcast, select, selected_from
+    from frequenz.sdk.actor import Actor, run
 
     class MyActor(Actor):
-        def __init__(self, updates_channel: Receiver[DispatchUpdate]):
-            super().__init__()
-            self._updates_channel = updates_channel
-            self._dry_run: bool
-            self._options : dict[str, Any]
-
+        def __init__(
+                self,
+                *,
+                name: str | None = None,
+        ) -> None:
+            super().__init__(name=name)
+            self._dispatch_updates_receiver: Receiver[DispatchUpdate] | None = None
+            self._dry_run: bool = False
+            self._options: dict[str, Any] = {}
+
+        @classmethod
+        def new_with_dispatch(
+                cls,
+                initial_dispatch: DispatchUpdate,
+                dispatch_updates_receiver: Receiver[DispatchUpdate],
+                *,
+                name: str | None = None,
+        ) -> "Self":
+            self = cls(name=name)
+            self._dispatch_updates_receiver = dispatch_updates_receiver
+            self._update_dispatch_information(initial_dispatch)
+            return self
+
+        @override
         async def _run(self) -> None:
-            while True:
-                update = await self._updates_channel.receive()
-                print("Received update:", update)
-
-                self.set_components(update.components)
-                self._dry_run = update.dry_run
-                self._options = update.options
+            other_recv: Receiver[Any] = ...
 
-        def set_components(self, components: TargetComponents) -> None:
-            match components:
+            if self._dispatch_updates_receiver is None:
+                async for msg in other_recv:
+                    # do stuff
+                    ...
+            else:
+                await self._run_with_dispatch(other_recv)
+
+        async def _run_with_dispatch(self, other_recv: Receiver[Any]) -> None:
+            async for selected in select(self._dispatch_updates_receiver, other_recv):
+                if selected_from(selected, self._dispatch_updates_receiver):
+                    self._update_dispatch_information(selected.message)
+                elif selected_from(selected, other_recv):
+                    # do stuff
+                    ...
+                else:
+                    assert False, f"Unexpected selected receiver: {selected}"
+
+        def _update_dispatch_information(self, dispatch_update: DispatchUpdate) -> None:
+            print("Received update:", dispatch_update)
+            self._dry_run = dispatch_update.dry_run
+            self._options = dispatch_update.options
+            match dispatch_update.components:
                 case []:
                     print("Dispatch: Using all components")
                 case list() as ids if isinstance(ids[0], int):
@@ -75,7 +110,7 @@ def set_components(self, components: TargetComponents) -> None:
                         unsupported,
                     )
 
-    async def run():
+    async def main():
         url = os.getenv("DISPATCH_API_URL", "grpc://fz-0004.frequenz.io:50051")
         key  = os.getenv("DISPATCH_API_KEY", "some-key")
 
@@ -86,64 +121,92 @@ async def run():
             server_url=url,
             key=key
         )
-
-        # Create update channel to receive dispatch update events pre-start and mid-run
-        dispatch_updates_channel = Broadcast[DispatchUpdate](name="dispatch_updates_channel")
-
-        # Start actor and give it an dispatch updates channel receiver
-        my_actor = MyActor(dispatch_updates_channel.new_receiver())
+        dispatcher.start()
 
         status_receiver = dispatcher.new_running_state_event_receiver("EXAMPLE_TYPE")
 
+        def actor_factory(initial_dispatch, dispatch_receiver):
+            return MyActor.new_with_dispatch(initial_dispatch, dispatch_receiver)
+
         managing_actor = DispatchManagingActor(
-            actor=my_actor,
+            actor_factory=actor_factory,
             running_status_receiver=status_receiver,
-            updates_sender=dispatch_updates_channel.new_sender(),
         )
 
-        await asyncio.gather(dispatcher.start(), managing_actor.start())
+        await run(managing_actor)
     ```
     """
 
     def __init__(
         self,
-        actor: Actor | Set[Actor],
+        actor_factory: Callable[[DispatchUpdate, Receiver[DispatchUpdate]], Actor],
         running_status_receiver: Receiver[Dispatch],
-        updates_sender: Sender[DispatchUpdate] | None = None,
     ) -> None:
         """Initialize the dispatch handler.
 
         Args:
-            actor: A set of actors or a single actor to manage.
+            actor_factory: A callable that creates an actor with some initial dispatch
+                information.
             running_status_receiver: The receiver for dispatch running status changes.
-            updates_sender: The sender for dispatch events
         """
         super().__init__()
         self._dispatch_rx = running_status_receiver
-        self._actors: frozenset[Actor] = frozenset(
-            [actor] if isinstance(actor, Actor) else actor
+        self._actor_factory = actor_factory
+        self._actor: Actor | None = None
+        self._updates_channel = Broadcast[DispatchUpdate](
+            name="dispatch_updates_channel", resend_latest=True
         )
-        self._updates_sender = updates_sender
+        self._updates_sender = self._updates_channel.new_sender()
+
+    def start(self) -> None:
+        """Start the background service."""
+        self._tasks.add(asyncio.create_task(self._run()))
+
+    def new_receiver(self) -> Receiver[DispatchUpdate]:
+        """Create a new receiver for dispatch updates.
+
+        Returns:
+            A new receiver for dispatch updates.
+        """
+        return self._updates_channel.new_receiver()
 
-    def _start_actors(self) -> None:
+    async def _start_actor(self, dispatch: Dispatch) -> None:
         """Start all actors."""
-        for actor in self._actors:
-            if actor.is_running:
-                _logger.warning("Actor %s is already running", actor.name)
-            else:
-                actor.start()
+        dispatch_update = DispatchUpdate(
+            components=dispatch.target,
+            dry_run=dispatch.dry_run,
+            options=dispatch.payload,
+        )
+
+        if self._actor:
+            sent_str = ""
+            if self._updates_sender is not None:
+                sent_str = ", sent a dispatch update instead of creating a new actor"
+                await self._updates_sender.send(dispatch_update)
+            _logger.warning(
+                "Actor for dispatch type %r is already running%s",
+                dispatch.type,
+                sent_str,
+            )
+        else:
+            _logger.info("Starting actor for dispatch type %r", dispatch.type)
+            self._actor = self._actor_factory(dispatch_update, self.new_receiver())
+            self._actor.start()
 
-    async def _stop_actors(self, msg: str) -> None:
+    async def _stop_actor(self, stopping_dispatch: Dispatch, msg: str) -> None:
         """Stop all actors.
 
         Args:
+            stopping_dispatch: The dispatch that is stopping the actor.
             msg: The message to be passed to the actors being stopped.
         """
-        for actor in self._actors:
-            if actor.is_running:
-                await actor.stop(msg)
-            else:
-                _logger.warning("Actor %s is not running", actor.name)
+        if self._actor is None:
+            _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."""
@@ -157,18 +220,6 @@ async def _handle_dispatch(self, dispatch: Dispatch) -> None:
             dispatch: The dispatch to handle.
         """
         if dispatch.started:
-            if self._updates_sender is not None:
-                _logger.info("Updated by dispatch %s", dispatch.id)
-                await self._updates_sender.send(
-                    DispatchUpdate(
-                        components=dispatch.target,
-                        dry_run=dispatch.dry_run,
-                        options=dispatch.payload,
-                    )
-                )
-
-            _logger.info("Started by dispatch %s", dispatch.id)
-            self._start_actors()
+            await self._start_actor(dispatch)
         else:
-            _logger.info("Stopped by dispatch %s", dispatch.id)
-            await self._stop_actors("Dispatch stopped")
+            await self._stop_actor(dispatch, "Dispatch stopped")