Discord Rate Limiting (429 Error code) handling (#901)

apprise/plugins/NotifyDiscord.py

@@ -50,6 +50,9 @@
 import re
 import requests
 from json import dumps
+from datetime import timedelta
+from datetime import datetime
+from datetime import timezone
 
 from .NotifyBase import NotifyBase
 from ..common import NotifyImageSize
@@ -84,6 +87,17 @@ class NotifyDiscord(NotifyBase):
     # Allows the user to specify the NotifyImageSize object
     image_size = NotifyImageSize.XY_256
 
+    # Discord is kind enough to return how many more requests we're allowed to
+    # continue to make within it's header response as:
+    # X-RateLimit-Reset: The epoc time (in seconds) we can expect our
+    #                    rate-limit to be reset.
+    # X-RateLimit-Remaining: an integer identifying how many requests we're
+    #                        still allow to make.
+    request_rate_per_sec = 0
+
+    # Taken right from google.auth.helpers:
+    clock_skew = timedelta(seconds=10)
+
     # The maximum allowable characters allowed in the body per message
     body_maxlen = 2000
 
@@ -215,6 +229,12 @@ class NotifyDiscord(NotifyBase):
         # dynamically generated avatar url images
         self.avatar_url = avatar_url
 
+        # For Tracking Purposes
+        self.ratelimit_reset = datetime.now(timezone.utc).replace(tzinfo=None)
+
+        # Default to 1.0
+        self.ratelimit_remaining = 1.0
+
         return
 
     def send(self, body, title='', notify_type=NotifyType.INFO, attach=None,
@@ -343,7 +363,8 @@ class NotifyDiscord(NotifyBase):
         # Otherwise return
         return True
 
-    def _send(self, payload, attach=None, params=None, **kwargs):
+    def _send(self, payload, attach=None, params=None, rate_limit=1,
+              **kwargs):
         """
         Wrapper to the requests (post) object
         """
@@ -365,8 +386,25 @@ class NotifyDiscord(NotifyBase):
         ))
         self.logger.debug('Discord Payload: %s' % str(payload))
 
-        # Always call throttle before any remote server i/o is made
-        self.throttle()
+        # By default set wait to None
+        wait = None
+
+        if self.ratelimit_remaining <= 0.0:
+            # Determine how long we should wait for or if we should wait at
+            # all. This isn't fool-proof because we can't be sure the client
+            # time (calling this script) is completely synced up with the
+            # Gitter server.  One would hope we're on NTP and our clocks are
+            # the same allowing this to role smoothly:
+
+            now = datetime.now(timezone.utc).replace(tzinfo=None)
+            if now < self.ratelimit_reset:
+                # We need to throttle for the difference in seconds
+                wait = abs(
+                    (self.ratelimit_reset - now + self.clock_skew)
+                    .total_seconds())
+
+        # Always call throttle before any remote server i/o is made;
+        self.throttle(wait=wait)
 
         # Perform some simple error checking
         if isinstance(attach, AttachBase):
@@ -401,6 +439,22 @@ class NotifyDiscord(NotifyBase):
                 verify=self.verify_certificate,
                 timeout=self.request_timeout,
             )
+
+            # Handle rate limiting (if specified)
+            try:
+                # Store our rate limiting (if provided)
+                self.ratelimit_remaining = \
+                    float(r.headers.get(
+                        'X-RateLimit-Remaining'))
+                self.ratelimit_reset = datetime.fromtimestamp(
+                    int(r.headers.get('X-RateLimit-Reset')),
+                    timezone.utc).replace(tzinfo=None)
+
+            except (TypeError, ValueError):
+                # This is returned if we could not retrieve this
+                # information gracefully accept this state and move on
+                pass
+
             if r.status_code not in (
                     requests.codes.ok, requests.codes.no_content):
 
@@ -408,6 +462,20 @@ class NotifyDiscord(NotifyBase):
                 status_str = \
                     NotifyBase.http_response_code_lookup(r.status_code)
 
+                if r.status_code == requests.codes.too_many_requests \
+                        and rate_limit > 0:
+
+                    # handle rate limiting
+                    self.logger.warning(
+                        'Discord rate limiting in effect; '
+                        'blocking for %.2f second(s)',
+                        self.ratelimit_remaining)
+
+                    # Try one more time before failing
+                    return self._send(
+                        payload=payload, attach=attach, params=params,
+                        rate_limit=rate_limit - 1, **kwargs)
+
                 self.logger.warning(
                     'Failed to send {}to Discord notification: '
                     '{}{}error={}.'.format(

test/test_plugin_discord.py

@@ -32,7 +32,8 @@
 
 import os
 from unittest import mock
-
+from datetime import datetime, timedelta
+from datetime import timezone
 import pytest
 import requests
 
@@ -182,6 +183,11 @@ def test_plugin_discord_general(mock_post):
 
     """
 
+    # Turn off clock skew for local testing
+    NotifyDiscord.clock_skew = timedelta(seconds=0)
+    # Epoch time:
+    epoch = datetime.fromtimestamp(0, timezone.utc)
+
     # Initialize some generic (but valid) tokens
     webhook_id = 'A' * 24
     webhook_token = 'B' * 64
@@ -189,6 +195,12 @@ def test_plugin_discord_general(mock_post):
     # Prepare Mock
     mock_post.return_value = requests.Request()
     mock_post.return_value.status_code = requests.codes.ok
+    mock_post.return_value.content = ''
+    mock_post.return_value.headers = {
+        'X-RateLimit-Reset': (
+            datetime.now(timezone.utc) - epoch).total_seconds(),
+        'X-RateLimit-Remaining': 1,
+    }
 
     # Invalid webhook id
     with pytest.raises(TypeError):
@@ -208,10 +220,85 @@ def test_plugin_discord_general(mock_post):
         webhook_id=webhook_id,
         webhook_token=webhook_token,
         footer=True, thumbnail=False)
+    assert obj.ratelimit_remaining == 1
 
     # Test that we get a string response
     assert isinstance(obj.url(), str) is True
 
+    # This call includes an image with it's payload:
+    assert obj.notify(
+        body='body', title='title', notify_type=NotifyType.INFO) is True
+
+    # Force a case where there are no more remaining posts allowed
+    mock_post.return_value.headers = {
+        'X-RateLimit-Reset': (
+            datetime.now(timezone.utc) - epoch).total_seconds(),
+        'X-RateLimit-Remaining': 0,
+    }
+
+    # This call includes an image with it's payload:
+    assert obj.notify(
+        body='body', title='title', notify_type=NotifyType.INFO) is True
+
+    # behind the scenes, it should cause us to update our rate limit
+    assert obj.send(body="test") is True
+    assert obj.ratelimit_remaining == 0
+
+    # This should cause us to block
+    mock_post.return_value.headers = {
+        'X-RateLimit-Reset': (
+            datetime.now(timezone.utc) - epoch).total_seconds(),
+        'X-RateLimit-Remaining': 10,
+    }
+    assert obj.send(body="test") is True
+    assert obj.ratelimit_remaining == 10
+
+    # Reset our variable back to 1
+    mock_post.return_value.headers = {
+        'X-RateLimit-Reset': (
+            datetime.now(timezone.utc) - epoch).total_seconds(),
+        'X-RateLimit-Remaining': 1,
+    }
+    # Handle cases where our epoch time is wrong
+    del mock_post.return_value.headers['X-RateLimit-Reset']
+    assert obj.send(body="test") is True
+
+    # Return our object, but place it in the future forcing us to block
+    mock_post.return_value.headers = {
+        'X-RateLimit-Reset': (
+            datetime.now(timezone.utc) - epoch).total_seconds() + 1,
+        'X-RateLimit-Remaining': 0,
+    }
+
+    obj.ratelimit_remaining = 0
+    assert obj.send(body="test") is True
+
+    # Test 429 error response
+    mock_post.return_value.status_code = requests.codes.too_many_requests
+
+    # The below will attempt a second transmission and fail (because we didn't
+    # set up a second post request to pass) :)
+    assert obj.send(body="test") is False
+
+    # Return our object, but place it in the future forcing us to block
+    mock_post.return_value.status_code = requests.codes.ok
+    mock_post.return_value.headers = {
+        'X-RateLimit-Reset': (
+            datetime.now(timezone.utc) - epoch).total_seconds() - 1,
+        'X-RateLimit-Remaining': 0,
+    }
+    assert obj.send(body="test") is True
+
+    # Return our limits to always work
+    obj.ratelimit_remaining = 1
+
+    # Return our headers to normal
+    mock_post.return_value.headers = {
+        'X-RateLimit-Reset': (
+            datetime.now(timezone.utc) - epoch).total_seconds(),
+        'X-RateLimit-Remaining': 1,
+    }
+
     # This call includes an image with it's payload:
     assert obj.notify(
         body='body', title='title', notify_type=NotifyType.INFO) is True