From b2e19d333b64d6855aca6f4ded535a6867ac0be5 Mon Sep 17 00:00:00 2001
From: Douglas Raillard <douglas.raillard@arm.com>
Date: Fri, 26 Apr 2024 13:41:40 +0100
Subject: [PATCH] utils/asyn: Factor out the calls to asyncio.run

Prepare for providing our own implementation of asyncio.run() to work
without nest_asyncio package.
---
 devlib/module/cgroups.py |  5 ++---
 devlib/utils/asyn.py     | 16 ++++++++++++----
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/devlib/module/cgroups.py b/devlib/module/cgroups.py
index ece52c278..a7edf879c 100644
--- a/devlib/module/cgroups.py
+++ b/devlib/module/cgroups.py
@@ -19,13 +19,12 @@
 from shlex import quote
 import itertools
 import warnings
-import asyncio
 
 from devlib.module import Module
 from devlib.exception import TargetStableError
 from devlib.utils.misc import list_to_ranges, isiterable
 from devlib.utils.types import boolean
-from devlib.utils.asyn import asyncf
+from devlib.utils.asyn import asyncf, run
 
 
 class Controller(object):
@@ -410,7 +409,7 @@ async def register_controller(ss):
                              controller.mount_point)
             self.controllers[ss.name] = controller
 
-        asyncio.run(
+        run(
             target.async_manager.map_concurrently(
                 register_controller,
                 subsys,
diff --git a/devlib/utils/asyn.py b/devlib/utils/asyn.py
index 771f3d5f0..b3fe8892b 100644
--- a/devlib/utils/asyn.py
+++ b/devlib/utils/asyn.py
@@ -292,6 +292,14 @@ def __set_name__(self, owner, name):
         self._name = name
 
 
+def run(coro):
+    """
+    Similar to :func:`asyncio.run` but can be called while an event loop is
+    running.
+    """
+    return asyncio.run(coro)
+
+
 def asyncf(f):
     """
     Decorator used to turn a coroutine into a blocking function, with an
@@ -328,14 +336,14 @@ def genf():
                     asyncgen = x.__aiter__()
                     while True:
                         try:
-                            yield asyncio.run(asyncgen.__anext__())
+                            yield run(asyncgen.__anext__())
                         except StopAsyncIteration:
                             return
 
                 return genf()
             else:
                 return await x
-        return asyncio.run(wrapper())
+        return run(wrapper())
 
     return _AsyncPolymorphicFunction(
         asyn=f,
@@ -358,10 +366,10 @@ def __aexit__(self, *args, **kwargs):
         return self.cm.__aexit__(*args, **kwargs)
 
     def __enter__(self, *args, **kwargs):
-        return asyncio.run(self.cm.__aenter__(*args, **kwargs))
+        return run(self.cm.__aenter__(*args, **kwargs))
 
     def __exit__(self, *args, **kwargs):
-        return asyncio.run(self.cm.__aexit__(*args, **kwargs))
+        return run(self.cm.__aexit__(*args, **kwargs))
 
 
 def asynccontextmanager(f):