From cc40e681e2214c11324d2533ab97d1d501e45d1b Mon Sep 17 00:00:00 2001 From: Raman Gupta <7243222+raman325@users.noreply.github.com> Date: Mon, 12 Apr 2021 20:26:49 -0400 Subject: [PATCH] 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 --- homeassistant/components/zwave_js/__init__.py | 118 +++++++++++------- homeassistant/components/zwave_js/const.py | 12 +- 2 files changed, 72 insertions(+), 58 deletions(-) diff --git a/homeassistant/components/zwave_js/__init__.py b/homeassistant/components/zwave_js/__init__.py index 10cc2543921..45aef87bf80 100644 --- a/homeassistant/components/zwave_js/__init__.py +++ b/homeassistant/components/zwave_js/__init__.py @@ -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: diff --git a/homeassistant/components/zwave_js/const.py b/homeassistant/components/zwave_js/const.py index 1c9f78b1751..afd899e0ee0 100644 --- a/homeassistant/components/zwave_js/const.py +++ b/homeassistant/components/zwave_js/const.py @@ -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"