mirror of
https://github.com/Rapptz/discord.py.git
synced 2026-03-05 03:02:49 +00:00
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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user