Fix memory leak in ESPHome disconnect callbacks (#104149)

homeassistant/components/esphome/bluetooth/client.py

@@ -136,7 +136,7 @@ class ESPHomeClientData:
     api_version: APIVersion
     title: str
     scanner: ESPHomeScanner | None
-    disconnect_callbacks: list[Callable[[], None]] = field(default_factory=list)
+    disconnect_callbacks: set[Callable[[], None]] = field(default_factory=set)
 
 
 class ESPHomeClient(BaseBleakClient):
@@ -215,6 +215,7 @@ class ESPHomeClient(BaseBleakClient):
             if not future.done():
                 future.set_result(None)
         self._disconnected_futures.clear()
+        self._disconnect_callbacks.discard(self._async_esp_disconnected)
         self._unsubscribe_connection_state()
 
     def _async_ble_device_disconnected(self) -> None:
@@ -228,7 +229,9 @@ class ESPHomeClient(BaseBleakClient):
     def _async_esp_disconnected(self) -> None:
         """Handle the esp32 client disconnecting from us."""
         _LOGGER.debug("%s: ESP device disconnected", self._description)
-        self._disconnect_callbacks.remove(self._async_esp_disconnected)
+        # Calling _async_ble_device_disconnected calls
+        # _async_disconnected_cleanup which will also remove
+        # the disconnect callbacks
         self._async_ble_device_disconnected()
 
     def _async_call_bleak_disconnected_callback(self) -> None:
@@ -289,7 +292,7 @@ class ESPHomeClient(BaseBleakClient):
             "%s: connected, registering for disconnected callbacks",
             self._description,
         )
-        self._disconnect_callbacks.append(self._async_esp_disconnected)
+        self._disconnect_callbacks.add(self._async_esp_disconnected)
         connected_future.set_result(connected)
 
     @api_error_as_bleak_error

homeassistant/components/esphome/entry_data.py

@@ -107,7 +107,7 @@ class RuntimeEntryData:
     bluetooth_device: ESPHomeBluetoothDevice | None = None
     api_version: APIVersion = field(default_factory=APIVersion)
     cleanup_callbacks: list[Callable[[], None]] = field(default_factory=list)
-    disconnect_callbacks: list[Callable[[], None]] = field(default_factory=list)
+    disconnect_callbacks: set[Callable[[], None]] = field(default_factory=set)
     state_subscriptions: dict[
         tuple[type[EntityState], int], Callable[[], None]
     ] = field(default_factory=dict)
@@ -427,3 +427,19 @@ class RuntimeEntryData:
         if self.original_options == entry.options:
             return
         hass.async_create_task(hass.config_entries.async_reload(entry.entry_id))
+
+    @callback
+    def async_on_disconnect(self) -> None:
+        """Call when the entry has been disconnected.
+
+        Safe to call multiple times.
+        """
+        self.available = False
+        # Make a copy since calling the disconnect callbacks
+        # may also try to discard/remove themselves.
+        for disconnect_cb in self.disconnect_callbacks.copy():
+            disconnect_cb()
+        # Make sure to clear the set to give up the reference
+        # to it and make sure all the callbacks can be GC'd.
+        self.disconnect_callbacks.clear()
+        self.disconnect_callbacks = set()

homeassistant/components/esphome/manager.py

@@ -295,7 +295,7 @@ class ESPHomeManager:
                 event.data["entity_id"], attribute, new_state
             )
 
-        self.entry_data.disconnect_callbacks.append(
+        self.entry_data.disconnect_callbacks.add(
             async_track_state_change_event(
                 hass, [entity_id], send_home_assistant_state_event
             )
@@ -440,7 +440,7 @@ class ESPHomeManager:
             reconnect_logic.name = device_info.name
 
         if device_info.bluetooth_proxy_feature_flags_compat(cli.api_version):
-            entry_data.disconnect_callbacks.append(
+            entry_data.disconnect_callbacks.add(
                 await async_connect_scanner(
                     hass, entry, cli, entry_data, self.domain_data.bluetooth_cache
                 )
@@ -462,7 +462,7 @@ class ESPHomeManager:
             )
 
             if device_info.voice_assistant_version:
-                entry_data.disconnect_callbacks.append(
+                entry_data.disconnect_callbacks.add(
                     await cli.subscribe_voice_assistant(
                         self._handle_pipeline_start,
                         self._handle_pipeline_stop,
@@ -490,10 +490,7 @@ class ESPHomeManager:
             host,
             expected_disconnect,
         )
-        for disconnect_cb in entry_data.disconnect_callbacks:
-            disconnect_cb()
-        entry_data.disconnect_callbacks = []
-        entry_data.available = False
+        entry_data.async_on_disconnect()
         entry_data.expected_disconnect = expected_disconnect
         # Mark state as stale so that we will always dispatch
         # the next state update of that type when the device reconnects
@@ -758,10 +755,7 @@ async def cleanup_instance(hass: HomeAssistant, entry: ConfigEntry) -> RuntimeEn
     """Cleanup the esphome client if it exists."""
     domain_data = DomainData.get(hass)
     data = domain_data.pop_entry_data(entry)
-    data.available = False
-    for disconnect_cb in data.disconnect_callbacks:
-        disconnect_cb()
-    data.disconnect_callbacks = []
+    data.async_on_disconnect()
     for cleanup_callback in data.cleanup_callbacks:
         cleanup_callback()
     await data.async_cleanup()