Improve performance of async_get_integration_with_requirements (#110770)

* Improve performance of async_get_integration_with_requirements

- Migrate to the future pattern instead of using asyncio.Event
- Use sets in a few places to avoid linear searching
- Check the cache when processing deps so we do not
  create tasks to process requirements for deps that
  have already been processed

* name

* add concurrency test

* Update homeassistant/requirements.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Update homeassistant/requirements.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Update homeassistant/requirements.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Update homeassistant/requirements.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* reset_mock

---------

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
This commit is contained in:
J. Nick Koston 2024-02-17 12:26:41 -06:00 committed by GitHub
parent aa8d8402b4
commit 53944235d2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 171 additions and 18 deletions

View file

@ -129,7 +129,7 @@ class RequirementsManager:
self.hass = hass
self.pip_lock = asyncio.Lock()
self.integrations_with_reqs: dict[
str, Integration | asyncio.Event | None | UndefinedType
str, Integration | asyncio.Future[None] | None | UndefinedType
] = {}
self.install_failure_history: set[str] = set()
self.is_installed_cache: set[str] = set()
@ -155,31 +155,33 @@ class RequirementsManager:
return integration
cache = self.integrations_with_reqs
int_or_evt = cache.get(domain, UNDEFINED)
int_or_fut = cache.get(domain, UNDEFINED)
if isinstance(int_or_evt, asyncio.Event):
await int_or_evt.wait()
if isinstance(int_or_fut, asyncio.Future):
await int_or_fut
# When we have waited and it's UNDEFINED, it doesn't exist
# We don't cache that it doesn't exist, or else people can't fix it
# and then restart, because their config will never be valid.
if (int_or_evt := cache.get(domain, UNDEFINED)) is UNDEFINED:
if (int_or_fut := cache.get(domain, UNDEFINED)) is UNDEFINED:
raise IntegrationNotFound(domain)
if int_or_evt is not UNDEFINED:
return cast(Integration, int_or_evt)
if int_or_fut is not UNDEFINED:
return cast(Integration, int_or_fut)
event = cache[domain] = asyncio.Event()
event = cache[domain] = self.hass.loop.create_future()
try:
await self._async_process_integration(integration, done)
except Exception:
del cache[domain]
event.set()
if not event.done():
event.set_result(None)
raise
cache[domain] = integration
event.set()
if not event.done():
event.set_result(None)
return integration
async def _async_process_integration(
@ -191,19 +193,35 @@ class RequirementsManager:
integration.domain, integration.requirements
)
deps_to_check = [
cache = self.integrations_with_reqs
deps_to_check = {
dep
for dep in integration.dependencies + integration.after_dependencies
if dep not in done
]
# If the dep is in the cache and it's an Integration
# it's already been checked for the requirements and we should
# not check it again.
and (
not (cached_integration := cache.get(dep))
or type(cached_integration) is not Integration
)
}
for check_domain, to_check in DISCOVERY_INTEGRATIONS.items():
if (
check_domain not in done
and check_domain not in deps_to_check
# If the integration is in the cache and it's an Integration
# it's already been checked for the requirements and we should
# not check it again.
and (
not (cached_integration := cache.get(check_domain))
or type(cached_integration) is not Integration
)
and any(check in integration.manifest for check in to_check)
):
deps_to_check.append(check_domain)
deps_to_check.add(check_domain)
if not deps_to_check:
return
@ -233,11 +251,11 @@ class RequirementsManager:
if an requirement can't be satisfied.
"""
if self.hass.config.skip_pip_packages:
skipped_requirements = [
skipped_requirements = {
req
for req in requirements
if Requirement(req).name in self.hass.config.skip_pip_packages
]
}
for req in skipped_requirements:
_LOGGER.warning("Skipping requirement %s. This may cause issues", req)
@ -249,9 +267,8 @@ class RequirementsManager:
self._raise_for_failed_requirements(name, missing)
async with self.pip_lock:
# Recaculate missing again now that we have the lock
missing = self._find_missing_requirements(requirements)
if missing:
# Recalculate missing again now that we have the lock
if missing := self._find_missing_requirements(requirements):
await self._async_process_requirements(name, missing)
def _find_missing_requirements(self, requirements: list[str]) -> list[str]:

View file

@ -1,4 +1,5 @@
"""Test requirements module."""
import asyncio
import logging
import os
from unittest.mock import call, patch
@ -7,9 +8,11 @@ import pytest
from homeassistant import loader, setup
from homeassistant.core import HomeAssistant
from homeassistant.loader import async_get_integration
from homeassistant.requirements import (
CONSTRAINT_FILE,
RequirementsNotFound,
_async_get_manager,
async_clear_install_history,
async_get_integration_with_requirements,
async_process_requirements,
@ -156,6 +159,139 @@ async def test_get_integration_with_requirements(hass: HomeAssistant) -> None:
]
async def test_get_integration_with_requirements_cache(hass: HomeAssistant) -> None:
"""Check getting an integration with loaded requirements considers cache.
We want to make sure that we do not check requirements for dependencies
that we have already checked.
"""
hass.config.skip_pip = False
mock_integration(
hass, MockModule("test_component_dep", requirements=["test-comp-dep==1.0.0"])
)
mock_integration(
hass,
MockModule(
"test_component_after_dep", requirements=["test-comp-after-dep==1.0.0"]
),
)
mock_integration(
hass,
MockModule(
"test_component",
requirements=["test-comp==1.0.0"],
dependencies=["test_component_dep"],
partial_manifest={"after_dependencies": ["test_component_after_dep"]},
),
)
mock_integration(
hass,
MockModule(
"test_component2",
requirements=["test-comp2==1.0.0"],
dependencies=["test_component_dep"],
partial_manifest={"after_dependencies": ["test_component_after_dep"]},
),
)
with patch(
"homeassistant.util.package.is_installed", return_value=False
) as mock_is_installed, patch(
"homeassistant.util.package.install_package", return_value=True
) as mock_inst, patch(
"homeassistant.requirements.async_get_integration", wraps=async_get_integration
) as mock_async_get_integration:
integration = await async_get_integration_with_requirements(
hass, "test_component"
)
assert integration
assert integration.domain == "test_component"
assert len(mock_is_installed.mock_calls) == 3
assert sorted(
mock_call[1][0] for mock_call in mock_is_installed.mock_calls
) == [
"test-comp-after-dep==1.0.0",
"test-comp-dep==1.0.0",
"test-comp==1.0.0",
]
assert len(mock_inst.mock_calls) == 3
assert sorted(mock_call[1][0] for mock_call in mock_inst.mock_calls) == [
"test-comp-after-dep==1.0.0",
"test-comp-dep==1.0.0",
"test-comp==1.0.0",
]
# The dependent integrations should be fetched since
assert len(mock_async_get_integration.mock_calls) == 3
assert sorted(
mock_call[1][1] for mock_call in mock_async_get_integration.mock_calls
) == ["test_component", "test_component_after_dep", "test_component_dep"]
# test_component2 has the same deps as test_component and we should
# not check the requirements for the deps again
mock_is_installed.reset_mock()
mock_inst.reset_mock()
mock_async_get_integration.reset_mock()
integration = await async_get_integration_with_requirements(
hass, "test_component2"
)
assert integration
assert integration.domain == "test_component2"
assert len(mock_is_installed.mock_calls) == 1
assert sorted(mock_call[1][0] for mock_call in mock_is_installed.mock_calls) == [
"test-comp2==1.0.0",
]
assert len(mock_inst.mock_calls) == 1
assert sorted(mock_call[1][0] for mock_call in mock_inst.mock_calls) == [
"test-comp2==1.0.0",
]
# The dependent integrations should not be fetched again
assert len(mock_async_get_integration.mock_calls) == 1
assert sorted(
mock_call[1][1] for mock_call in mock_async_get_integration.mock_calls
) == [
"test_component2",
]
async def test_get_integration_with_requirements_concurrency(
hass: HomeAssistant,
) -> None:
"""Test that we don't install the same requirement concurrently."""
hass.config.skip_pip = False
mock_integration(
hass, MockModule("test_component_dep", requirements=["test-comp-dep==1.0.0"])
)
process_integration_calls = 0
async def _async_process_integration_slowed(*args, **kwargs):
nonlocal process_integration_calls
process_integration_calls += 1
await asyncio.sleep(0)
manager = _async_get_manager(hass)
with patch.object(
manager, "_async_process_integration", _async_process_integration_slowed
):
tasks = [
async_get_integration_with_requirements(hass, "test_component_dep")
for _ in range(10)
]
results = await asyncio.gather(*tasks)
assert all(result.domain == "test_component_dep" for result in results)
assert process_integration_calls == 1
async def test_get_integration_with_requirements_pip_install_fails_two_passes(
hass: HomeAssistant,
) -> None: