Fix memory leak with the view store when removing items

The previous code would maintain items in the dispatch mapping if
nested children were removed between calls because it would only
remove items that are live in the view at the point of removal. This
meant that calling something like ActionRow.clear_items() would keep
all the removed items within the mapping and would not be evicted.

This attempts to fix it by maintaining a cache state snapshot and
making a diff between the two versions to know which keys are now safe
to delete since they are no longer in the live view at all.
This commit is contained in:
Rapptz
2026-03-02 11:31:05 -05:00
parent 9798e5921a
commit 616137875b

View File

@@ -207,6 +207,24 @@ class _ViewWeights:
self.weights = [0, 0, 0, 0, 0]
class _ViewCacheSnapshot:
__slots__ = ('items', 'dynamic_items')
def __init__(self) -> None:
self.items: Set[Tuple[int, str]] = set()
self.dynamic_items: Set[re.Pattern[str]] = set()
@classmethod
def diff(cls, older: _ViewCacheSnapshot, newer: _ViewCacheSnapshot) -> Self:
self = cls()
self.items = older.items - newer.items
self.dynamic_items = older.dynamic_items - newer.dynamic_items
return self
def __repr__(self) -> str:
return f'<_ViewCacheSnapshot items={self.items!r} dynamic_items={self.dynamic_items!r}>'
class BaseView:
__discord_ui_view__: ClassVar[bool] = False
__discord_ui_modal__: ClassVar[bool] = False
@@ -220,6 +238,7 @@ class BaseView:
self.__cancel_callback: Optional[Callable[[BaseView], None]] = None
self.__timeout_expiry: Optional[float] = None
self.__timeout_task: Optional[asyncio.Task[None]] = None
self.__snapshot: Optional[_ViewCacheSnapshot] = None
try:
loop = asyncio.get_running_loop()
@@ -326,6 +345,31 @@ class BaseView:
def _add_count(self, value: int) -> None:
self._total_children = max(0, self._total_children + value)
@property
def _snapshot(self) -> Optional[_ViewCacheSnapshot]:
return self.__snapshot
def _get_snapshot_diff(self) -> Optional[_ViewCacheSnapshot]:
if self.__snapshot is None:
self.__snapshot = self._get_snapshot()
return None
newer = self._get_snapshot()
diff = _ViewCacheSnapshot.diff(older=self.__snapshot, newer=newer)
# Update our snapshot to the newer version after diffing it
self.__snapshot = newer
return diff
def _get_snapshot(self) -> _ViewCacheSnapshot:
snapshot = _ViewCacheSnapshot()
for item in self.walk_children():
if isinstance(item, DynamicItem):
snapshot.dynamic_items.add(item.__discord_ui_compiled_template__)
elif item.is_dispatchable():
custom_id = item.custom_id # type: ignore
snapshot.items.add((item.type.value, custom_id))
return snapshot
@property
def children(self) -> List[Item[Self]]:
"""List[:class:`Item`]: The list of children attached to this view."""
@@ -901,6 +945,7 @@ class ViewStore:
dispatch_info = self._views.get(message_id, {})
is_fully_dynamic = True
snapshot = view._get_snapshot_diff()
for item in view.walk_children():
if isinstance(item, DynamicItem):
pattern = item.__discord_ui_compiled_template__
@@ -909,6 +954,12 @@ class ViewStore:
dispatch_info[(item.type.value, item.custom_id)] = item # type: ignore
is_fully_dynamic = False
if snapshot is not None:
for key in snapshot.items:
dispatch_info.pop(key, None)
for key in snapshot.dynamic_items:
self._dynamic_items.pop(key, None)
view._cache_key = message_id
if dispatch_info:
self._views[message_id] = dispatch_info
@@ -922,13 +973,12 @@ class ViewStore:
return
dispatch_info = self._views.get(view._cache_key)
if dispatch_info:
for item in view.walk_children():
if isinstance(item, DynamicItem):
pattern = item.__discord_ui_compiled_template__
self._dynamic_items.pop(pattern, None)
elif item.is_dispatchable():
dispatch_info.pop((item.type.value, item.custom_id), None) # type: ignore
snapshot = view._snapshot
if dispatch_info and snapshot:
for key in snapshot.items:
dispatch_info.pop(key, None)
for key in snapshot.dynamic_items:
self._dynamic_items.pop(key, None)
if dispatch_info is not None and len(dispatch_info) == 0:
self._views.pop(view._cache_key, None)