Context timer

Author: Marenz

src/frequenz/dispatch/_actor_dispatcher.py

@@ -12,8 +12,7 @@
 
 from frequenz.channels import Broadcast, Receiver, Sender, select
 from frequenz.channels.timer import SkipMissedAndDrift, Timer
-from frequenz.client.common.microgrid.components import ComponentCategory
-from frequenz.client.microgrid import ComponentId
+from frequenz.client.common.microgrid.components import ComponentCategory, ComponentId
 from frequenz.sdk.actor import Actor, BackgroundService
 
 from ._dispatch import Dispatch

src/frequenz/dispatch/_bg_service.py

@@ -10,6 +10,7 @@
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
+from contextlib import closing
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
 from heapq import heappop, heappush
@@ -20,7 +21,6 @@
 from frequenz.client.dispatch import DispatchApiClient
 from frequenz.client.dispatch.types import Event
 from frequenz.sdk.actor import BackgroundService
-from typing_extensions import override
 
 from ._dispatch import Dispatch
 from ._event import Created, Deleted, DispatchEvent, Updated
@@ -192,6 +192,10 @@ async def new_running_state_event_receiver(
         if not self._tasks:
             raise RuntimeError("Dispatch service not started")
 
+        (task,) = self._tasks
+        if task.done():
+            raise RuntimeError("Dispatch service not running")
+
         # Find all matching dispatches based on the type and collect them
         dispatches = [
             dispatch for dispatch in self._dispatches.values() if dispatch.type == type
@@ -220,15 +224,8 @@ async def new_running_state_event_receiver(
 
     # pylint: enable=redefined-builtin
 
-    @override
-    async def stop(self, msg: str | None = None) -> None:
-        """Stop the background service."""
-        self._next_event_timer.close()
-        await super().stop(msg)
-
     def start(self) -> None:
         """Start the background service."""
-        self._next_event_timer.reset(interval=timedelta(seconds=1))
         self._tasks.add(asyncio.create_task(self._run()))
 
     async def _run(self) -> None:
@@ -244,32 +241,43 @@ async def _run(self) -> None:
         stream = self._client.stream(microgrid_id=self._microgrid_id)
 
         # Streaming updates
-        async for selected in select(self._next_event_timer, stream):
-            if selected_from(selected, self._next_event_timer):
-                if not self._scheduled_events:
-                    continue
-                await self._execute_scheduled_event(
-                    heappop(self._scheduled_events).dispatch
-                )
-            elif selected_from(selected, stream):
-                _logger.debug("Received dispatch event: %s", selected.message)
-                dispatch = Dispatch(selected.message.dispatch)
-                match selected.message.event:
-                    case Event.CREATED:
-                        self._dispatches[dispatch.id] = dispatch
-                        await self._update_dispatch_schedule_and_notify(dispatch, None)
-                        await self._lifecycle_events_tx.send(Created(dispatch=dispatch))
-                    case Event.UPDATED:
-                        await self._update_dispatch_schedule_and_notify(
-                            dispatch, self._dispatches[dispatch.id]
-                        )
-                        self._dispatches[dispatch.id] = dispatch
-                        await self._lifecycle_events_tx.send(Updated(dispatch=dispatch))
-                    case Event.DELETED:
-                        self._dispatches.pop(dispatch.id)
-                        await self._update_dispatch_schedule_and_notify(None, dispatch)
-
-                        await self._lifecycle_events_tx.send(Deleted(dispatch=dispatch))
+        with closing(self._next_event_timer) as next_event_timer:
+            async for selected in select(next_event_timer, stream):
+                if selected_from(selected, next_event_timer):
+                    if not self._scheduled_events:
+                        continue
+                    await self._execute_scheduled_event(
+                        heappop(self._scheduled_events).dispatch
+                    )
+                elif selected_from(selected, stream):
+                    _logger.debug("Received dispatch event: %s", selected.message)
+                    dispatch = Dispatch(selected.message.dispatch)
+                    match selected.message.event:
+                        case Event.CREATED:
+                            self._dispatches[dispatch.id] = dispatch
+                            await self._update_dispatch_schedule_and_notify(
+                                dispatch, None
+                            )
+                            await self._lifecycle_events_tx.send(
+                                Created(dispatch=dispatch)
+                            )
+                        case Event.UPDATED:
+                            await self._update_dispatch_schedule_and_notify(
+                                dispatch, self._dispatches[dispatch.id]
+                            )
+                            self._dispatches[dispatch.id] = dispatch
+                            await self._lifecycle_events_tx.send(
+                                Updated(dispatch=dispatch)
+                            )
+                        case Event.DELETED:
+                            self._dispatches.pop(dispatch.id)
+                            await self._update_dispatch_schedule_and_notify(
+                                None, dispatch
+                            )
+
+                            await self._lifecycle_events_tx.send(
+                                Deleted(dispatch=dispatch)
+                            )
 
     async def _execute_scheduled_event(self, dispatch: Dispatch) -> None:
         """Execute a scheduled event.

tests/test_frequenz_dispatch.py

@@ -461,6 +461,9 @@ async def test_dispatch_new_but_finished(
     test_env.client.set_dispatches(test_env.microgrid_id, [finished_dispatch])
     await test_env.service.stop()
     test_env.service.start()
+    fake_time.shift(timedelta(seconds=1))
+    await asyncio.sleep(1)
+
     test_env = replace(
         test_env,
         lifecycle_events=test_env.service.new_lifecycle_events_receiver("TEST_TYPE"),
@@ -470,8 +473,8 @@ async def test_dispatch_new_but_finished(
             )
         ),
     )
+    await asyncio.sleep(1)
 
-    fake_time.shift(timedelta(seconds=1))
     # Process the lifecycle event caused by the old dispatch at startup
     await test_env.lifecycle_events.receive()