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] 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: class BaseView:
__discord_ui_view__: ClassVar[bool] = False __discord_ui_view__: ClassVar[bool] = False
__discord_ui_modal__: ClassVar[bool] = False __discord_ui_modal__: ClassVar[bool] = False
@@ -220,6 +238,7 @@ class BaseView:
self.__cancel_callback: Optional[Callable[[BaseView], None]] = None self.__cancel_callback: Optional[Callable[[BaseView], None]] = None
self.__timeout_expiry: Optional[float] = None self.__timeout_expiry: Optional[float] = None
self.__timeout_task: Optional[asyncio.Task[None]] = None self.__timeout_task: Optional[asyncio.Task[None]] = None
self.__snapshot: Optional[_ViewCacheSnapshot] = None
try: try:
loop = asyncio.get_running_loop() loop = asyncio.get_running_loop()
@@ -326,6 +345,31 @@ class BaseView:
def _add_count(self, value: int) -> None: def _add_count(self, value: int) -> None:
self._total_children = max(0, self._total_children + value) 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 @property
def children(self) -> List[Item[Self]]: def children(self) -> List[Item[Self]]:
"""List[:class:`Item`]: The list of children attached to this view.""" """List[:class:`Item`]: The list of children attached to this view."""
@@ -901,6 +945,7 @@ class ViewStore:
dispatch_info = self._views.get(message_id, {}) dispatch_info = self._views.get(message_id, {})
is_fully_dynamic = True is_fully_dynamic = True
snapshot = view._get_snapshot_diff()
for item in view.walk_children(): for item in view.walk_children():
if isinstance(item, DynamicItem): if isinstance(item, DynamicItem):
pattern = item.__discord_ui_compiled_template__ pattern = item.__discord_ui_compiled_template__
@@ -909,6 +954,12 @@ class ViewStore:
dispatch_info[(item.type.value, item.custom_id)] = item # type: ignore dispatch_info[(item.type.value, item.custom_id)] = item # type: ignore
is_fully_dynamic = False 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 view._cache_key = message_id
if dispatch_info: if dispatch_info:
self._views[message_id] = dispatch_info self._views[message_id] = dispatch_info
@@ -922,13 +973,12 @@ class ViewStore:
return return
dispatch_info = self._views.get(view._cache_key) dispatch_info = self._views.get(view._cache_key)
if dispatch_info: snapshot = view._snapshot
for item in view.walk_children(): if dispatch_info and snapshot:
if isinstance(item, DynamicItem): for key in snapshot.items:
pattern = item.__discord_ui_compiled_template__ dispatch_info.pop(key, None)
self._dynamic_items.pop(pattern, None) for key in snapshot.dynamic_items:
elif item.is_dispatchable(): self._dynamic_items.pop(key, None)
dispatch_info.pop((item.type.value, item.custom_id), None) # type: ignore
if dispatch_info is not None and len(dispatch_info) == 0: if dispatch_info is not None and len(dispatch_info) == 0:
self._views.pop(view._cache_key, None) self._views.pop(view._cache_key, None)