Stable entity registry id when a deleted entity is restored (#77710)

* Stable entity_id and registry id when a deleted entity is restored

* Don't restore area_id

* Don't restore entity_id

* Address review comments

homeassistant/helpers/entity_registry.py

@@ -11,7 +11,9 @@ from __future__ import annotations
 
 from collections import UserDict
 from collections.abc import Callable, Iterable, Mapping, ValuesView
+from datetime import datetime, timedelta
 import logging
+import time
 from typing import TYPE_CHECKING, Any, TypeVar, cast
 
 import attr
@@ -26,6 +28,7 @@ from homeassistant.const import (
     ATTR_SUPPORTED_FEATURES,
     ATTR_UNIT_OF_MEASUREMENT,
     EVENT_HOMEASSISTANT_START,
+    EVENT_HOMEASSISTANT_STOP,
     MAX_LENGTH_STATE_DOMAIN,
     MAX_LENGTH_STATE_ENTITY_ID,
     STATE_UNAVAILABLE,
@@ -61,9 +64,12 @@ SAVE_DELAY = 10
 _LOGGER = logging.getLogger(__name__)
 
 STORAGE_VERSION_MAJOR = 1
-STORAGE_VERSION_MINOR = 10
+STORAGE_VERSION_MINOR = 11
 STORAGE_KEY = "core.entity_registry"
 
+CLEANUP_INTERVAL = 3600 * 24
+ORPHANED_ENTITY_KEEP_SECONDS = 3600 * 24 * 30
+
 ENTITY_CATEGORY_VALUE_TO_INDEX: dict[EntityCategory | None, int] = {
     # mypy does not understand strenum
     val: idx  # type: ignore[misc]
@@ -138,7 +144,10 @@ class RegistryEntry:
     entity_category: EntityCategory | None = attr.ib(default=None)
     hidden_by: RegistryEntryHider | None = attr.ib(default=None)
     icon: str | None = attr.ib(default=None)
-    id: str = attr.ib(factory=uuid_util.random_uuid_hex)
+    id: str = attr.ib(
+        default=None,
+        converter=attr.converters.default_if_none(factory=uuid_util.random_uuid_hex),  # type: ignore[misc]
+    )
     has_entity_name: bool = attr.ib(default=False)
     name: str | None = attr.ib(default=None)
     options: ReadOnlyEntityOptionsType = attr.ib(
@@ -297,6 +306,24 @@ class RegistryEntry:
         hass.states.async_set(self.entity_id, STATE_UNAVAILABLE, attrs)
 
 
+@attr.s(slots=True, frozen=True)
+class DeletedRegistryEntry:
+    """Deleted Entity Registry Entry."""
+
+    entity_id: str = attr.ib()
+    unique_id: str = attr.ib()
+    platform: str = attr.ib()
+    config_entry_id: str | None = attr.ib()
+    domain: str = attr.ib(init=False, repr=False)
+    id: str = attr.ib()
+    orphaned_timestamp: float | None = attr.ib()
+
+    @domain.default
+    def _domain_default(self) -> str:
+        """Compute domain value."""
+        return split_entity_id(self.entity_id)[0]
+
+
 class EntityRegistryStore(storage.Store[dict[str, list[dict[str, Any]]]]):
     """Store entity registry data."""
 
@@ -372,6 +399,10 @@ class EntityRegistryStore(storage.Store[dict[str, list[dict[str, Any]]]]):
             for entity in data["entities"]:
                 entity["aliases"] = []
 
+        if old_major_version == 1 and old_minor_version < 11:
+            # Version 1.11 adds deleted_entities
+            data["deleted_entities"] = data.get("deleted_entities", [])
+
         if old_major_version > 1:
             raise NotImplementedError
         return data
@@ -424,6 +455,7 @@ class EntityRegistryItems(UserDict[str, "RegistryEntry"]):
 class EntityRegistry:
     """Class to hold a registry of entities."""
 
+    deleted_entities: dict[tuple[str, str, str], DeletedRegistryEntry]
     entities: EntityRegistryItems
     _entities_data: dict[str, RegistryEntry]
 
@@ -496,6 +528,9 @@ class EntityRegistry:
         - It's not registered
         - It's not known by the entity component adding the entity
         - It's not in the state machine
+
+        Note that an entity_id which belongs to a deleted entity is considered
+        available.
         """
         if known_object_ids is None:
             known_object_ids = {}
@@ -591,8 +626,16 @@ class EntityRegistry:
                 unit_of_measurement=unit_of_measurement,
             )
 
+        entity_registry_id: str | None = None
+        deleted_entity = self.deleted_entities.pop((domain, platform, unique_id), None)
+        if deleted_entity is not None:
+            # Restore id
+            entity_registry_id = deleted_entity.id
+
         entity_id = self.async_generate_entity_id(
-            domain, suggested_object_id or f"{platform}_{unique_id}", known_object_ids
+            domain,
+            suggested_object_id or f"{platform}_{unique_id}",
+            known_object_ids,
         )
 
         if disabled_by and not isinstance(disabled_by, RegistryEntryDisabler):
@@ -630,6 +673,7 @@ class EntityRegistry:
             entity_id=entity_id,
             hidden_by=hidden_by,
             has_entity_name=none_if_undefined(has_entity_name) or False,
+            id=entity_registry_id,
             options=initial_options,
             original_device_class=none_if_undefined(original_device_class),
             original_icon=none_if_undefined(original_icon),
@@ -653,7 +697,19 @@ class EntityRegistry:
     @callback
     def async_remove(self, entity_id: str) -> None:
         """Remove an entity from registry."""
-        self.entities.pop(entity_id)
+        entity = self.entities.pop(entity_id)
+        config_entry_id = entity.config_entry_id
+        key = (entity.domain, entity.platform, entity.unique_id)
+        # If the entity does not belong to a config entry, mark it as orphaned
+        orphaned_timestamp = None if config_entry_id else time.time()
+        self.deleted_entities[key] = DeletedRegistryEntry(
+            config_entry_id=config_entry_id,
+            entity_id=entity_id,
+            id=entity.id,
+            orphaned_timestamp=orphaned_timestamp,
+            platform=entity.platform,
+            unique_id=entity.unique_id,
+        )
         self.hass.bus.async_fire(
             EVENT_ENTITY_REGISTRY_UPDATED, {"action": "remove", "entity_id": entity_id}
         )
@@ -954,10 +1010,12 @@ class EntityRegistry:
 
     async def async_load(self) -> None:
         """Load the entity registry."""
-        async_setup_entity_restore(self.hass, self)
+        _async_setup_cleanup(self.hass, self)
+        _async_setup_entity_restore(self.hass, self)
 
         data = await self._store.async_load()
         entities = EntityRegistryItems()
+        deleted_entities: dict[tuple[str, str, str], DeletedRegistryEntry] = {}
 
         if data is not None:
             for entity in data["entities"]:
@@ -996,7 +1054,22 @@ class EntityRegistry:
                     unique_id=entity["unique_id"],
                     unit_of_measurement=entity["unit_of_measurement"],
                 )
+            for entity in data["deleted_entities"]:
+                key = (
+                    split_entity_id(entity["entity_id"])[0],
+                    entity["platform"],
+                    entity["unique_id"],
+                )
+                deleted_entities[key] = DeletedRegistryEntry(
+                    config_entry_id=entity["config_entry_id"],
+                    entity_id=entity["entity_id"],
+                    id=entity["id"],
+                    orphaned_timestamp=entity["orphaned_timestamp"],
+                    platform=entity["platform"],
+                    unique_id=entity["unique_id"],
+                )
 
+        self.deleted_entities = deleted_entities
         self.entities = entities
         self._entities_data = entities.data
 
@@ -1038,18 +1111,54 @@ class EntityRegistry:
             }
             for entry in self.entities.values()
         ]
+        data["deleted_entities"] = [
+            {
+                "config_entry_id": entry.config_entry_id,
+                "entity_id": entry.entity_id,
+                "id": entry.id,
+                "orphaned_timestamp": entry.orphaned_timestamp,
+                "platform": entry.platform,
+                "unique_id": entry.unique_id,
+            }
+            for entry in self.deleted_entities.values()
+        ]
 
         return data
 
     @callback
-    def async_clear_config_entry(self, config_entry: str) -> None:
+    def async_clear_config_entry(self, config_entry_id: str) -> None:
         """Clear config entry from registry entries."""
+        now_time = time.time()
         for entity_id in [
             entity_id
             for entity_id, entry in self.entities.items()
-            if config_entry == entry.config_entry_id
+            if config_entry_id == entry.config_entry_id
         ]:
             self.async_remove(entity_id)
+        for key, deleted_entity in list(self.deleted_entities.items()):
+            if config_entry_id != deleted_entity.config_entry_id:
+                continue
+            # Add a time stamp when the deleted entity became orphaned
+            self.deleted_entities[key] = attr.evolve(
+                deleted_entity, orphaned_timestamp=now_time, config_entry_id=None
+            )
+            self.async_schedule_save()
+
+    @callback
+    def async_purge_expired_orphaned_entities(self) -> None:
+        """Purge expired orphaned entities from the registry.
+
+        We need to purge these periodically to avoid the database
+        growing without bound.
+        """
+        now_time = time.time()
+        for key, deleted_entity in list(self.deleted_entities.items()):
+            if (orphaned_timestamp := deleted_entity.orphaned_timestamp) is None:
+                continue
+
+            if orphaned_timestamp + ORPHANED_ENTITY_KEEP_SECONDS < now_time:
+                self.deleted_entities.pop(key)
+                self.async_schedule_save()
 
     @callback
     def async_clear_area_id(self, area_id: str) -> None:
@@ -1136,7 +1245,31 @@ def async_config_entry_disabled_by_changed(
 
 
 @callback
-def async_setup_entity_restore(hass: HomeAssistant, registry: EntityRegistry) -> None:
+def _async_setup_cleanup(hass: HomeAssistant, registry: EntityRegistry) -> None:
+    """Clean up device registry when entities removed."""
+    from . import event  # pylint: disable=import-outside-toplevel
+
+    @callback
+    def cleanup(_: datetime) -> None:
+        """Clean up entity registry."""
+        # Periodic purge of orphaned entities to avoid the registry
+        # growing without bounds when there are lots of deleted entities
+        registry.async_purge_expired_orphaned_entities()
+
+    cancel = event.async_track_time_interval(
+        hass, cleanup, timedelta(seconds=CLEANUP_INTERVAL)
+    )
+
+    @callback
+    def _on_homeassistant_stop(event: Event) -> None:
+        """Cancel cleanup."""
+        cancel()
+
+    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _on_homeassistant_stop)
+
+
+@callback
+def _async_setup_entity_restore(hass: HomeAssistant, registry: EntityRegistry) -> None:
     """Set up the entity restore mechanism."""
 
     @callback

tests/common.py

@@ -511,6 +511,7 @@ def mock_registry(
     registry = er.EntityRegistry(hass)
     if mock_entries is None:
         mock_entries = {}
+    registry.deleted_entities = {}
     registry.entities = er.EntityRegistryItems()
     registry._entities_data = registry.entities.data
     for key, entry in mock_entries.items():

tests/helpers/test_entity_registry.py

@@ -1,7 +1,9 @@
 """Tests for the Entity Registry."""
+from datetime import timedelta
 from typing import Any
 from unittest.mock import patch
 
+import attr
 import pytest
 import voluptuous as vol
 
@@ -15,7 +17,7 @@ from homeassistant.core import CoreState, HomeAssistant, callback
 from homeassistant.exceptions import MaxLengthExceeded
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 
-from tests.common import MockConfigEntry, flush_store
+from tests.common import MockConfigEntry, async_fire_time_changed, flush_store
 
 YAML__OPEN_PATH = "homeassistant.util.yaml.loader.open"
 
@@ -276,8 +278,13 @@ async def test_loading_saving_data(
         orig_entry2.entity_id, "light", {"minimum_brightness": 20}
     )
     orig_entry2 = entity_registry.async_get(orig_entry2.entity_id)
+    orig_entry3 = entity_registry.async_get_or_create("light", "hue", "ABCD")
+    orig_entry4 = entity_registry.async_get_or_create("light", "hue", "EFGH")
+    entity_registry.async_remove(orig_entry3.entity_id)
+    entity_registry.async_remove(orig_entry4.entity_id)
 
     assert len(entity_registry.entities) == 2
+    assert len(entity_registry.deleted_entities) == 2
 
     # Now load written data in new registry
     registry2 = er.EntityRegistry(hass)
@@ -286,11 +293,16 @@ async def test_loading_saving_data(
 
     # Ensure same order
     assert list(entity_registry.entities) == list(registry2.entities)
+    assert list(entity_registry.deleted_entities) == list(registry2.deleted_entities)
     new_entry1 = entity_registry.async_get_or_create("light", "hue", "1234")
     new_entry2 = entity_registry.async_get_or_create("light", "hue", "5678")
+    new_entry3 = entity_registry.async_get_or_create("light", "hue", "ABCD")
+    new_entry4 = entity_registry.async_get_or_create("light", "hue", "EFGH")
 
     assert orig_entry1 == new_entry1
     assert orig_entry2 == new_entry2
+    assert orig_entry3 == new_entry3
+    assert orig_entry4 == new_entry4
 
     assert new_entry2.area_id == "mock-area-id"
     assert new_entry2.capabilities == {"max": 100}
@@ -485,6 +497,42 @@ async def test_removing_config_entry_id(
     assert update_events[1]["entity_id"] == entry.entity_id
 
 
+async def test_deleted_entity_removing_config_entry_id(
+    hass, entity_registry: er.EntityRegistry
+):
+    """Test that we update config entry id in registry on deleted entity."""
+    mock_config = MockConfigEntry(domain="light", entry_id="mock-id-1")
+
+    entry = entity_registry.async_get_or_create(
+        "light", "hue", "5678", config_entry=mock_config
+    )
+    assert entry.config_entry_id == "mock-id-1"
+    entity_registry.async_remove(entry.entity_id)
+
+    assert len(entity_registry.entities) == 0
+    assert len(entity_registry.deleted_entities) == 1
+    assert (
+        entity_registry.deleted_entities[("light", "hue", "5678")].config_entry_id
+        == "mock-id-1"
+    )
+    assert (
+        entity_registry.deleted_entities[("light", "hue", "5678")].orphaned_timestamp
+        is None
+    )
+
+    entity_registry.async_clear_config_entry("mock-id-1")
+    assert len(entity_registry.entities) == 0
+    assert len(entity_registry.deleted_entities) == 1
+    assert (
+        entity_registry.deleted_entities[("light", "hue", "5678")].config_entry_id
+        is None
+    )
+    assert (
+        entity_registry.deleted_entities[("light", "hue", "5678")].orphaned_timestamp
+        is not None
+    )
+
+
 async def test_removing_area_id(entity_registry: er.EntityRegistry) -> None:
     """Make sure we can clear area id."""
     entry = entity_registry.async_get_or_create("light", "hue", "5678")
@@ -1537,3 +1585,95 @@ def test_migrate_entity_to_new_platform(
             new_unique_id=new_unique_id,
             new_config_entry_id=new_config_entry.entry_id,
         )
+
+
+async def test_restore_entity(hass, update_events, freezer):
+    """Make sure entity registry id is stable and entity_id is reused if possible."""
+    registry = er.async_get(hass)  # We need the real entity registry for this test
+    config_entry = MockConfigEntry(domain="light")
+    entry1 = registry.async_get_or_create(
+        "light", "hue", "1234", config_entry=config_entry
+    )
+    entry2 = registry.async_get_or_create(
+        "light", "hue", "5678", config_entry=config_entry
+    )
+
+    entry1 = registry.async_update_entity(
+        entry1.entity_id, new_entity_id="light.custom_1"
+    )
+
+    registry.async_remove(entry1.entity_id)
+    registry.async_remove(entry2.entity_id)
+    assert len(registry.entities) == 0
+    assert len(registry.deleted_entities) == 2
+
+    # Re-add entities
+    entry1_restored = registry.async_get_or_create(
+        "light", "hue", "1234", config_entry=config_entry
+    )
+    entry2_restored = registry.async_get_or_create("light", "hue", "5678")
+
+    assert len(registry.entities) == 2
+    assert len(registry.deleted_entities) == 0
+    assert entry1 != entry1_restored
+    # entity_id is not restored
+    assert attr.evolve(entry1, entity_id="light.hue_1234") == entry1_restored
+    assert entry2 != entry2_restored
+    # Config entry is not restored
+    assert attr.evolve(entry2, config_entry_id=None) == entry2_restored
+
+    # Remove two of the entities again, then bump time
+    registry.async_remove(entry1_restored.entity_id)
+    registry.async_remove(entry2.entity_id)
+    assert len(registry.entities) == 0
+    assert len(registry.deleted_entities) == 2
+    freezer.tick(timedelta(seconds=er.ORPHANED_ENTITY_KEEP_SECONDS + 1))
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    # Re-add two entities, expect to get a new id after the purge for entity w/o config entry
+    entry1_restored = registry.async_get_or_create(
+        "light", "hue", "1234", config_entry=config_entry
+    )
+    entry2_restored = registry.async_get_or_create("light", "hue", "5678")
+    assert len(registry.entities) == 2
+    assert len(registry.deleted_entities) == 0
+    assert entry1.id == entry1_restored.id
+    assert entry2.id != entry2_restored.id
+
+    # Remove the first entity, then its config entry, finally bump time
+    registry.async_remove(entry1_restored.entity_id)
+    assert len(registry.entities) == 1
+    assert len(registry.deleted_entities) == 1
+    registry.async_clear_config_entry(config_entry.entry_id)
+    freezer.tick(timedelta(seconds=er.ORPHANED_ENTITY_KEEP_SECONDS + 1))
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    # Re-add the entity, expect to get a new id after the purge
+    entry1_restored = registry.async_get_or_create(
+        "light", "hue", "1234", config_entry=config_entry
+    )
+    assert len(registry.entities) == 2
+    assert len(registry.deleted_entities) == 0
+    assert entry1.id != entry1_restored.id
+
+    # Check the events
+    await hass.async_block_till_done()
+    assert len(update_events) == 13
+    assert update_events[0] == {"action": "create", "entity_id": "light.hue_1234"}
+    assert update_events[1] == {"action": "create", "entity_id": "light.hue_5678"}
+    assert update_events[2]["action"] == "update"
+    assert update_events[3] == {"action": "remove", "entity_id": "light.custom_1"}
+    assert update_events[4] == {"action": "remove", "entity_id": "light.hue_5678"}
+    # Restore entities the 1st time
+    assert update_events[5] == {"action": "create", "entity_id": "light.hue_1234"}
+    assert update_events[6] == {"action": "create", "entity_id": "light.hue_5678"}
+    assert update_events[7] == {"action": "remove", "entity_id": "light.hue_1234"}
+    assert update_events[8] == {"action": "remove", "entity_id": "light.hue_5678"}
+    # Restore entities the 2nd time
+    assert update_events[9] == {"action": "create", "entity_id": "light.hue_1234"}
+    assert update_events[10] == {"action": "create", "entity_id": "light.hue_5678"}
+    assert update_events[11] == {"action": "remove", "entity_id": "light.hue_1234"}
+    # Restore entities the 3rd time
+    assert update_events[12] == {"action": "create", "entity_id": "light.hue_1234"}