Use a PersistentTaskGroup to implement ServiceBase

The service base class now delegates all sub-task management to a
`PersistentTaskGroup`, and has a special task called `main()` that
drives the service.

Writing single-task services should be much easier now, as only a
`main()` method needs to be implemented.

For more complex, multi-task, services, the internal task group can be
used to monitor sub-tasks, for example by using
`task_group.as_completed()`.

Signed-off-by: Leandro Lucarella <[email protected]>

Author: llucax

src/frequenz/core/asyncio/_service.py

@@ -8,16 +8,14 @@
 import asyncio
 import collections.abc
 import contextvars
-import logging
 from types import TracebackType
 from typing import Any, Self
 
 from typing_extensions import override
 
+from ._task_group import PersistentTaskGroup
 from ._util import TaskCreator, TaskReturnT
 
-_logger = logging.getLogger(__name__)
-
 
 class Service(abc.ABC):
     """A service running in the background.
@@ -62,14 +60,11 @@ def unique_id(self) -> str:
     @property
     @abc.abstractmethod
     def is_running(self) -> bool:
-        """Whether this service is running.
-
-        A service is considered running when at least one task is running.
-        """
+        """Whether this service is running."""
 
     @abc.abstractmethod
     def cancel(self, msg: str | None = None) -> None:
-        """Cancel all running tasks spawned by this service.
+        """Cancel this service.
 
         Args:
             msg: The message to be passed to the tasks being cancelled.
@@ -79,8 +74,7 @@ def cancel(self, msg: str | None = None) -> None:
     async def stop(self, msg: str | None = None) -> None:  # noqa: DOC502
         """Stop this service.
 
-        This method cancels all running tasks spawned by this service and waits for them
-        to finish.
+        This method cancels the service and waits for it to finish.
 
         Args:
             msg: The message to be passed to the tasks being cancelled.
@@ -149,22 +143,19 @@ class ServiceBase(Service, abc.ABC):
     [`stop()`][frequenz.core.asyncio.ServiceBase.stop] method, as the base
     implementation does not collect any results and re-raises all exceptions.
 
-    Example:
+    Example: Simple single-task example
         ```python
         import datetime
         import asyncio
+        from typing_extensions import override
 
         class Clock(ServiceBase):
             def __init__(self, resolution_s: float, *, unique_id: str | None = None) -> None:
                 super().__init__(unique_id=unique_id)
                 self._resolution_s = resolution_s
 
-            def start(self) -> None:
-                # Managed tasks are automatically saved, so there is no need to hold a
-                # reference to them if you don't need to further interact with them.
-                self.create_task(self._tick())
-
-            async def _tick(self) -> None:
+            @override
+            async def main(self) -> None:
                 while True:
                     await asyncio.sleep(self._resolution_s)
                     print(datetime.datetime.now())
@@ -182,6 +173,49 @@ async def main() -> None:
 
         asyncio.run(main())
         ```
+
+    Example: Multi-tasks example
+        ```python
+        import asyncio
+        import datetime
+        from typing_extensions import override
+
+        class MultiTaskService(ServiceBase):
+
+            async def _print_every(self, *, seconds: float) -> None:
+                while True:
+                    await asyncio.sleep(seconds)
+                    print(datetime.datetime.now())
+
+            async def _fail_after(self, *, seconds: float) -> None:
+                await asyncio.sleep(seconds)
+                raise ValueError("I failed")
+
+            @override
+            async def main(self) -> None:
+                self.create_task(self._print_every(seconds=1), name="print_1")
+                self.create_task(self._print_every(seconds=11), name="print_11")
+                failing = self.create_task(self._fail_after(seconds=5), name=f"fail_5")
+
+                async for task in self.task_group.as_completed():
+                    assert task.done()  # For demonstration purposes only
+                    try:
+                        task.result()
+                    except ValueError as error:
+                        if failing == task:
+                            failing = self.create_task(
+                                self._fail_after(seconds=5), name=f"fail_5"
+                            )
+                        else:
+                            raise
+
+        async def main() -> None:
+            async with MultiTaskService():
+                await asyncio.sleep(11)
+
+        asyncio.run(main())
+        ```
+
     """
 
     def __init__(
@@ -201,13 +235,10 @@ def __init__(
         # [2:] is used to remove the '0x' prefix from the hex representation of the id,
         # as it doesn't add any uniqueness to the string.
         self._unique_id: str = hex(id(self))[2:] if unique_id is None else unique_id
-        self._tasks: set[asyncio.Task[Any]] = set()
-        self._task_creator: TaskCreator = task_creator
-
-    @override
-    @abc.abstractmethod
-    def start(self) -> None:
-        """Start this service."""
+        self._main_task: asyncio.Task[None] | None = None
+        self._task_group: PersistentTaskGroup = PersistentTaskGroup(
+            unique_id=self._unique_id, task_creator=task_creator
+        )
 
     @property
     @override
@@ -216,9 +247,22 @@ def unique_id(self) -> str:
         return self._unique_id
 
     @property
-    def tasks(self) -> collections.abc.Set[asyncio.Task[Any]]:
-        """The set of running tasks spawned by this service."""
-        return self._tasks
+    def task_group(self) -> PersistentTaskGroup:
+        """The task group managing the tasks of this service."""
+        return self._task_group
+
+    @abc.abstractmethod
+    async def main(self) -> None:
+        """Execute the service logic."""
+
+    @override
+    def start(self) -> None:
+        """Start this service."""
+        if self.is_running:
+            return
+        self._main_task = self._task_group.task_creator.create_task(
+            self.main(), name=str(self)
+        )
 
     @property
     @override
@@ -227,7 +271,7 @@ def is_running(self) -> bool:
 
         A service is considered running when at least one task is running.
         """
-        return any(not task.done() for task in self._tasks)
+        return self._main_task is not None and not self._main_task.done()
 
     def create_task(
         self,
@@ -242,8 +286,8 @@ def create_task(
         A reference to the task will be held by the service, so there is no need to save
         the task object.
 
-        Tasks can be retrieved via the
-        [`tasks`][frequenz.core.asyncio.ServiceBase.tasks] property.
+        Tasks are created using the
+        [`task_group`][frequenz.core.asyncio.ServiceBase.task_group].
 
         Managed tasks always have a `name` including information about the service
         itself. If you need to retrieve the final name of the task you can always do so
@@ -268,24 +312,9 @@ def create_task(
         """
         if not name:
             name = hex(id(coro))[2:]
-        task = self._task_creator.create_task(
-            coro, name=f"{self}:{name}", context=context
+        return self._task_group.create_task(
+            coro, name=f"{self}:{name}", context=context, log_exception=log_exception
         )
-        self._tasks.add(task)
-        task.add_done_callback(self._tasks.discard)
-
-        if log_exception:
-
-            def _log_exception(task: asyncio.Task[TaskReturnT]) -> None:
-                try:
-                    task.result()
-                except asyncio.CancelledError:
-                    pass
-                except BaseException:  # pylint: disable=broad-except
-                    _logger.exception("%s: Task %r raised an exception", self, task)
-
-            task.add_done_callback(_log_exception)
-        return task
 
     @override
     def cancel(self, msg: str | None = None) -> None:
@@ -294,8 +323,9 @@ def cancel(self, msg: str | None = None) -> None:
         Args:
             msg: The message to be passed to the tasks being cancelled.
         """
-        for task in self._tasks:
-            task.cancel(msg)
+        if self._main_task is not None:
+            self._main_task.cancel(msg)
+        self._task_group.cancel(msg)
 
     @override
     async def stop(self, msg: str | None = None) -> None:
@@ -311,8 +341,6 @@ async def stop(self, msg: str | None = None) -> None:
             BaseExceptionGroup: If any of the tasks spawned by this service raised an
                 exception.
         """
-        if not self._tasks:
-            return
         self.cancel(msg)
         try:
             await self
@@ -369,28 +397,21 @@ async def _wait(self) -> None:
                 exception (`CancelError` is not considered an error and not returned in
                 the exception group).
         """
-        # We need to account for tasks that were created between when we started
-        # awaiting and we finished awaiting.
-        while self._tasks:
-            done, pending = await asyncio.wait(self._tasks)
-            assert not pending
-
-            # We remove the done tasks, but there might be new ones created after we
-            # started waiting.
-            self._tasks = self._tasks - done
-
-            exceptions: list[BaseException] = []
-            for task in done:
-                try:
-                    # This will raise a CancelledError if the task was cancelled or any
-                    # other exception if the task raised one.
-                    _ = task.result()
-                except BaseException as error:  # pylint: disable=broad-except
-                    exceptions.append(error)
-            if exceptions:
-                raise BaseExceptionGroup(
-                    f"Error while stopping service {self}", exceptions
-                )
+        exceptions: list[BaseException] = []
+
+        if self._main_task is not None:
+            try:
+                await self._main_task
+            except BaseException as error:  # pylint: disable=broad-except
+                exceptions.append(error)
+
+        try:
+            await self._task_group
+        except BaseExceptionGroup as exc_group:
+            exceptions.append(exc_group)
+
+        if exceptions:
+            raise BaseExceptionGroup(f"Error while stopping {self}", exceptions)
 
     @override
     def __await__(self) -> collections.abc.Generator[None, None, None]:
@@ -416,7 +437,13 @@ def __repr__(self) -> str:
         Returns:
             A string representation of this instance.
         """
-        return f"{type(self).__name__}<{self._unique_id} tasks={self._tasks!r}>"
+        details = "main"
+        if not self.is_running:
+            details += " not"
+        details += " running"
+        if self._task_group.is_running:
+            details += f", {len(self._task_group.tasks)} extra tasks"
+        return f"{type(self).__name__}<{self._unique_id} {details}>"
 
     def __str__(self) -> str:
         """Return a string representation of this instance.

tests/asyncio/test_service.py

@@ -8,6 +8,7 @@
 
 import async_solipsism
 import pytest
+from typing_extensions import override
 
 from frequenz.core.asyncio import ServiceBase
 
@@ -34,33 +35,32 @@ def __init__(
         self._sleep = sleep
         self._exc = exc
 
-    def start(self) -> None:
-        """Start this service."""
-
-        async def nop() -> None:
-            if self._sleep is not None:
-                await asyncio.sleep(self._sleep)
-            if self._exc is not None:
-                raise self._exc
-
-        self._tasks.add(asyncio.create_task(nop(), name="nop"))
+    @override
+    async def main(self) -> None:
+        """Run this service."""
+        if self._sleep is not None:
+            await asyncio.sleep(self._sleep)
+        if self._exc is not None:
+            raise self._exc
 
 
 async def test_construction_defaults() -> None:
     """Test the construction of a service with default arguments."""
     fake_service = FakeService()
     assert fake_service.unique_id == hex(id(fake_service))[2:]
-    assert fake_service.tasks == set()
+    assert fake_service.task_group.tasks == set()
     assert fake_service.is_running is False
     assert str(fake_service) == f"FakeService:{fake_service.unique_id}"
-    assert repr(fake_service) == f"FakeService<{fake_service.unique_id} tasks=set()>"
+    assert (
+        repr(fake_service) == f"FakeService<{fake_service.unique_id} main not running>"
+    )
 
 
 async def test_construction_custom() -> None:
     """Test the construction of a service with a custom unique ID."""
     fake_service = FakeService(unique_id="test")
     assert fake_service.unique_id == "test"
-    assert fake_service.tasks == set()
+    assert fake_service.task_group.tasks == set()
     assert fake_service.is_running is False