Release discarded registries after cached queries

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Cached queries no longer prevent discarded registries and their entities from being garbage-collected.
 - Queries requiring and excluding the same condition now return no matches instead of raising `AssertionError`.
 - Component reads and tag cleanup no longer retain unused entities through empty registry rows.
 - Failed component removals no longer create empty rows, and clearing removes existing empty rows.

tcod/ecs/query.py

@@ -6,7 +6,7 @@
 import warnings
 from collections import defaultdict
 from typing import TYPE_CHECKING, Any, Protocol, TypeVar, overload
-from weakref import WeakKeyDictionary, WeakSet
+from weakref import WeakSet
 
 import attrs
 from typing_extensions import Self, deprecated
@@ -29,10 +29,6 @@
 _T5 = TypeVar("_T5")
 
 
-_query_caches: WeakKeyDictionary[Registry, _QueryCache] = WeakKeyDictionary()
-"""The master table of cached queries."""
-
-
 @attrs.define
 class _QueryCache:
     """Main data structure for the query cache."""
@@ -59,7 +55,7 @@ def _drop_cached_query(cache: _QueryCache, query: _Query) -> None:
     """Drop a cached query and all of its dependant queries."""
     cache.queries.pop(query, None)
     for sub_registry, sub_query in cache.dependencies.pop(query, ()):
-        _drop_cached_query(_query_caches[sub_registry], sub_query)
+        _drop_cached_query(_get_query_cache(sub_registry), sub_query)
 
 
 def _touch_component(registry: Registry, component: ComponentKey[object]) -> None:
@@ -126,11 +122,8 @@ def _fetch_relation_table(registry: Registry, relation: _RelationQuery) -> Abstr
 
 
 def _get_query_cache(registry: Registry) -> _QueryCache:
-    """Return the global cache for the given registry, creating it if it does not exist."""
-    cache = _query_caches.get(registry)
-    if cache is None:
-        cache = _query_caches[registry] = _QueryCache()
-    return cache
+    """Return the cache owned by the given registry."""
+    return registry._query_cache
 
 
 def _get_query(registry: Registry, query: _Query) -> AbstractSet[Entity]:

tcod/ecs/registry.py

@@ -13,6 +13,7 @@
 import tcod.ecs._converter
 import tcod.ecs.query
 from tcod.ecs.entity import Entity
+from tcod.ecs.query import _QueryCache
 
 if TYPE_CHECKING:
     from tcod.ecs.typing import ComponentKey, _RelationTargetLookup
@@ -152,6 +153,9 @@ class Registry:
     dict[Entity] = entities_name
     """
 
+    _query_cache: _QueryCache = attrs.field(init=False, factory=_QueryCache, repr=False)
+    """Transient query cache, owned here so cached entities do not keep registries alive globally."""
+
     @property
     @deprecated(
         "The 'registry.global_' attribute has been deprecated. Use 'registry[None]' to access this entity.",
@@ -178,6 +182,7 @@ def global_(self) -> Entity:
 
     def __setstate__(self, state: dict[str, Any]) -> None:
         """Unpickle this object and handle state migration."""
+        self._query_cache = _QueryCache()
         global_: Entity | None = state.pop("global_", None)  # Migrate from version <=1.2.0
 
         # These attributes contain redundant data and will be removed

tests/test_entity_lifecycle.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import gc
+import pickle
 import weakref
 
 import pytest
@@ -81,3 +82,50 @@ def test_incoming_relation_retains_cleared_entity() -> None:
     del owner.relation_tag["target"]
     gc.collect()
     assert reference() is None
+
+
+@pytest.mark.parametrize("query_kind", ["component", "tag", "relation", "compound", "subquery"])
+@pytest.mark.parametrize("restored", [False, True])
+def test_queried_registry_is_released(query_kind: str, *, restored: bool) -> None:
+    """Query caches must not keep a discarded registry and its entities alive."""
+    registry = Registry()
+    registry["player"].components[int] = 3
+    registry["player"].tags.add("actor")
+    registry["player"].relation_tag["room"] = registry["room"]
+    registry["room"].tags.add("active")
+    if restored:
+        registry = pickle.loads(pickle.dumps(registry))  # noqa: S301
+    if query_kind == "component":
+        query = registry.Q.all_of(components=[int])
+    elif query_kind == "tag":
+        query = registry.Q.all_of(tags=["actor"])
+    elif query_kind == "relation":
+        query = registry.Q.all_of(relations=[("room", registry["room"])])
+    elif query_kind == "compound":
+        query = registry.Q.all_of(components=[int], tags=["actor"])
+    else:
+        query = registry.Q.all_of(relations=[("room", registry.Q.all_of(tags=["active"]))])
+    assert set(query) == {registry["player"]}
+    registry_ref = weakref.ref(registry)
+    entity_ref = weakref.ref(registry["player"])
+    del query, registry
+    gc.collect()
+    assert registry_ref() is None
+    assert entity_ref() is None
+
+
+def test_retained_query_keeps_registry_alive() -> None:
+    """An application-held query continues to own its registry."""
+    registry = Registry()
+    registry["player"].tags.add("actor")
+    query = registry.Q.all_of(tags=["actor"])
+    assert set(query) == {registry["player"]}
+    reference = weakref.ref(registry)
+    del registry
+    gc.collect()
+    assert reference() is query.registry
+    query.registry["player"].tags.clear()
+    assert not query
+    del query
+    gc.collect()
+    assert reference() is None