Skip to content

Commit

Permalink
Move global user storage from WeakValueDictionary to dict
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
Rapptz committed Jul 8, 2021
1 parent 88d825a commit cb2363f
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 6 deletions.
31 changes: 26 additions & 5 deletions discord/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
import datetime
import itertools
import logging
from typing import Optional
from typing import Dict, Optional
import weakref
import warnings
import inspect
Expand Down Expand Up @@ -58,7 +58,6 @@
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):
Expand Down Expand Up @@ -181,6 +180,7 @@ def __init__(self, *, dispatch, handlers, hooks, http, loop, **options):

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):
Expand All @@ -191,7 +191,19 @@ def __init__(self, *, dispatch, handlers, hooks, http, loop, **options):

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)
Expand Down Expand Up @@ -265,19 +277,25 @@ def _update_references(self, ws):
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]
except KeyError:
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)

Expand Down Expand Up @@ -461,7 +479,7 @@ def parse_ready(self, data):
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:
Expand Down Expand Up @@ -632,6 +650,9 @@ def parse_presence_update(self, data):

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)
Expand Down
19 changes: 18 additions & 1 deletion discord/user.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'<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):
ch = await self.create_dm()
return ch
Expand Down

0 comments on commit cb2363f

Please sign in to comment.