Bootstrap to start registry loading early (#22321)

* Registries store directly in data on loading.

* Loading registries concurent with stage 1.

* Removed comments

homeassistant/bootstrap.py

@@ -1,4 +1,5 @@
 """Provide methods to bootstrap a Home Assistant instance."""
+import asyncio
 import logging
 import logging.handlers
 import os
@@ -157,6 +158,12 @@ async def async_from_config_dict(config: Dict[str, Any],
 
     await hass.async_block_till_done()
 
+    # Kick off loading the registries. They don't need to be awaited.
+    asyncio.gather(
+        hass.helpers.device_registry.async_get_registry(),
+        hass.helpers.entity_registry.async_get_registry(),
+        hass.helpers.area_registry.async_get_registry())
+
     # stage 1
     for component in components:
         if component in FIRST_INIT_COMPONENT:

homeassistant/helpers/area_registry.py

@@ -1,6 +1,7 @@
 """Provide a way to connect devices to one physical location."""
 import logging
 import uuid
+from asyncio import Event
 from collections import OrderedDict
 from typing import MutableMapping  # noqa: F401
 from typing import Iterable, Optional, cast
@@ -9,6 +10,7 @@ import attr
 
 from homeassistant.core import callback
 from homeassistant.loader import bind_hass
+
 from .typing import HomeAssistantType
 
 _LOGGER = logging.getLogger(__name__)
@@ -133,14 +135,21 @@ class AreaRegistry:
 @bind_hass
 async def async_get_registry(hass: HomeAssistantType) -> AreaRegistry:
     """Return area registry instance."""
-    task = hass.data.get(DATA_REGISTRY)
+    reg_or_evt = hass.data.get(DATA_REGISTRY)
 
-    if task is None:
-        async def _load_reg() -> AreaRegistry:
-            registry = AreaRegistry(hass)
-            await registry.async_load()
-            return registry
+    if not reg_or_evt:
+        evt = hass.data[DATA_REGISTRY] = Event()
 
-        task = hass.data[DATA_REGISTRY] = hass.async_create_task(_load_reg())
+        reg = AreaRegistry(hass)
+        await reg.async_load()
 
-    return cast(AreaRegistry, await task)
+        hass.data[DATA_REGISTRY] = reg
+        evt.set()
+        return reg
+
+    if isinstance(reg_or_evt, Event):
+        evt = reg_or_evt
+        await evt.wait()
+        return cast(AreaRegistry, hass.data.get(DATA_REGISTRY))
+
+    return cast(AreaRegistry, reg_or_evt)

homeassistant/helpers/device_registry.py

@@ -1,15 +1,17 @@
 """Provide a way to connect entities belonging to one device."""
 import logging
 import uuid
-from typing import List, Optional
-
+from asyncio import Event
 from collections import OrderedDict
+from typing import List, Optional, cast
 
 import attr
 
 from homeassistant.core import callback
 from homeassistant.loader import bind_hass
 
+from .typing import HomeAssistantType
+
 _LOGGER = logging.getLogger(__name__)
 _UNDEF = object()
 
@@ -273,19 +275,26 @@ class DeviceRegistry:
 
 
 @bind_hass
-async def async_get_registry(hass) -> DeviceRegistry:
+async def async_get_registry(hass: HomeAssistantType) -> DeviceRegistry:
     """Return device registry instance."""
-    task = hass.data.get(DATA_REGISTRY)
+    reg_or_evt = hass.data.get(DATA_REGISTRY)
 
-    if task is None:
-        async def _load_reg():
-            registry = DeviceRegistry(hass)
-            await registry.async_load()
-            return registry
+    if not reg_or_evt:
+        evt = hass.data[DATA_REGISTRY] = Event()
 
-        task = hass.data[DATA_REGISTRY] = hass.async_create_task(_load_reg())
+        reg = DeviceRegistry(hass)
+        await reg.async_load()
 
-    return await task
+        hass.data[DATA_REGISTRY] = reg
+        evt.set()
+        return reg
+
+    if isinstance(reg_or_evt, Event):
+        evt = reg_or_evt
+        await evt.wait()
+        return cast(DeviceRegistry, hass.data.get(DATA_REGISTRY))
+
+    return cast(DeviceRegistry, reg_or_evt)
 
 
 @callback

homeassistant/helpers/entity_registry.py

@@ -7,10 +7,11 @@ The Entity Registry will persist itself 10 seconds after a new entity is
 registered. Registering a new entity while a timer is in progress resets the
 timer.
 """
+from asyncio import Event
 from collections import OrderedDict
 from itertools import chain
 import logging
-from typing import Optional, List
+from typing import List, Optional, cast
 import weakref
 
 import attr
@@ -20,6 +21,8 @@ from homeassistant.loader import bind_hass
 from homeassistant.util import ensure_unique_string, slugify
 from homeassistant.util.yaml import load_yaml
 
+from .typing import HomeAssistantType
+
 PATH_REGISTRY = 'entity_registry.yaml'
 DATA_REGISTRY = 'entity_registry'
 SAVE_DELAY = 10
@@ -277,19 +280,26 @@ class EntityRegistry:
 
 
 @bind_hass
-async def async_get_registry(hass) -> EntityRegistry:
+async def async_get_registry(hass: HomeAssistantType) -> EntityRegistry:
     """Return entity registry instance."""
-    task = hass.data.get(DATA_REGISTRY)
+    reg_or_evt = hass.data.get(DATA_REGISTRY)
 
-    if task is None:
-        async def _load_reg():
-            registry = EntityRegistry(hass)
-            await registry.async_load()
-            return registry
+    if not reg_or_evt:
+        evt = hass.data[DATA_REGISTRY] = Event()
 
-        task = hass.data[DATA_REGISTRY] = hass.async_create_task(_load_reg())
+        reg = EntityRegistry(hass)
+        await reg.async_load()
 
-    return await task
+        hass.data[DATA_REGISTRY] = reg
+        evt.set()
+        return reg
+
+    if isinstance(reg_or_evt, Event):
+        evt = reg_or_evt
+        await evt.wait()
+        return cast(EntityRegistry, hass.data.get(DATA_REGISTRY))
+
+    return cast(EntityRegistry, reg_or_evt)
 
 
 @callback

tests/common.py

@@ -327,11 +327,7 @@ def mock_registry(hass, mock_entries=None):
     registry = entity_registry.EntityRegistry(hass)
     registry.entities = mock_entries or OrderedDict()
 
-    async def _get_reg():
-        return registry
-
-    hass.data[entity_registry.DATA_REGISTRY] = \
-        hass.loop.create_task(_get_reg())
+    hass.data[entity_registry.DATA_REGISTRY] = registry
     return registry
 
 
@@ -340,11 +336,7 @@ def mock_area_registry(hass, mock_entries=None):
     registry = area_registry.AreaRegistry(hass)
     registry.areas = mock_entries or OrderedDict()
 
-    async def _get_reg():
-        return registry
-
-    hass.data[area_registry.DATA_REGISTRY] = \
-        hass.loop.create_task(_get_reg())
+    hass.data[area_registry.DATA_REGISTRY] = registry
     return registry
 
 
@@ -353,11 +345,7 @@ def mock_device_registry(hass, mock_entries=None):
     registry = device_registry.DeviceRegistry(hass)
     registry.devices = mock_entries or OrderedDict()
 
-    async def _get_reg():
-        return registry
-
-    hass.data[device_registry.DATA_REGISTRY] = \
-        hass.loop.create_task(_get_reg())
+    hass.data[device_registry.DATA_REGISTRY] = registry
     return registry