From 64b72ee72763f717ca7917a5f3f2d8792d133d01 Mon Sep 17 00:00:00 2001 From: Robert Gildein Date: Wed, 29 May 2024 14:40:48 +0200 Subject: [PATCH] Use gather to improve performance of wait_for_idle (#421) We found that the wait_for_idle execution time grows proportionally to the number of applications, so we used asyncio.gather. see [1] fixes: #419 --- [1]: https://github.com/juju/python-libjuju/issues/1055 --- cou/utils/juju_utils.py | 25 +++++++++++------- tests/unit/utils/test_juju_utils.py | 41 ++++++++++++++++++++--------- 2 files changed, 45 insertions(+), 21 deletions(-) diff --git a/cou/utils/juju_utils.py b/cou/utils/juju_utils.py index 6986980e..306af5da 100644 --- a/cou/utils/juju_utils.py +++ b/cou/utils/juju_utils.py @@ -605,6 +605,8 @@ async def wait_for_active_idle( WaitForApplicationsTimeout, defaults to False. :type raise_on_blocked: bool """ + if apps is None: + apps = await self._get_supported_apps() @retry(timeout=timeout, no_retry_exceptions=(WaitForApplicationsTimeout,)) @wraps(self.wait_for_active_idle) @@ -612,12 +614,20 @@ async def _wait_for_active_idle() -> None: # NOTE(rgildein): Defining wrapper so we can use retry with proper timeout model = await self._get_model() try: - await model.wait_for_idle( - apps=apps, - timeout=timeout, - idle_period=idle_period, - raise_on_blocked=raise_on_blocked, - status="active", + # 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, + status="active", + ) + for app in apps + ) ) except (asyncio.exceptions.TimeoutError, JujuAppError, JujuUnitError) as error: # NOTE(rgildein): Catching TimeoutError raised as exception when wait_for_idle @@ -631,7 +641,4 @@ async def _wait_for_active_idle() -> None: msg = str(error).replace("\n", "\n ", 1) raise WaitForApplicationsTimeout(msg) from error - if apps is None: - apps = await self._get_supported_apps() - await _wait_for_active_idle() diff --git a/tests/unit/utils/test_juju_utils.py b/tests/unit/utils/test_juju_utils.py index 518b5fef..54aafad6 100644 --- a/tests/unit/utils/test_juju_utils.py +++ b/tests/unit/utils/test_juju_utils.py @@ -46,6 +46,7 @@ def mocked_model(mocker): model.connection.return_value.is_open = True # simulate already connected model model.disconnect = AsyncMock() model.connect = AsyncMock() + model.wait_for_idle = AsyncMock() yield model @@ -506,12 +507,23 @@ async def test_coumodel_wait_for_active_idle(mock_get_supported_apps, mocked_mod await model.wait_for_active_idle(timeout) - mocked_model.wait_for_idle.assert_awaited_once_with( - apps=["app1", "app2"], - timeout=timeout, - idle_period=juju_utils.DEFAULT_MODEL_IDLE_PERIOD, - raise_on_blocked=False, - status="active", + mocked_model.wait_for_idle.assert_has_awaits( + [ + call( + apps=["app1"], + timeout=timeout, + idle_period=juju_utils.DEFAULT_MODEL_IDLE_PERIOD, + raise_on_blocked=False, + status="active", + ), + call( + apps=["app2"], + timeout=timeout, + idle_period=juju_utils.DEFAULT_MODEL_IDLE_PERIOD, + raise_on_blocked=False, + status="active", + ), + ] ) mock_get_supported_apps.assert_awaited_once_with() @@ -547,12 +559,17 @@ async def test_coumodel_wait_for_active_idle_timeout(mock_get_supported_apps, mo with pytest.raises(WaitForApplicationsTimeout): await model.wait_for_active_idle(timeout, apps=exp_apps) - mocked_model.wait_for_idle.assert_awaited_once_with( - apps=exp_apps, - timeout=timeout, - idle_period=juju_utils.DEFAULT_MODEL_IDLE_PERIOD, - raise_on_blocked=False, - status="active", + mocked_model.wait_for_idle.assert_has_awaits( + [ + call( + apps=[app], + timeout=timeout, + idle_period=juju_utils.DEFAULT_MODEL_IDLE_PERIOD, + raise_on_blocked=False, + status="active", + ) + for app in exp_apps + ] ) mock_get_supported_apps.assert_not_awaited()