Discord API rate limiting (#799)

Added discord too many requests handling decorator
Added tests for core Discord manager functions
Added discord backoff retry for celery
Added tests for update groups backoff
Support per-route and global rate limiting

alliance_auth/tests/test_settings.py

@@ -398,12 +398,12 @@ TEAMSPEAK3_PUBLIC_URL = os.environ.get('AA_TEAMSPEAK3_PUBLIC_URL', 'example.com'
 # DISCORD_CALLBACK_URL - oauth callback url
 # DISCORD_SYNC_NAMES - enable to force discord nicknames to be set to eve char name (bot needs Manage Nicknames permission)
 ######################################
-DISCORD_GUILD_ID = os.environ.get('AA_DISCORD_GUILD_ID', '')
-DISCORD_BOT_TOKEN = os.environ.get('AA_DISCORD_BOT_TOKEN', '')
-DISCORD_INVITE_CODE = os.environ.get('AA_DISCORD_INVITE_CODE', '')
-DISCORD_APP_ID = os.environ.get('AA_DISCORD_APP_ID', '')
-DISCORD_APP_SECRET = os.environ.get('AA_DISCORD_APP_SECRET', '')
-DISCORD_CALLBACK_URL = os.environ.get('AA_DISCORD_CALLBACK_URL', 'http://example.com/discord_callback')
+DISCORD_GUILD_ID = os.environ.get('AA_DISCORD_GUILD_ID', '0118999')
+DISCORD_BOT_TOKEN = os.environ.get('AA_DISCORD_BOT_TOKEN', 'bottoken')
+DISCORD_INVITE_CODE = os.environ.get('AA_DISCORD_INVITE_CODE', 'invitecode')
+DISCORD_APP_ID = os.environ.get('AA_DISCORD_APP_ID', 'appid')
+DISCORD_APP_SECRET = os.environ.get('AA_DISCORD_APP_SECRET', 'secret')
+DISCORD_CALLBACK_URL = os.environ.get('AA_DISCORD_CALLBACK_URL', 'http://example.com/discord/callback')
 DISCORD_SYNC_NAMES = 'True' == os.environ.get('AA_DISCORD_SYNC_NAMES', 'False')
 
 ######################################

services/modules/discord/manager.py

@@ -5,9 +5,12 @@ import re
 from django.conf import settings
 from services.models import GroupCache
 from requests_oauthlib import OAuth2Session
+from functools import wraps
 import logging
 import datetime
+import time
 from django.utils import timezone
+from django.core.cache import cache
 
 logger = logging.getLogger(__name__)
 
@@ -31,6 +34,118 @@ SCOPES = [
 GROUP_CACHE_MAX_AGE = datetime.timedelta(minutes=30)
 
 
+class DiscordApiException(Exception):
+    def __init__(self):
+        super(Exception, self).__init__()
+
+
+class DiscordApiTooBusy(DiscordApiException):
+    def __init__(self):
+        super(DiscordApiException, self).__init__()
+        self.message = "The Discord API is too busy to process this request now, please try again later."
+
+
+class DiscordApiBackoff(DiscordApiException):
+    def __init__(self, retry_after, global_ratelimit):
+        super(DiscordApiException, self).__init__()
+        self.retry_after = retry_after
+        self.global_ratelimit = global_ratelimit
+
+
+cache_time_format = '%Y-%m-%d %H:%M:%S'
+
+
+def api_backoff(func):
+    """
+    Decorator, Handles HTTP 429 "Too Many Requests" messages from the Discord API
+    If blocking=True is specified, this function will block and retry
+    the function up to max_retries=n times, or 3 if retries is not specified.
+    If the API call still recieves a backoff timer this function will raise
+    a <DiscordApiTooBusy> exception.
+    If the caller chooses blocking=False, the decorator will raise a DiscordApiBackoff
+    exception and the caller can choose to retry after the given timespan available in
+    the retry_after property in seconds.
+    """
+
+    class PerformBackoff(Exception):
+        def __init__(self, retry_after, retry_datetime, global_ratelimit):
+            super(Exception, self).__init__()
+            self.retry_after = int(retry_after)
+            self.retry_datetime = retry_datetime
+            self.global_ratelimit = global_ratelimit
+
+    @wraps(func)
+    def decorated(*args, **kwargs):
+        blocking = kwargs.get('blocking', False)
+        retries = kwargs.get('max_retries', 3)
+
+        # Strip our parameters
+        if 'max_retries' in kwargs:
+            del kwargs['max_retries']
+        if 'blocking' in kwargs:
+            del kwargs['blocking']
+
+        cache_key = 'DISCORD_BACKOFF_' + func.__name__
+        cache_global_key = 'DISCORD_BACKOFF_GLOBAL'
+
+        while retries > 0:
+            try:
+                try:
+                    # Check global backoff first, then route backoff
+                    existing_global_backoff = cache.get(cache_global_key)
+                    existing_backoff = existing_global_backoff or cache.get(cache_key)
+                    if existing_backoff:
+                        backoff_timer = datetime.datetime.strptime(existing_backoff, cache_time_format)
+                        if backoff_timer > datetime.datetime.utcnow():
+                            backoff_seconds = (backoff_timer - datetime.datetime.utcnow()).total_seconds()
+                            logger.debug("Still under backoff for {} seconds, backing off" % backoff_seconds)
+                            # Still under backoff
+                            raise PerformBackoff(
+                                retry_after=backoff_seconds,
+                                retry_datetime=backoff_timer,
+                                global_ratelimit=bool(existing_global_backoff)
+                            )
+                    logger.debug("Calling API calling function")
+                    func(*args, **kwargs)
+                    break
+                except requests.HTTPError as e:
+                    if e.response.status_code == 429:
+                        if 'Retry-After' in e.response.headers:
+                            retry_after = e.response.headers['Retry-After']
+                        else:
+                            # Pick some random time
+                            retry_after = 5
+
+                        logger.info("Received backoff from API of %s seconds, handling" % retry_after)
+                        # Store value in redis
+                        backoff_until = (datetime.datetime.utcnow() +
+                                         datetime.timedelta(seconds=int(retry_after)))
+                        global_backoff = bool(e.response.headers.get('X-RateLimit-Global', False))
+                        if global_backoff:
+                            logger.info("Global backoff!!")
+                            cache.set(cache_global_key, backoff_until.strftime(cache_time_format), retry_after)
+                        else:
+                            cache.set(cache_key, backoff_until.strftime(cache_time_format), retry_after)
+                        raise PerformBackoff(retry_after=retry_after, retry_datetime=backoff_until,
+                                             global_ratelimit=global_backoff)
+                    else:
+                        # Not 429, re-raise
+                        raise e
+            except PerformBackoff as bo:
+                # Sleep if we're blocking
+                if blocking:
+                    logger.info("Blocking Back off from API calls for %s seconds" % bo.retry_after)
+                    time.sleep(10 if bo.retry_after > 10 else bo.retry_after)
+                else:
+                    # Otherwise raise exception and let caller handle the backoff
+                    raise DiscordApiBackoff(retry_after=bo.retry_after, global_ratelimit=bo.global_ratelimit)
+            finally:
+                retries -= 1
+        if retries == 0:
+            raise DiscordApiTooBusy()
+    return decorated
+
+
 class DiscordOAuthManager:
     def __init__(self):
         pass
@@ -191,6 +306,7 @@ class DiscordOAuthManager:
         DiscordOAuthManager.__update_group_cache()
 
     @staticmethod
+    @api_backoff
     def update_groups(user_id, groups):
         custom_headers = {'content-type': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN}
         group_ids = [DiscordOAuthManager.__group_name_to_id(DiscordOAuthManager._sanitize_groupname(g)) for g in groups]
@@ -199,3 +315,4 @@ class DiscordOAuthManager:
         r = requests.patch(path, headers=custom_headers, json=data)
         logger.debug("Received status code %s after setting user roles" % r.status_code)
         r.raise_for_status()
+

services/modules/discord/tasks.py

@@ -9,7 +9,7 @@ from django.core.exceptions import ObjectDoesNotExist
 
 from eveonline.managers import EveManager
 from notifications import notify
-from services.modules.discord.manager import DiscordOAuthManager
+from services.modules.discord.manager import DiscordOAuthManager, DiscordApiBackoff
 from services.tasks import only_one
 from .models import DiscordUser
 
@@ -74,6 +74,10 @@ class DiscordTasks:
             logger.debug("Updating user %s discord groups to %s" % (user, groups))
             try:
                 DiscordOAuthManager.update_groups(user.discord.uid, groups)
+            except DiscordApiBackoff as bo:
+                logger.info("Discord group sync API back off for %s, "
+                            "retrying in %s seconds" % (user, bo.retry_after))
+                raise task_self.retry(countdown=bo.retry_after)
             except Exception as e:
                 if task_self:
                     logger.exception("Discord group sync failed for %s, retrying in 10 mins" % user)

services/modules/discord/tests.py

@@ -17,6 +17,10 @@ from alliance_auth.tests.auth_utils import AuthUtils
 from .auth_hooks import DiscordService
 from .models import DiscordUser
 from .tasks import DiscordTasks
+from .manager import DiscordOAuthManager
+
+import requests_mock
+import datetime
 
 MODULE_PATH = 'services.modules.discord'
 
@@ -198,3 +202,254 @@ class DiscordViewsTestCase(TestCase):
         self.assertRedirects(response, expected_url='/en/services/', target_status_code=200)
         with self.assertRaises(ObjectDoesNotExist):
             discord_user = User.objects.get(pk=self.member.pk).discord
+
+
+class DiscordManagerTestCase(TestCase):
+    def setUp(self):
+        pass
+
+    def test__sanitize_groupname(self):
+        test_group_name = ' Group Name_Test_'
+        group_name = DiscordOAuthManager._sanitize_groupname(test_group_name)
+
+        self.assertEqual(group_name, 'GroupName_Test')
+
+    def test_generate_Bot_add_url(self):
+        from . import manager
+        bot_add_url = DiscordOAuthManager.generate_bot_add_url()
+
+        auth_url = manager.AUTH_URL
+        real_bot_add_url = '{}?client_id=appid&scope=bot&permissions={}'.format(auth_url, manager.BOT_PERMISSIONS)
+        self.assertEqual(bot_add_url, real_bot_add_url)
+
+    def test_generate_oauth_redirect_url(self):
+        from . import manager
+        import urllib
+        import sys
+        oauth_url = DiscordOAuthManager.generate_oauth_redirect_url()
+
+        self.assertIn(manager.AUTH_URL, oauth_url)
+        self.assertIn('+'.join(manager.SCOPES), oauth_url)
+        self.assertIn(settings.DISCORD_APP_ID, oauth_url)
+        if sys.version_info[0] < 3:
+            # Py2
+            self.assertIn(urllib.quote_plus(settings.DISCORD_CALLBACK_URL), oauth_url)
+        else: # Py3
+            self.assertIn(urllib.parse.quote_plus(settings.DISCORD_CALLBACK_URL), oauth_url)
+
+    @mock.patch(MODULE_PATH + '.manager.OAuth2Session')
+    def test__process_callback_code(self, oauth):
+        from . import manager
+        instance = oauth.return_value
+        instance.fetch_token.return_value = {'access_token': 'mywonderfultoken'}
+
+        token = DiscordOAuthManager._process_callback_code('12345')
+
+        self.assertTrue(oauth.called)
+        args, kwargs = oauth.call_args
+        self.assertEqual(args[0], settings.DISCORD_APP_ID)
+        self.assertEqual(kwargs['redirect_uri'], settings.DISCORD_CALLBACK_URL)
+        self.assertTrue(instance.fetch_token.called)
+        args, kwargs = instance.fetch_token.call_args
+        self.assertEqual(args[0], manager.TOKEN_URL)
+        self.assertEqual(kwargs['client_secret'], settings.DISCORD_APP_SECRET)
+        self.assertEqual(kwargs['code'], '12345')
+        self.assertEqual(token['access_token'], 'mywonderfultoken')
+
+    @mock.patch(MODULE_PATH + '.manager.DiscordOAuthManager._process_callback_code')
+    @requests_mock.Mocker()
+    def test_add_user(self, oauth_token, m):
+        from . import manager
+        import json
+
+        # Arrange
+        oauth_token.return_value = {'access_token': 'accesstoken'}
+
+        headers = {'accept': 'application/json', 'authorization': 'Bearer accesstoken'}
+
+        m.register_uri('POST',
+                       manager.DISCORD_URL + '/invites/'+str(settings.DISCORD_INVITE_CODE),
+                       request_headers=headers,
+                       text='{}')
+
+        m.register_uri('GET',
+                       manager.DISCORD_URL + "/users/@me",
+                       request_headers=headers,
+                       text=json.dumps({'id': "123456"}))
+
+        # Act
+        return_value = DiscordOAuthManager.add_user('abcdef')
+
+        # Assert
+        self.assertEqual(return_value, '123456')
+        self.assertEqual(m.call_count, 2)
+
+    @requests_mock.Mocker()
+    def test_delete_user(self, m):
+        from . import manager
+        import json
+
+        # Arrange
+        headers = {'accept': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN}
+
+        user_id = 12345
+        request_url = '{}/guilds/{}/members/{}'.format(manager.DISCORD_URL, settings.DISCORD_GUILD_ID, user_id)
+        m.register_uri('DELETE',
+                       request_url,
+                       request_headers=headers,
+                       text=json.dumps({}))
+
+        # Act
+        result = DiscordOAuthManager.delete_user(user_id)
+
+        # Assert
+        self.assertTrue(result)
+
+        ###
+        # Test 404 (already deleted)
+        # Arrange
+        m.register_uri('DELETE',
+                       request_url,
+                       request_headers=headers,
+                       status_code=404)
+
+        # Act
+        result = DiscordOAuthManager.delete_user(user_id)
+
+        # Assert
+        self.assertTrue(result)
+
+        ###
+        # Test 500 (some random API error)
+        # Arrange
+        m.register_uri('DELETE',
+                       request_url,
+                       request_headers=headers,
+                       status_code=500)
+
+        # Act
+        result = DiscordOAuthManager.delete_user(user_id)
+
+        # Assert
+        self.assertFalse(result)
+
+    @requests_mock.Mocker()
+    def test_update_nickname(self, m):
+        from . import manager
+        import json
+        # Arrange
+        headers = {'content-type': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN}
+
+        user_id = 12345
+        request_url = '{}/guilds/{}/members/{}'.format(manager.DISCORD_URL, settings.DISCORD_GUILD_ID, user_id)
+        m.patch(request_url,
+                request_headers=headers)
+
+        # Act
+        result = DiscordOAuthManager.update_nickname(user_id, 'somenick')
+
+        # Assert
+        self.assertTrue(result)
+
+    @mock.patch(MODULE_PATH + '.manager.DiscordOAuthManager._DiscordOAuthManager__get_group_cache')
+    @requests_mock.Mocker()
+    def test_update_groups(self, group_cache, m):
+        from . import manager
+        import json
+
+        # Arrange
+        groups = ['Member', 'Blue', 'Special Group']
+
+        group_cache.return_value = [{'id': 111, 'name': 'Member'},
+                                    {'id': 222, 'name': 'Blue'},
+                                    {'id': 333, 'name': 'SpecialGroup'},
+                                    {'id': 444, 'name': 'NotYourGroup'}]
+
+        headers = {'content-type': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN}
+        user_id = 12345
+        request_url = '{}/guilds/{}/members/{}'.format(manager.DISCORD_URL, settings.DISCORD_GUILD_ID, user_id)
+
+        m.patch(request_url,
+                request_headers=headers)
+
+        # Act
+        DiscordOAuthManager.update_groups(user_id, groups)
+
+        # Assert
+        self.assertEqual(len(m.request_history), 1, 'Must be one HTTP call made')
+        history = json.loads(m.request_history[0].text)
+        self.assertIn('roles', history, "'The request must send JSON object with the 'roles' key")
+        self.assertIn(111, history['roles'], 'The group id 111 must be added to the request')
+        self.assertIn(222, history['roles'], 'The group id 222 must be added to the request')
+        self.assertIn(333, history['roles'], 'The group id 333 must be added to the request')
+        self.assertNotIn(444, history['roles'], 'The group id 444 must NOT be added to the request')
+
+    @mock.patch(MODULE_PATH + '.manager.cache')
+    @mock.patch(MODULE_PATH + '.manager.DiscordOAuthManager._DiscordOAuthManager__get_group_cache')
+    @requests_mock.Mocker()
+    def test_update_groups_backoff(self, group_cache, djcache, m):
+        from . import manager
+
+        # Arrange
+        groups = ['Member']
+        group_cache.return_value = [{'id': 111, 'name': 'Member'}]
+
+        headers = {'content-type': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN}
+        user_id = 12345
+        request_url = '{}/guilds/{}/members/{}'.format(manager.DISCORD_URL, settings.DISCORD_GUILD_ID, user_id)
+
+        djcache.get.return_value = None  # No existing backoffs in cache
+
+        m.patch(request_url,
+                request_headers=headers,
+                headers={'Retry-After': '200'},
+                status_code=429)
+
+        # Act & Assert
+        with self.assertRaises(manager.DiscordApiBackoff) as bo:
+            try:
+                DiscordOAuthManager.update_groups(user_id, groups, blocking=False)
+            except manager.DiscordApiBackoff as bo:
+                self.assertEqual(bo.retry_after, 200, 'Retry-After time must be equal to Retry-After set in header')
+                self.assertFalse(bo.global_ratelimit, 'global_ratelimit must be False')
+                raise bo
+
+        self.assertTrue(djcache.set.called)
+        args, kwargs = djcache.set.call_args
+        self.assertEqual(args[0], 'DISCORD_BACKOFF_update_groups')
+        self.assertTrue(datetime.datetime.strptime(args[1], manager.cache_time_format) > datetime.datetime.now())
+
+    @mock.patch(MODULE_PATH + '.manager.cache')
+    @mock.patch(MODULE_PATH + '.manager.DiscordOAuthManager._DiscordOAuthManager__get_group_cache')
+    @requests_mock.Mocker()
+    def test_update_groups_global_backoff(self, group_cache, djcache, m):
+        from . import manager
+
+        # Arrange
+        groups = ['Member']
+        group_cache.return_value = [{'id': 111, 'name': 'Member'}]
+
+        headers = {'content-type': 'application/json', 'authorization': 'Bot ' + settings.DISCORD_BOT_TOKEN}
+        user_id = 12345
+        request_url = '{}/guilds/{}/members/{}'.format(manager.DISCORD_URL, settings.DISCORD_GUILD_ID, user_id)
+
+        djcache.get.return_value = None  # No existing backoffs in cache
+
+        m.patch(request_url,
+                request_headers=headers,
+                headers={'Retry-After': '200', 'X-RateLimit-Global': 'true'},
+                status_code=429)
+
+        # Act & Assert
+        with self.assertRaises(manager.DiscordApiBackoff) as bo:
+            try:
+                DiscordOAuthManager.update_groups(user_id, groups, blocking=False)
+            except manager.DiscordApiBackoff as bo:
+                self.assertEqual(bo.retry_after, 200, 'Retry-After time must be equal to Retry-After set in header')
+                self.assertTrue(bo.global_ratelimit, 'global_ratelimit must be True')
+                raise bo
+
+        self.assertTrue(djcache.set.called)
+        args, kwargs = djcache.set.call_args
+        self.assertEqual(args[0], 'DISCORD_BACKOFF_GLOBAL')
+        self.assertTrue(datetime.datetime.strptime(args[1], manager.cache_time_format) > datetime.datetime.now())