diff --git a/discord/state.py b/discord/state.py index 6ebcfcec6..a04ce80fa 100644 --- a/discord/state.py +++ b/discord/state.py @@ -28,7 +28,7 @@ import copy import datetime import itertools import logging -from typing import Optional +from typing import Dict, Optional import weakref import warnings import inspect @@ -58,7 +58,6 @@ from .interactions import Interaction from .ui.view import ViewStore from .stage_instance import StageInstance from .threads import Thread, ThreadMember -from discord import guild class ChunkRequest: def __init__(self, guild_id, loop, resolver, *, cache=True): @@ -181,6 +180,7 @@ class ConnectionState: if not intents.members or cache_flags._empty: self.store_user = self.store_user_no_intents + self.deref_user = self.deref_user_no_intents self.parsers = parsers = {} for attr, func in inspect.getmembers(self): @@ -191,7 +191,19 @@ class ConnectionState: def clear(self): self.user = None - self._users = weakref.WeakValueDictionary() + # Originally, this code used WeakValueDictionary to maintain references to the + # global user mapping. + + # However, profiling showed that this came with two cons: + + # 1. The __weakref__ slot caused a non-trivial increase in memory + # 2. The performance of the mapping caused store_user to be a bottleneck. + + # Since this is undesirable, a mapping is now used instead with stored + # references now using a regular dictionary with eviction being done + # using __del__. Testing this for memory leaks led to no discernable leaks, + # though more testing will have to be done. + self._users: Dict[int, User] = {} self._emojis = {} self._guilds = {} self._view_store = ViewStore(self) @@ -265,7 +277,6 @@ class ConnectionState: vc.main_ws = ws def store_user(self, data): - # this way is 300% faster than `dict.setdefault`. user_id = int(data['id']) try: return self._users[user_id] @@ -273,11 +284,18 @@ class ConnectionState: user = User(state=self, data=data) if user.discriminator != '0000': self._users[user_id] = user + user._stored = True return user + def deref_user(self, user_id): + self._users.pop(user_id, None) + def store_user_no_intents(self, data): return User(state=self, data=data) + def deref_user_no_intents(self, user_id): + return + def get_user(self, id): return self._users.get(id) @@ -461,7 +479,7 @@ class ConnectionState: self._ready_state = asyncio.Queue() self.clear() self.user = user = ClientUser(state=self, data=data['user']) - self._users[user.id] = user + self.store_user(data['user']) if self.application_id is None: try: @@ -632,6 +650,9 @@ class ConnectionState: def parse_user_update(self, data): self.user._update(data) + ref = self._users.get(self.user.id) + if ref: + ref._update(data) def parse_invite_create(self, data): invite = Invite.from_gateway(state=self, data=data) diff --git a/discord/user.py b/discord/user.py index 4771df0a9..39c0f5b6d 100644 --- a/discord/user.py +++ b/discord/user.py @@ -322,11 +322,28 @@ class User(BaseUser, discord.abc.Messageable): Specifies if the user is a system user (i.e. represents Discord officially). """ - __slots__ = ('__weakref__',) + __slots__ = ('_stored',) + + def __init__(self, *, state, data): + super().__init__(state=state, data=data) + self._stored = False def __repr__(self): return f'' + def __del__(self) -> None: + try: + if self._stored: + self._state.deref_user(self.id) + except Exception: + pass + + @classmethod + def _copy(cls, user): + self = super()._copy(user) + self._stored = getattr(user, '_stored', False) + return self + async def _get_channel(self): ch = await self.create_dm() return ch