Handle CancelledErrors properly (#1027)

The method `exception()` in `Task` doesn't return `CancelledError` if
the task was cancelled, instead it raises the error, so if a task is
cancelled an exception would be raised unintentionally.

Instead we want to catch `CancelledError` and treat it like any other
task error for these tasks, are they are not supposed to be cancelled.

This PR also adds some improvements to the resampling actor, to avoid
unnecessary restarts.

Fixes #1026.

Author: llucax

RELEASE_NOTES.md

@@ -41,3 +41,4 @@
 - Re-expose `ComponentMetricId` to the docs.
 - Fixed typing ambiguities when building composite formulas on streaming data.
 - Fixed a bug that was causing the `PowerDistributor` to exit if power requests to PV inverters or EV chargers timeout.
+- Fix handling of cancelled tasks in the data sourcing and resampling actor.

src/frequenz/sdk/actor/_data_sourcing/microgrid_api_source.py

@@ -413,11 +413,13 @@ async def clean_tasks(
             ) -> set[asyncio.Task[None]]:
                 done, pending = await asyncio.wait(sending_tasks, timeout=0)
                 for task in done:
-                    if error := task.exception():
-                        _logger.error(
+                    try:
+                        task.result()
+                    # pylint: disable-next=broad-except
+                    except (asyncio.CancelledError, Exception):
+                        _logger.exception(
                             "Error while processing message in task %s",
                             task.get_name(),
-                            exc_info=error,
                         )
                 return pending
 

src/frequenz/sdk/actor/_resampling.py

@@ -106,58 +106,36 @@ async def _run(self) -> None:
 
         - One task to process incoming subscription requests to resample new metrics.
         - One task to run the resampler.
-
-        Raises:
-            RuntimeError: If there is some unexpected error while resampling or
-                handling requests.
         """
         tasks_to_cancel: set[asyncio.Task[None]] = set()
-        try:
-            subscriptions_task = asyncio.create_task(
-                self._process_resampling_requests()
-            )
-            tasks_to_cancel.add(subscriptions_task)
+        subscriptions_task: asyncio.Task[None] | None = None
+        resampling_task: asyncio.Task[None] | None = None
 
+        try:
             while True:
-                resampling_task = asyncio.create_task(self._resampler.resample())
-                tasks_to_cancel.add(resampling_task)
+                if subscriptions_task is None or subscriptions_task.done():
+                    subscriptions_task = asyncio.create_task(
+                        self._process_resampling_requests()
+                    )
+                    tasks_to_cancel.add(subscriptions_task)
+
+                if resampling_task is None or resampling_task.done():
+                    resampling_task = asyncio.create_task(self._resampler.resample())
+                    tasks_to_cancel.add(resampling_task)
+
                 done, _ = await asyncio.wait(
                     [resampling_task, subscriptions_task],
                     return_when=asyncio.FIRST_COMPLETED,
                 )
 
                 if subscriptions_task in done:
                     tasks_to_cancel.remove(subscriptions_task)
-                    raise RuntimeError(
-                        "There was a problem with the subscriptions channel."
-                    )
+                    self._log_subscriptions_task_error(subscriptions_task)
 
                 if resampling_task in done:
                     tasks_to_cancel.remove(resampling_task)
-                    # The resampler shouldn't end without an exception
-                    error = resampling_task.exception()
-                    assert (
-                        error is not None
-                    ), "The resample() function shouldn't exit normally."
-
-                    # We don't know what to do with something other than
-                    # ResamplingError, so propagate the exception if that is the
-                    # case.
-                    if not isinstance(error, ResamplingError):
-                        raise error
-                    for source, source_error in error.exceptions.items():
-                        _logger.error(
-                            "Error resampling source %s, removing source...", source
-                        )
-                        removed = self._resampler.remove_timeseries(source)
-                        if not removed:
-                            _logger.warning(
-                                "Got an exception from an unknown source: "
-                                "source=%r, exception=%r",
-                                source,
-                                source_error,
-                            )
-                    # The resampling_task will be re-created if we reached this point
+                    self._log_resampling_task_error(resampling_task)
+
         finally:
             await asyncio.gather(*[cancel_and_await(t) for t in tasks_to_cancel])
 
@@ -171,3 +149,55 @@ async def _run(self) -> None:
             # state would be mostly the same as not really leaving the run()
             # method and just swallow any exception, which doesn't look super
             # smart.
+
+    def _log_subscriptions_task_error(
+        self, subscriptions_task: asyncio.Task[None]
+    ) -> None:
+        """Log an error from a stopped subscriptions task.
+
+        Args:
+            subscriptions_task: The subscriptions task.
+        """
+        try:
+            subscriptions_task.result()
+        # pylint: disable-next=broad-except
+        except (Exception, asyncio.CancelledError):
+            _logger.exception(
+                "The subscriptions task ended with an exception, restarting..."
+            )
+        else:
+            _logger.error("The subscriptions task ended unexpectedly, restarting...")
+
+    def _log_resampling_task_error(self, resampling_task: asyncio.Task[None]) -> None:
+        """Log an error from a stopped resampling task.
+
+        Args:
+            resampling_task: The resampling task.
+        """
+        # The resampler shouldn't be cancelled or end without an exception
+        try:
+            resampling_task.result()
+        except ResamplingError as error:
+            for source, source_error in error.exceptions.items():
+                _logger.error("Error resampling source %s, removing source...", source)
+                removed = self._resampler.remove_timeseries(source)
+                if not removed:
+                    _logger.error(
+                        "Got an exception from an unknown source: "
+                        "source=%r, exception=%r",
+                        source,
+                        source_error,
+                    )
+        # pylint: disable-next=broad-except
+        except (Exception, asyncio.CancelledError):
+            # We don't know what to do with something other than
+            # ResamplingError, so we log it, restart, and hope for the best.
+            _logger.exception(
+                "The resample() function got an unexpected error, restarting..."
+            )
+        else:
+            # The resample function should not end normally, so we log it,
+            # restart, and hope for the best.
+            _logger.error(
+                "The resample() function ended without an exception, restarting..."
+            )

tests/timeseries/mock_microgrid.py

@@ -230,8 +230,12 @@ def start_mock_client(
         self.init_mock_client(initialize_cb)
 
         def _done_callback(task: asyncio.Task[None]) -> None:
-            if exc := task.exception():
-                raise SystemExit(f"Streaming task {task.get_name()!r} failed: {exc}")
+            try:
+                task.result()
+            except (asyncio.CancelledError, Exception) as exc:
+                raise SystemExit(
+                    f"Streaming task {task.get_name()!r} failed: {exc}"
+                ) from exc
 
         for component_id, coro in self._streaming_coros:
             task = asyncio.create_task(coro, name=f"component-id:{component_id}")