Replace asyncio gather() with wait()

This will be only useful in the future if
the actor decorator does no longer manage
the lifecycle of the actor tasks.

Logging the tasks done in this case might contain
useful debugging information as they include
cancellation information.

Signed-off-by: Daniel Zullo <[email protected]>

Author: daniel-zullo-frequenz

src/frequenz/sdk/actor/_run_utils.py

@@ -5,10 +5,13 @@
 
 
 import asyncio
+import logging
 from typing import Any
 
 from ._decorator import BaseActor
 
+_logger = logging.getLogger(__name__)
+
 
 async def run(*actors: Any) -> None:
     """Await the completion of all actors.
@@ -24,4 +27,27 @@ async def run(*actors: Any) -> None:
     for actor in actors:
         assert isinstance(actor, BaseActor), f"{actor} is not an instance of BaseActor"
 
-    await asyncio.gather(*(actor.join() for actor in actors))
+    pending_tasks = set()
+    for actor in actors:
+        pending_tasks.add(asyncio.create_task(actor.join(), name=str(actor)))
+
+    # Currently the actor decorator manages the life-cycle of the actor tasks
+    while pending_tasks:
+        done_tasks, pending_tasks = await asyncio.wait(
+            pending_tasks, return_when=asyncio.FIRST_COMPLETED
+        )
+
+        # This should always be only one task, but we handle many for extra safety
+        for task in done_tasks:
+            # Cancellation needs to be checked first, otherwise the other methods
+            # could raise a CancelledError
+            if task.cancelled():
+                _logger.info("The actor %s was cancelled", task.get_name())
+            elif exception := task.exception():
+                _logger.error(
+                    "The actor %s was finished due to an uncaught exception",
+                    task.get_name(),
+                    exc_info=exception,
+                )
+            else:
+                _logger.info("The actor %s finished normally", task.get_name())