Yadciel/Titan
1
0
Fork 0
mirror of https://github.com/TitanEmbeds/Titan.git synced 2024-12-24 14:07:03 +01:00
Code Issues Packages Projects Releases Wiki Activity

Use redis for discordapi bucket

Browse Source
This commit is contained in:
Jeremy Zhang 2017-04-22 05:10:46 -07:00
parent 1fe074a346
commit c497760f43
1 changed files with 26 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
titanembeds/discordrest.py
View File

@ -4,6 +4,7 @@ import time
import json import json
from functools import partial from functools import partial
from titanembeds.utils import cache from titanembeds.utils import cache
from redislite import Redis
_DISCORD_API_BASE = "https://discordapp.com/api/v6" _DISCORD_API_BASE = "https://discordapp.com/api/v6"
@ -15,11 +16,24 @@ def json_or_text(response):
class DiscordREST: class DiscordREST:
def __init__(self, bot_token): def __init__(self, bot_token):
self.global_redis_prefix = "discordapiratelimit/"
self.bot_token = bot_token self.bot_token = bot_token
self.user_agent = "TitanEmbeds (https://github.com/EndenDragon/Titan) Python/{} requests/{}".format(sys.version_info, requests.__version__) self.user_agent = "TitanEmbeds (https://github.com/EndenDragon/Titan) Python/{} requests/{}".format(sys.version_info, requests.__version__)
self.rate_limit_bucket = {} self.rate_limit_bucket = Redis("redislite.db")
self.global_limited = False
self.global_limit_expire = 0 if not self._bucket_contains("global_limited"):
self._set_bucket("global_limited", False)
self._set_bucket("global_limit_expire", 0)
def _get_bucket(self, key):
value = self.rate_limit_bucket.get(self.global_redis_prefix + key)
return value
def _set_bucket(self, key, value):
return self.rate_limit_bucket.set(self.global_redis_prefix + key, value)
def _bucket_contains(self, key):
return self.rate_limit_bucket.exists(self.global_redis_prefix + key)
def request(self, verb, url, **kwargs): def request(self, verb, url, **kwargs):
headers = { headers = {
@ -38,12 +52,12 @@ class DiscordREST:
for tries in range(5): for tries in range(5):
curepoch = time.time() curepoch = time.time()
if self.global_limited: if self._get_bucket("global_limited") == "True":
time.sleep(self.global_limit_expire - curepoch) time.sleep(int(self._get_bucket("global_limit_expire")) - curepoch)
curepoch = time.time() curepoch = time.time()
if url in self.rate_limit_bucket and self.rate_limit_bucket[url] > curepoch: if self._bucket_contains(url) and int(self._get_bucket(url)) > curepoch:
time.sleep(self.rate_limit_bucket[url] - curepoch) time.sleep(int(self._get_bucket(url)) - curepoch)
url_formatted = _DISCORD_API_BASE + url url_formatted = _DISCORD_API_BASE + url
req = requests.request(verb, url_formatted, params=params, data=data, headers=headers) req = requests.request(verb, url_formatted, params=params, data=data, headers=headers)
@ -52,10 +66,10 @@ class DiscordREST:
if 'X-RateLimit-Remaining' in req.headers: if 'X-RateLimit-Remaining' in req.headers:
remaining = req.headers['X-RateLimit-Remaining'] remaining = req.headers['X-RateLimit-Remaining']
if remaining == '0' and req.status_code != 429: if remaining == '0' and req.status_code != 429:
self.rate_limit_bucket[url] = int(req.headers['X-RateLimit-Reset']) self._set_bucket(url, int(req.headers['X-RateLimit-Reset']))
if 300 > req.status_code >= 200: if 300 > req.status_code >= 200:
self.global_limited = False self._set_bucket("global_limited", False)
return { return {
'success': True, 'success': True,
'content': json_or_text(req), 'content': json_or_text(req),
@ -64,9 +78,10 @@ class DiscordREST:
if req.status_code == 429: if req.status_code == 429:
if 'X-RateLimit-Global' not in req.headers: if 'X-RateLimit-Global' not in req.headers:
self.rate_limit_bucket[url] = int(req.headers['X-RateLimit-Reset']) self._set_bucket(url, int(req.headers['X-RateLimit-Reset']))
else: else:
self.global_limit_expire = time.time() + int(req.headers['Retry-After']) self._set_bucket("global_limited", True)
self._set_bucket("global_limit_expire", time.time() + int(req.headers['Retry-After']))
if req.status_code == 502 and tries <= 5: if req.status_code == 502 and tries <= 5:
time.sleep(1 + tries * 2) time.sleep(1 + tries * 2)