From e272aeacac9c7c94d4d89981785b0cc9dd1af134 Mon Sep 17 00:00:00 2001
From: james_lin <james.lin@canonical.com>
Date: Mon, 7 Oct 2024 13:23:57 +0800
Subject: [PATCH] perf: Limit the wait_for_idle concurrency task number (#573)

- Limit the wait_for_idle concurrency task number to 5 to avoid sending
    too much request to controller API
- Now the retry is base on single application, this somehow reduce a lot
    of time to unnecessary repeat checks on those active application.
---
 cou/utils/juju_utils.py | 33 +++++++++++++++------------------
 1 file changed, 15 insertions(+), 18 deletions(-)

diff --git a/cou/utils/juju_utils.py b/cou/utils/juju_utils.py
index 6cf4bea1..8833e1e2 100644
--- a/cou/utils/juju_utils.py
+++ b/cou/utils/juju_utils.py
@@ -680,30 +680,26 @@ async def wait_for_idle(  # pylint: disable=too-many-arguments,too-many-position
         """
         if apps is None:
             apps = await self._get_supported_apps()
+        semaphore = asyncio.Semaphore(5)
 
         @retry(timeout=timeout, no_retry_exceptions=(WaitForApplicationsTimeout,))
         @wraps(self.wait_for_idle)
-        async def _wait_for_idle() -> None:
+        async def _wait_for_idle(app: str) -> None:
             # NOTE(rgildein): Defining wrapper so we can use retry with proper timeout
             model = await self._get_model()
             try:
-                # NOTE(rgildein): Use asyncio.gather because we always go through the application
-                # list (apps will never be None) and libjuju is very slow here.
-                # https://github.com/juju/python-libjuju/issues/1055
-                await asyncio.gather(
-                    *(
-                        model.wait_for_idle(
-                            apps=[app],
-                            timeout=timeout,
-                            idle_period=idle_period,
-                            raise_on_blocked=raise_on_blocked,
-                            raise_on_error=raise_on_error,
-                            status=status,
-                            wait_for_at_least_units=0,  # don't hang if an app has no units
-                        )
-                        for app in apps
+                # NOTE(jneo8): Use semaphore to limit the number of tasks to avoid sending too much
+                # request to API server at the same time.
+                async with semaphore:
+                    await model.wait_for_idle(
+                        apps=[app],
+                        timeout=timeout,
+                        idle_period=idle_period,
+                        raise_on_blocked=raise_on_blocked,
+                        raise_on_error=raise_on_error,
+                        status=status,
+                        wait_for_at_least_units=0,  # don't hang if an app has no units
                     )
-                )
             except (asyncio.exceptions.TimeoutError, JujuAppError, JujuUnitError) as error:
                 # NOTE(rgildein): Catching TimeoutError raised as exception when wait_for_idle
                 # reached timeout. Also adding two spaces to make it more user friendly.
@@ -716,7 +712,8 @@ async def _wait_for_idle() -> None:
                 msg = str(error).replace("\n", "\n  ", 1)
                 raise WaitForApplicationsTimeout(msg) from error
 
-        await _wait_for_idle()
+        tasks = [_wait_for_idle(app=app) for app in apps]
+        await asyncio.gather(*tasks)
 
     async def resolve_all(self) -> None:
         """Resolve all the units in the model if they are in error status."""