Lazy load zwave_js platforms when the first entity needs to be created (#49016)

* Lazy load zwave_js platforms when the first entity needs to be created

* switch order to make things easier to understand

* await task instead of using wait_for_done callback

* gather tasks

* switch from asyncio.create_task to hass.async_create_task

* unsubscribe from callbacks before unloading platforms

* Clean up as much as possible during entry unload, even if a platform unload fails
This commit is contained in:
Raman Gupta 2021-04-12 20:26:49 -04:00 committed by GitHub
parent 53853f035d
commit cc40e681e2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 72 additions and 58 deletions

View File

@ -54,11 +54,11 @@ from .const import (
CONF_USB_PATH,
CONF_USE_ADDON,
DATA_CLIENT,
DATA_PLATFORM_SETUP,
DATA_UNSUBSCRIBE,
DOMAIN,
EVENT_DEVICE_ADDED_TO_REGISTRY,
LOGGER,
PLATFORMS,
ZWAVE_JS_NOTIFICATION_EVENT,
ZWAVE_JS_VALUE_NOTIFICATION_EVENT,
)
@ -113,49 +113,69 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
client = ZwaveClient(entry.data[CONF_URL], async_get_clientsession(hass))
dev_reg = device_registry.async_get(hass)
ent_reg = entity_registry.async_get(hass)
entry_hass_data: dict = hass.data[DOMAIN].setdefault(entry.entry_id, {})
@callback
def async_on_node_ready(node: ZwaveNode) -> None:
unsubscribe_callbacks: list[Callable] = []
entry_hass_data[DATA_CLIENT] = client
entry_hass_data[DATA_UNSUBSCRIBE] = unsubscribe_callbacks
entry_hass_data[DATA_PLATFORM_SETUP] = {}
async def async_on_node_ready(node: ZwaveNode) -> None:
"""Handle node ready event."""
LOGGER.debug("Processing node %s", node)
platform_setup_tasks = entry_hass_data[DATA_PLATFORM_SETUP]
# register (or update) node in device registry
register_node_in_dev_reg(hass, entry, dev_reg, client, node)
# run discovery on all node values and create/update entities
for disc_info in async_discover_values(node):
LOGGER.debug("Discovered entity: %s", disc_info)
# This migration logic was added in 2021.3 to handle a breaking change to
# the value_id format. Some time in the future, this call (as well as the
# helper functions) can be removed.
async_migrate_discovered_value(ent_reg, client, disc_info)
if disc_info.platform not in platform_setup_tasks:
platform_setup_tasks[disc_info.platform] = hass.async_create_task(
hass.config_entries.async_forward_entry_setup(
entry, disc_info.platform
)
)
await platform_setup_tasks[disc_info.platform]
LOGGER.debug("Discovered entity: %s", disc_info)
async_dispatcher_send(
hass, f"{DOMAIN}_{entry.entry_id}_add_{disc_info.platform}", disc_info
)
# add listener for stateless node value notification events
node.on(
"value notification",
lambda event: async_on_value_notification(event["value_notification"]),
unsubscribe_callbacks.append(
node.on(
"value notification",
lambda event: async_on_value_notification(event["value_notification"]),
)
)
# add listener for stateless node notification events
node.on(
"notification", lambda event: async_on_notification(event["notification"])
unsubscribe_callbacks.append(
node.on(
"notification",
lambda event: async_on_notification(event["notification"]),
)
)
@callback
def async_on_node_added(node: ZwaveNode) -> None:
async def async_on_node_added(node: ZwaveNode) -> None:
"""Handle node added event."""
# we only want to run discovery when the node has reached ready state,
# otherwise we'll have all kinds of missing info issues.
if node.ready:
async_on_node_ready(node)
await async_on_node_ready(node)
return
# if node is not yet ready, register one-time callback for ready state
LOGGER.debug("Node added: %s - waiting for it to become ready", node.node_id)
node.once(
"ready",
lambda event: async_on_node_ready(event["node"]),
lambda event: hass.async_create_task(async_on_node_ready(event["node"])),
)
# we do submit the node to device registry so user has
# some visual feedback that something is (in the process of) being added
@ -234,7 +254,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
hass.bus.async_fire(ZWAVE_JS_NOTIFICATION_EVENT, event_data)
entry_hass_data: dict = hass.data[DOMAIN].setdefault(entry.entry_id, {})
# connect and throw error if connection failed
try:
async with timeout(CONNECT_TIMEOUT):
@ -256,10 +275,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
entry_hass_data[DATA_CONNECT_FAILED_LOGGED] = False
entry_hass_data[DATA_INVALID_SERVER_VERSION_LOGGED] = False
unsubscribe_callbacks: list[Callable] = []
entry_hass_data[DATA_CLIENT] = client
entry_hass_data[DATA_UNSUBSCRIBE] = unsubscribe_callbacks
services = ZWaveServices(hass, ent_reg)
services.async_register()
@ -268,14 +283,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
async def start_platforms() -> None:
"""Start platforms and perform discovery."""
# wait until all required platforms are ready
await asyncio.gather(
*[
hass.config_entries.async_forward_entry_setup(entry, platform)
for platform in PLATFORMS
]
)
driver_ready = asyncio.Event()
async def handle_ha_shutdown(event: Event) -> None:
@ -313,17 +320,28 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
dev_reg.async_remove_device(device.id)
# run discovery on all ready nodes
for node in client.driver.controller.nodes.values():
async_on_node_added(node)
await asyncio.gather(
*[
async_on_node_added(node)
for node in client.driver.controller.nodes.values()
]
)
# listen for new nodes being added to the mesh
client.driver.controller.on(
"node added", lambda event: async_on_node_added(event["node"])
unsubscribe_callbacks.append(
client.driver.controller.on(
"node added",
lambda event: hass.async_create_task(
async_on_node_added(event["node"])
),
)
)
# listen for nodes being removed from the mesh
# NOTE: This will not remove nodes that were removed when HA was not running
client.driver.controller.on(
"node removed", lambda event: async_on_node_removed(event["node"])
unsubscribe_callbacks.append(
client.driver.controller.on(
"node removed", lambda event: async_on_node_removed(event["node"])
)
)
platform_task = hass.async_create_task(start_platforms())
@ -355,7 +373,7 @@ async def client_listen(
# All model instances will be replaced when the new state is acquired.
if should_reload:
LOGGER.info("Disconnected from server. Reloading integration")
asyncio.create_task(hass.config_entries.async_reload(entry.entry_id))
hass.async_create_task(hass.config_entries.async_reload(entry.entry_id))
async def disconnect_client(
@ -368,8 +386,13 @@ async def disconnect_client(
"""Disconnect client."""
listen_task.cancel()
platform_task.cancel()
platform_setup_tasks = (
hass.data[DOMAIN].get(entry.entry_id, {}).get(DATA_PLATFORM_SETUP, {}).values()
)
for task in platform_setup_tasks:
task.cancel()
await asyncio.gather(listen_task, platform_task)
await asyncio.gather(listen_task, platform_task, *platform_setup_tasks)
if client.connected:
await client.disconnect()
@ -378,22 +401,23 @@ async def disconnect_client(
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload a config entry."""
unload_ok = all(
await asyncio.gather(
*[
hass.config_entries.async_forward_entry_unload(entry, platform)
for platform in PLATFORMS
]
)
)
if not unload_ok:
return False
info = hass.data[DOMAIN].pop(entry.entry_id)
for unsub in info[DATA_UNSUBSCRIBE]:
unsub()
tasks = []
for platform, task in info[DATA_PLATFORM_SETUP].items():
if task.done():
tasks.append(
hass.config_entries.async_forward_entry_unload(entry, platform)
)
else:
task.cancel()
tasks.append(task)
unload_ok = all(await asyncio.gather(*tasks))
if DATA_CLIENT_LISTEN_TASK in info:
await disconnect_client(
hass,
@ -412,7 +436,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
LOGGER.error("Failed to stop the Z-Wave JS add-on: %s", err)
return False
return True
return unload_ok
async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:

View File

@ -8,20 +8,10 @@ CONF_NETWORK_KEY = "network_key"
CONF_USB_PATH = "usb_path"
CONF_USE_ADDON = "use_addon"
DOMAIN = "zwave_js"
PLATFORMS = [
"binary_sensor",
"climate",
"cover",
"fan",
"light",
"lock",
"number",
"sensor",
"switch",
]
DATA_CLIENT = "client"
DATA_UNSUBSCRIBE = "unsubs"
DATA_PLATFORM_SETUP = "platform_setup"
EVENT_DEVICE_ADDED_TO_REGISTRY = f"{DOMAIN}_device_added_to_registry"