Browse Source

Revert "Move global user storage from WeakValueDictionary to dict"

This reverts commit cb2363f0fd.

This lead to memory leaks due to insufficient tracking, assuming that
the members intent was enabled.
pull/7494/head
Rapptz 3 years ago
parent
commit
940bdb988a
  1. 48
      discord/state.py
  2. 21
      discord/user.py

48
discord/state.py

@ -31,6 +31,7 @@ import datetime
import itertools import itertools
import logging import logging
from typing import Dict, Optional, TYPE_CHECKING, Union, Callable, Any, List, TypeVar, Coroutine, Sequence, Tuple, Deque from typing import Dict, Optional, TYPE_CHECKING, Union, Callable, Any, List, TypeVar, Coroutine, Sequence, Tuple, Deque
import weakref
import inspect import inspect
import os import os
@ -225,8 +226,7 @@ class ConnectionState:
self._intents: Intents = intents self._intents: Intents = intents
if not intents.members or cache_flags._empty: if not intents.members or cache_flags._empty:
self.store_user = self.create_user # type: ignore self.store_user = self.store_user_no_intents # type: ignore
self.deref_user = self.deref_user_no_intents # type: ignore
self.parsers = parsers = {} self.parsers = parsers = {}
for attr, func in inspect.getmembers(self): for attr, func in inspect.getmembers(self):
@ -237,19 +237,7 @@ class ConnectionState:
def clear(self, *, views: bool = True) -> None: def clear(self, *, views: bool = True) -> None:
self.user: Optional[ClientUser] = None self.user: Optional[ClientUser] = None
# Originally, this code used WeakValueDictionary to maintain references to the self._users: weakref.WeakValueDictionary[int, User] = weakref.WeakValueDictionary()
# 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: Dict[int, Emoji] = {} self._emojis: Dict[int, Emoji] = {}
self._stickers: Dict[int, GuildSticker] = {} self._stickers: Dict[int, GuildSticker] = {}
self._guilds: Dict[int, Guild] = {} self._guilds: Dict[int, Guild] = {}
@ -324,7 +312,8 @@ class ConnectionState:
for vc in self.voice_clients: for vc in self.voice_clients:
vc.main_ws = ws # type: ignore vc.main_ws = ws # type: ignore
def store_user(self, data: UserPayload) -> User: def store_user(self, data):
# this way is 300% faster than `dict.setdefault`.
user_id = int(data['id']) user_id = int(data['id'])
try: try:
return self._users[user_id] return self._users[user_id]
@ -332,21 +321,16 @@ class ConnectionState:
user = User(state=self, data=data) user = User(state=self, data=data)
if user.discriminator != '0000': if user.discriminator != '0000':
self._users[user_id] = user self._users[user_id] = user
user._stored = True
return user return user
def deref_user(self, user_id: int) -> None: def store_user_no_intents(self, data):
self._users.pop(user_id, None) return User(state=self, data=data)
def create_user(self, data: UserPayload) -> User: def create_user(self, data: UserPayload) -> User:
return User(state=self, data=data) return User(state=self, data=data)
def deref_user_no_intents(self, user_id: int) -> None: def get_user(self, id):
return return self._users.get(id)
def get_user(self, id: Optional[int]) -> Optional[User]:
# the keys of self._users are ints
return self._users.get(id) # type: ignore
def store_emoji(self, guild: Guild, data: EmojiPayload) -> Emoji: def store_emoji(self, guild: Guild, data: EmojiPayload) -> Emoji:
# the id will be present here # the id will be present here
@ -551,8 +535,8 @@ class ConnectionState:
self._ready_state = asyncio.Queue() self._ready_state = asyncio.Queue()
self.clear(views=False) self.clear(views=False)
self.user = ClientUser(state=self, data=data['user']) self.user = user = ClientUser(state=self, data=data['user'])
self.store_user(data['user']) self._users[user.id] = user # type: ignore
if self.application_id is None: if self.application_id is None:
try: try:
@ -729,13 +713,9 @@ class ConnectionState:
self.dispatch('presence_update', old_member, member) self.dispatch('presence_update', old_member, member)
def parse_user_update(self, data) -> None: def parse_user_update(self, data):
# self.user is *always* cached when this is called if self.user:
user: ClientUser = self.user # type: ignore self.user._update(data)
user._update(data)
ref = self._users.get(user.id)
if ref:
ref._update(data)
def parse_invite_create(self, data) -> None: def parse_invite_create(self, data) -> None:
invite = Invite.from_gateway(state=self, data=data) invite = Invite.from_gateway(state=self, data=data)

21
discord/user.py

@ -429,29 +429,12 @@ class User(BaseUser, discord.abc.Messageable):
Specifies if the user is a system user (i.e. represents Discord officially). Specifies if the user is a system user (i.e. represents Discord officially).
""" """
__slots__ = ('_stored',) __slots__ = ('__weakref__',)
def __init__(self, *, state: ConnectionState, data: UserPayload) -> None:
super().__init__(state=state, data=data)
self._stored: bool = False
def __repr__(self) -> str: def __repr__(self) -> str:
return f'<User id={self.id} name={self.name!r} discriminator={self.discriminator!r} bot={self.bot}>' return f'<User id={self.id} name={self.name!r} discriminator={self.discriminator!r} bot={self.bot}>'
def __del__(self) -> None: async def _get_channel(self):
try:
if self._stored:
self._state.deref_user(self.id)
except Exception:
pass
@classmethod
def _copy(cls, user: User):
self = super()._copy(user)
self._stored = False
return self
async def _get_channel(self) -> DMChannel:
ch = await self.create_dm() ch = await self.create_dm()
return ch return ch

Loading…
Cancel
Save