Add stop method to FormulaEvaluator

Author: ela-kotulska-frequenz

src/frequenz/sdk/timeseries/formula_engine/_formula_engine.py

@@ -314,7 +314,8 @@ async def _run(self) -> None:
             try:
                 msg = await evaluator.apply()
             except asyncio.CancelledError:
-                _logger.debug("FormulaEngine task cancelled: %s", self._name)
+                _logger.debug("FormulaEngine has been cancelled, stopping evaluator")
+                await evaluator.stop()
                 raise
             except Exception as err:  # pylint: disable=broad-except
                 _logger.warning(

src/frequenz/sdk/timeseries/formula_engine/_formula_evaluator.py

@@ -4,13 +4,14 @@
 """A post-fix formula evaluator that operates on `Sample` receivers."""
 
 import asyncio
-from collections.abc import Callable
+from collections.abc import Callable, Awaitable
 from datetime import datetime
 from math import isinf, isnan
-from typing import Generic
+from typing import Generic, Any
 
 from .._base_types import QuantityT, Sample
 from ._formula_steps import FormulaStep, MetricFetcher
+from frequenz.sdk._internal._asyncio import cancel_and_await
 
 
 class FormulaEvaluator(Generic[QuantityT]):
@@ -38,6 +39,7 @@ def __init__(
         self._metric_fetchers: dict[str, MetricFetcher[QuantityT]] = metric_fetchers
         self._first_run = True
         self._create_method: Callable[[float], QuantityT] = create_method
+        self._pending_fetch_tasks: list[asyncio.Task[Any]] = []
 
     async def _synchronize_metric_timestamps(
         self, metrics: set[asyncio.Task[Sample[QuantityT] | None]]
@@ -87,6 +89,14 @@ async def _synchronize_metric_timestamps(
         self._first_run = False
         return latest_ts
 
+    async def stop(self) -> None:
+        tasks_to_cancel = [
+            asyncio.create_task(cancel_and_await(task))
+            for task in self._pending_fetch_tasks
+        ]
+        await asyncio.wait(tasks_to_cancel, return_when=asyncio.ALL_COMPLETED)
+        self._pending_fetch_tasks = []
+
     async def apply(self) -> Sample[QuantityT]:
         """Fetch the latest metrics, apply the formula once and return the result.
 
@@ -98,11 +108,13 @@ async def apply(self) -> Sample[QuantityT]:
                 failed.
         """
         eval_stack: list[float] = []
+        self._pending_fetch_tasks = [
+            asyncio.create_task(fetcher.fetch_next(), name=name)
+            for name, fetcher in self._metric_fetchers.items()
+        ]
+
         ready_metrics, pending = await asyncio.wait(
-            [
-                asyncio.create_task(fetcher.fetch_next(), name=name)
-                for name, fetcher in self._metric_fetchers.items()
-            ],
+            self._pending_fetch_tasks,
             return_when=asyncio.ALL_COMPLETED,
         )