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>

homeassistant/requirements.py

@@ -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):
+        if isinstance(int_or_fut, asyncio.Future):
-            await int_or_evt.wait()
+            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:
+        if int_or_fut is not UNDEFINED:
-            return cast(Integration, int_or_evt)
+            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
+            # Recalculate missing again now that we have the lock
-            missing = self._find_missing_requirements(requirements)
+            if missing := self._find_missing_requirements(requirements):
-            if missing:
                 await self._async_process_requirements(name, missing)
 
     def _find_missing_requirements(self, requirements: list[str]) -> list[str]:

tests/test_requirements.py

@@ -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: