Improve MQTT reload performance (#73313)

* Improve MQTT reload performance

* Update homeassistant/components/mqtt/__init__.py

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

* Update homeassistant/components/mqtt/mixins.py

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

* Update homeassistant/components/mqtt/__init__.py

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

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

homeassistant/components/mqtt/__init__.py

@@ -28,7 +28,14 @@ from homeassistant.data_entry_flow import BaseServiceInfo
 from homeassistant.exceptions import TemplateError, Unauthorized
 from homeassistant.helpers import config_validation as cv, event, template
 from homeassistant.helpers.device_registry import DeviceEntry
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.dispatcher import (
+    async_dispatcher_connect,
+    async_dispatcher_send,
+)
+from homeassistant.helpers.reload import (
+    async_integration_yaml_config,
+    async_setup_reload_service,
+)
 from homeassistant.helpers.typing import ConfigType
 
 # Loading the config flow file will register the flow
@@ -60,12 +67,14 @@ from .const import (  # noqa: F401
     DATA_MQTT,
     DATA_MQTT_CONFIG,
     DATA_MQTT_RELOAD_NEEDED,
+    DATA_MQTT_UPDATED_CONFIG,
     DEFAULT_ENCODING,
     DEFAULT_QOS,
     DEFAULT_RETAIN,
     DOMAIN,
     MQTT_CONNECTED,
     MQTT_DISCONNECTED,
+    MQTT_RELOADED,
     PLATFORMS,
 )
 from .models import (  # noqa: F401
@@ -227,7 +236,9 @@ async def _async_config_entry_updated(hass: HomeAssistant, entry: ConfigEntry) -
         await _async_setup_discovery(hass, mqtt_client.conf, entry)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(  # noqa: C901
+    hass: HomeAssistant, entry: ConfigEntry
+) -> bool:
     """Load a config entry."""
     # Merge basic configuration, and add missing defaults for basic options
     _merge_basic_config(hass, entry, hass.data.get(DATA_MQTT_CONFIG, {}))
@@ -364,6 +375,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     hass.data[DATA_CONFIG_ENTRY_LOCK] = asyncio.Lock()
     hass.data[CONFIG_ENTRY_IS_SETUP] = set()
 
+    # Setup reload service. Once support for legacy config is removed in 2022.9, we
+    # should no longer call async_setup_reload_service but instead implement a custom
+    # service
+    await async_setup_reload_service(hass, DOMAIN, PLATFORMS)
+
+    async def _async_reload_platforms(_: Event | None) -> None:
+        """Discover entities for a platform."""
+        config_yaml = await async_integration_yaml_config(hass, DOMAIN) or {}
+        hass.data[DATA_MQTT_UPDATED_CONFIG] = config_yaml.get(DOMAIN, {})
+        async_dispatcher_send(hass, MQTT_RELOADED)
+
     async def async_forward_entry_setup():
         """Forward the config entry setup to the platforms."""
         async with hass.data[DATA_CONFIG_ENTRY_LOCK]:
@@ -374,6 +396,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
                     await hass.config_entries.async_forward_entry_setup(
                         entry, component
                     )
+        # Setup reload service after all platforms have loaded
+        entry.async_on_unload(
+            hass.bus.async_listen("event_mqtt_reloaded", _async_reload_platforms)
+        )
 
     hass.async_create_task(async_forward_entry_setup())
 

homeassistant/components/mqtt/const.py

@@ -35,6 +35,7 @@ DATA_CONFIG_ENTRY_LOCK = "mqtt_config_entry_lock"
 DATA_MQTT = "mqtt"
 DATA_MQTT_CONFIG = "mqtt_config"
 DATA_MQTT_RELOAD_NEEDED = "mqtt_reload_needed"
+DATA_MQTT_UPDATED_CONFIG = "mqtt_updated_config"
 
 DEFAULT_PREFIX = "homeassistant"
 DEFAULT_BIRTH_WILL_TOPIC = DEFAULT_PREFIX + "/status"
@@ -63,6 +64,7 @@ DOMAIN = "mqtt"
 
 MQTT_CONNECTED = "mqtt_connected"
 MQTT_DISCONNECTED = "mqtt_disconnected"
+MQTT_RELOADED = "mqtt_reloaded"
 
 PAYLOAD_EMPTY_JSON = "{}"
 PAYLOAD_NONE = "None"

homeassistant/components/mqtt/mixins.py

@@ -48,10 +48,6 @@ from homeassistant.helpers.entity import (
     async_generate_entity_id,
 )
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.reload import (
-    async_integration_yaml_config,
-    async_setup_reload_service,
-)
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
 from . import debug_info, subscription
@@ -67,13 +63,14 @@ from .const import (
     DATA_MQTT,
     DATA_MQTT_CONFIG,
     DATA_MQTT_RELOAD_NEEDED,
+    DATA_MQTT_UPDATED_CONFIG,
     DEFAULT_ENCODING,
     DEFAULT_PAYLOAD_AVAILABLE,
     DEFAULT_PAYLOAD_NOT_AVAILABLE,
     DOMAIN,
     MQTT_CONNECTED,
     MQTT_DISCONNECTED,
-    PLATFORMS,
+    MQTT_RELOADED,
 )
 from .debug_info import log_message, log_messages
 from .discovery import (
@@ -270,14 +267,11 @@ async def async_setup_platform_discovery(
 ) -> CALLBACK_TYPE:
     """Set up platform discovery for manual config."""
 
-    async def _async_discover_entities(event: Event | None) -> None:
+    async def _async_discover_entities() -> None:
         """Discover entities for a platform."""
-        if event:
+        if DATA_MQTT_UPDATED_CONFIG in hass.data:
             # The platform has been reloaded
-            config_yaml = await async_integration_yaml_config(hass, DOMAIN)
+            config_yaml = hass.data[DATA_MQTT_UPDATED_CONFIG]
-            if not config_yaml:
-                return
-            config_yaml = config_yaml.get(DOMAIN, {})
         else:
             config_yaml = hass.data.get(DATA_MQTT_CONFIG, {})
         if not config_yaml:
@@ -293,8 +287,8 @@ async def async_setup_platform_discovery(
             )
         )
 
-    unsub = hass.bus.async_listen("event_mqtt_reloaded", _async_discover_entities)
+    unsub = async_dispatcher_connect(hass, MQTT_RELOADED, _async_discover_entities)
-    await _async_discover_entities(None)
+    await _async_discover_entities()
     return unsub
 
 
@@ -359,7 +353,6 @@ async def async_setup_platform_helper(
     async_setup_entities: SetupEntity,
 ) -> None:
     """Return true if platform setup should be aborted."""
-    await async_setup_reload_service(hass, DOMAIN, PLATFORMS)
     if not bool(hass.config_entries.async_entries(DOMAIN)):
         hass.data[DATA_MQTT_RELOAD_NEEDED] = None
         _LOGGER.warning(