Better UX for rate limits (fixes #4) 🥳

cogs/emote.py

@@ -37,8 +37,9 @@ from discord.ext import commands
 import utils
 import utils.image
 from utils import errors
-from utils.converter import emote_type_filter_default
 from utils.paginator import ListPaginator
+from utils.emote_client import EmoteClient
+from utils.converter import emote_type_filter_default
 
 logger = logging.getLogger(__name__)
 
@@ -75,6 +76,8 @@ class Emotes(commands.Cog):
 					+ self.bot.http.user_agent
 			})
 
+		self.emote_client = EmoteClient(token=self.bot.config['tokens']['discord'])
+
 		with open('data/ec-emotes-final.json') as f:
 			self.ec_emotes = json.load(f)
 
@@ -84,6 +87,7 @@ class Emotes(commands.Cog):
 	def cog_unload(self):
 		async def close():
 			await self.http.close()
+			await self.emote_client.close()
 
 			for paginator in self.paginators:
 				await paginator.stop()
@@ -326,6 +330,8 @@ class Emotes(commands.Cog):
 		if not url and not context.message.attachments:
 			raise commands.BadArgument('A URL or attachment must be given.')
 
+		self.emote_client.check_create(context.guild.id)
+
 		url = url or context.message.attachments[0].url
 		async with context.typing():
 			archive = await self.fetch_safe(url, valid_mimetypes=self.ARCHIVE_MIMETYPES)
@@ -363,12 +369,13 @@ class Emotes(commands.Cog):
 
 	async def add_safe(self, context, name, url, author_id, *, reason=None):
 		"""Try to add an emote. Returns a string that should be sent to the user."""
+		self.emote_client.check_create(context.guild.id)
 		try:
 			image_data = await self.fetch_safe(url)
 		except errors.InvalidFileError:
 			raise errors.InvalidImageError
 
-		if type(image_data) is str:  # error case
+		if type(image_data) is str:  # error case (shitty i know)
 			return image_data
 		return await self.add_safe_bytes(context, name, author_id, image_data, reason=reason)
 
@@ -436,8 +443,8 @@ class Emotes(commands.Cog):
 	async def create_emote_from_bytes(self, guild, name, author_id, image_data: bytes, *, reason=None):
 		image_data = await utils.image.resize_in_subprocess(image_data)
 		if reason is None:
-			reason = f'Created by {utils.format_user(self.bot, author_id)}'
+			reason = 'Created by ' + utils.format_user(self.bot, author_id)
-		return await guild.create_custom_emoji(name=name, image=image_data, reason=reason)
+		return await self.emote_client.create(guild_id=guild.id, name=name, image=image_data, reason=reason)
 
 	@commands.command(aliases=('delete', 'delet', 'rm'))
 	async def remove(self, context, emote, *emotes):
@@ -447,7 +454,11 @@ class Emotes(commands.Cog):
 		"""
 		if not emotes:
 			emote = await self.parse_emote(context, emote)
-			await emote.delete(reason=f'Removed by {utils.format_user(self.bot, context.author.id)}')
+			await self.emote_client.delete(
+				guild_id=context.guild.id,
+				emote_id=emote.id,
+				reason='Removed by ' + utils.format_user(self.bot, context.author.id),
+			)
 			await context.send(fr'Emote \:{emote.name}: successfully removed.')
 		else:
 			for emote in (emote,) + emotes:

utils/__init__.py

@@ -5,5 +5,6 @@ from . import archive
 from . import emote
 from . import errors
 from . import paginator
+from . import emote_client
 # note: do not import .image in case the user doesn't want it
 # since importing image can take a long time.

utils/emote_client.py

@@ -0,0 +1,156 @@
+# © lambda#0987
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+import sys
+import json
+import asyncio
+import aiohttp
+import platform
+import datetime
+import urllib.parse
+from typing import Dict
+from http import HTTPStatus
+import utils.image as image_utils
+from utils.errors import RateLimitedError
+from discord import HTTPException, Forbidden, NotFound, DiscordServerError
+
+class GuildRetryTimes:
+	"""Holds the times, for a particular guild,
+	that we have to wait until for the rate limit for a particular HTTP method to elapse.
+	"""
+	__slots__ = frozenset({'POST', 'DELETE'})
+
+	def __init__(self, POST=None, DELETE=None):
+		self.POST = POST
+		self.DELETE = DELETE
+
+	def validate(self):
+		now = datetime.datetime.now(tz=datetime.timezone.utc).timestamp()
+		if self.POST and self.POST < now:
+			self.POST = None
+		if self.DELETE and self.DELETE < now:
+			self.DELETE = None
+		return self.POST or self.DELETE
+
+GuildId = int
+
+async def json_or_text(resp):
+	text = await resp.text(encoding='utf-8')
+	try:
+		if resp.headers['content-type'] == 'application/json':
+			return json.loads(text)
+	except KeyError:
+		# Thanks Cloudflare
+		pass
+
+	return text
+
+class EmoteClient:
+	BASE_URL = 'https://discord.com/api/v7'
+	HTTP_ERROR_CLASSES = {
+		HTTPStatus.FORBIDDEN: Forbidden,
+		HTTPStatus.NOT_FOUND: NotFound,
+		HTTPStatus.SERVICE_UNAVAILABLE: DiscordServerError,
+	}
+
+	def __init__(self, *, token):
+		self.guild_rls: Dict[GuildId, GuildRetryTimes] = {}
+		user_agent = (
+			'EmoteManager-EmoteClient; '
+			f'aiohttp/{aiohttp.__version__}; '
+			f'{platform.python_implementation()}/{".".join(map(str, sys.version_info))}'
+		)
+		self.http = aiohttp.ClientSession(headers={
+			'User-Agent': user_agent,
+			'Authorization': 'Bot ' + token,
+			'X-Ratelimit-Precision': 'millisecond',
+		})
+
+	async def request(self, method, path, guild_id, **kwargs):
+		self._check_rl(method, guild_id)
+		print('post check rl')
+
+		headers = {}
+		# Emote Manager shouldn't use walrus op until Debian adopts 3.8 :(
+		reason = kwargs.pop('reason', None)
+		if reason:
+			headers['X-Audit-Log-Reason'] = urllib.parse.quote(reason, safe='/ ')
+		kwargs['headers'] = headers
+
+		# TODO handle OSError and 500/502, like dpy does
+		async with self.http.request(method, self.BASE_URL + path, **kwargs) as resp:
+			if resp.status == HTTPStatus.TOO_MANY_REQUESTS:
+				print('handling rl')
+				return await self._handle_rl(resp, method, path, guild_id, **kwargs)
+
+			data = await json_or_text(resp)
+			if resp.status in range(200, 300):
+				return data
+
+			error_cls = self.HTTP_ERROR_CLASSES.get(resp.status, HTTPException)
+			raise error_cls(resp, data)
+
+	def _check_rl(self, method, guild_id):
+		try:
+			rls = self.guild_rls[guild_id]
+		except KeyError:
+			print('guild not found', repr(guild_id))
+			return
+
+		if not rls.validate():
+			print('guild rls invalid')
+			del self.guild_rls[guild_id]
+			return
+
+		retry_at = getattr(rls, method, None)
+		if retry_at:
+			print('retry later')
+			raise RateLimitedError(retry_at)
+
+	async def _handle_rl(self, resp, method, path, guild_id, **kwargs):
+		retry_after = (await resp.json())['retry_after'] / 1000.0
+		retry_at = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(seconds=retry_after)
+
+		# cache unconditionally in case request() is called again while we're sleeping
+		try:
+			rls = self.guild_rls[guild_id]
+		except KeyError:
+			self.guild_rls[guild_id] = rls = GuildRetryTimes()
+
+		setattr(rls, resp.method, retry_at.timestamp())
+
+		if resp.method not in GuildRetryTimes.__slots__ or retry_after < 10.0:
+			await asyncio.sleep(retry_after)
+			# woo mutual recursion
+			return await self.request(method, path, guild_id, **kwargs)
+
+		# we've been hit with one of those crazy high rate limits, which only occur for specific methods
+		raise RateLimitedError(retry_at)
+
+	# optimization methods that let us check the RLs before downloading the user's image
+	def check_create(self, guild_id):
+		self._check_rl('POST', guild_id)
+
+	def check_delete(self, guild_id):
+		self._check_rl('DELETE', guild_id)
+
+	async def create(self, *, guild_id, name, image: bytes, role_ids=(), reason=None):
+		return await self.request(
+			'POST', f'/guilds/{guild_id}/emojis',
+			guild_id,
+			json=dict(name=name, image=image_utils.image_to_base64_url(image), roles=role_ids),
+			reason=reason,
+		)
+
+	async def delete(self, *, guild_id, emote_id, reason=None):
+		return await self.request('DELETE', f'/guilds/{guild_id}/emojis/{emote_id}', guild_id, reason=reason)
+
+	async def __aenter__(self):
+		self.http = await self.http.__aenter__()
+		return self
+
+	async def __aexit__(self, *excinfo):
+		return await self.http.__aexit__(*excinfo)
+
+	async def close(self):
+		return await self.http.close()

utils/errors.py

@@ -13,11 +13,11 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with Emote Manager. If not, see <https://www.gnu.org/licenses/>.
 
-import asyncio
-
-from discord.ext import commands
-
 import utils
+import asyncio
+import humanize
+import datetime
+from discord.ext import commands
 
 class MissingManageEmojisPermission(commands.MissingPermissions):
 	"""The invoker or the bot doesn't have permissions to manage server emojis."""
@@ -45,10 +45,19 @@ class ImageConversionTimeoutError(ImageProcessingTimeoutError):
 		super().__init__('Error: converting the image to a GIF took too long.')
 
 class HTTPException(EmoteManagerError):
-	"""The server did not respond with an OK status code."""
+	"""The server did not respond with an OK status code. This is only for non-Discord HTTP requests."""
 	def __init__(self, status):
 		super().__init__(f'URL error: server returned error code {status}')
 
+class RateLimitedError(EmoteManagerError):
+	def __init__(self, retry_at):
+		if isinstance(retry_at, float):
+			# it took me about an HOUR to realize i had to pass tz because utcfromtimestamp returns a NAÏVE time obj!
+			retry_at = datetime.datetime.fromtimestamp(retry_at, tz=datetime.timezone.utc)
+		# humanize.naturaltime is annoying to work with due to timezones so we use this
+		delta = humanize.naturaldelta(retry_at, when=datetime.datetime.now(tz=datetime.timezone.utc))
+		super().__init__(f'Error: Discord told me to slow down! Please retry this command in {delta}.')
+
 class EmoteNotFoundError(EmoteManagerError):
 	"""An emote with that name was not found"""
 	def __init__(self, name):