Browse Source

Move global user storage from WeakValueDictionary to dict

Profiling showed that WeakValueDictionary caused rather significant
and noticeable slowdowns during startup. Since the only thing it was
used for was to automatically remove the key from the mapping when
the reference count reaches zero, the same could theoretically be
accomplished by using the __del__ special method. There is a chance
that this could lead to a memory leak since the __del__ method is not
always called, but the only instances of this happening are during
interpreter shutdown to my knowledge and at that point the mapping
is the least of my concern.
pull/7190/head
Rapptz 4 years ago
parent
commit
cb2363f0fd
  1. 31
      discord/state.py
  2. 19
      discord/user.py

31
discord/state.py

@ -28,7 +28,7 @@ import copy
import datetime import datetime
import itertools import itertools
import logging import logging
from typing import Optional from typing import Dict, Optional
import weakref import weakref
import warnings import warnings
import inspect import inspect
@ -58,7 +58,6 @@ from .interactions import Interaction
from .ui.view import ViewStore from .ui.view import ViewStore
from .stage_instance import StageInstance from .stage_instance import StageInstance
from .threads import Thread, ThreadMember from .threads import Thread, ThreadMember
from discord import guild
class ChunkRequest: class ChunkRequest:
def __init__(self, guild_id, loop, resolver, *, cache=True): def __init__(self, guild_id, loop, resolver, *, cache=True):
@ -181,6 +180,7 @@ class ConnectionState:
if not intents.members or cache_flags._empty: if not intents.members or cache_flags._empty:
self.store_user = self.store_user_no_intents self.store_user = self.store_user_no_intents
self.deref_user = self.deref_user_no_intents
self.parsers = parsers = {} self.parsers = parsers = {}
for attr, func in inspect.getmembers(self): for attr, func in inspect.getmembers(self):
@ -191,7 +191,19 @@ class ConnectionState:
def clear(self): def clear(self):
self.user = None 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._emojis = {}
self._guilds = {} self._guilds = {}
self._view_store = ViewStore(self) self._view_store = ViewStore(self)
@ -265,7 +277,6 @@ class ConnectionState:
vc.main_ws = ws vc.main_ws = ws
def store_user(self, data): 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]
@ -273,11 +284,18 @@ 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):
self._users.pop(user_id, None)
def store_user_no_intents(self, data): def store_user_no_intents(self, data):
return User(state=self, data=data) return User(state=self, data=data)
def deref_user_no_intents(self, user_id):
return
def get_user(self, id): def get_user(self, id):
return self._users.get(id) return self._users.get(id)
@ -461,7 +479,7 @@ class ConnectionState:
self._ready_state = asyncio.Queue() self._ready_state = asyncio.Queue()
self.clear() self.clear()
self.user = user = ClientUser(state=self, data=data['user']) 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: if self.application_id is None:
try: try:
@ -632,6 +650,9 @@ class ConnectionState:
def parse_user_update(self, data): def parse_user_update(self, data):
self.user._update(data) self.user._update(data)
ref = self._users.get(self.user.id)
if ref:
ref._update(data)
def parse_invite_create(self, data): def parse_invite_create(self, data):
invite = Invite.from_gateway(state=self, data=data) invite = Invite.from_gateway(state=self, data=data)

19
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). 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): def __repr__(self):
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:
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): async def _get_channel(self):
ch = await self.create_dm() ch = await self.create_dm()
return ch return ch

Loading…
Cancel
Save