rate_limiter: Move email_mirror limiter to use rate_limit_entity.

We change the rate limiting code in the email mirror to use the new,
general rate_limit_entity function.
This commit is contained in:
Mateusz Mandera 2019-03-23 18:50:05 +01:00 committed by Tim Abbott
parent 40763070b7
commit c1ceba9037
2 changed files with 11 additions and 16 deletions

View File

@ -18,10 +18,9 @@ from zerver.lib.email_notifications import convert_html_to_markdown
from zerver.lib.queue import queue_json_publish from zerver.lib.queue import queue_json_publish
from zerver.lib.redis_utils import get_redis_client from zerver.lib.redis_utils import get_redis_client
from zerver.lib.upload import upload_message_file from zerver.lib.upload import upload_message_file
from zerver.lib.utils import generate_random_token, statsd from zerver.lib.utils import generate_random_token
from zerver.lib.send_email import FromAddress from zerver.lib.send_email import FromAddress
from zerver.lib.rate_limiter import RateLimitedObject, RateLimiterLockingException, \ from zerver.lib.rate_limiter import RateLimitedObject, rate_limit_entity
is_ratelimited, incr_ratelimit
from zerver.lib.exceptions import RateLimited from zerver.lib.exceptions import RateLimited
from zerver.models import Stream, Recipient, \ from zerver.models import Stream, Recipient, \
get_user_profile_by_id, get_display_recipient, get_personal_recipient, \ get_user_profile_by_id, get_display_recipient, get_personal_recipient, \
@ -412,20 +411,12 @@ class RateLimitedRealmMirror(RateLimitedObject):
def rules(self) -> List[Tuple[int, int]]: def rules(self) -> List[Tuple[int, int]]:
return settings.RATE_LIMITING_MIRROR_REALM_RULES return settings.RATE_LIMITING_MIRROR_REALM_RULES
def __str__(self) -> str:
return self.realm.string_id
def rate_limit_mirror_by_realm(recipient_realm: Realm) -> None: def rate_limit_mirror_by_realm(recipient_realm: Realm) -> None:
# Code based on the rate_limit_user function:
entity = RateLimitedRealmMirror(recipient_realm) entity = RateLimitedRealmMirror(recipient_realm)
ratelimited, time = is_ratelimited(entity) ratelimited = rate_limit_entity(entity)[0]
if ratelimited: if ratelimited:
statsd.incr("ratelimiter.limited.%s.%s" % (type(recipient_realm),
recipient_realm.id))
raise RateLimited()
try:
incr_ratelimit(entity)
except RateLimiterLockingException:
logger.warning("Email mirror rate limiter: Deadlock trying to "
"incr_ratelimit for realm %s" % (recipient_realm.name,))
raise RateLimited() raise RateLimited()

View File

@ -237,9 +237,11 @@ class WorkerTest(ZulipTestCase):
self.assertEqual(mock_mirror_email.call_count, 3) self.assertEqual(mock_mirror_email.call_count, 3)
@patch('zerver.lib.rate_limiter.logger.warning')
@patch('zerver.worker.queue_processors.mirror_email') @patch('zerver.worker.queue_processors.mirror_email')
@override_settings(RATE_LIMITING_MIRROR_REALM_RULES=[(10, 2)]) @override_settings(RATE_LIMITING_MIRROR_REALM_RULES=[(10, 2)])
def test_mirror_worker_rate_limiting(self, mock_mirror_email: MagicMock) -> None: def test_mirror_worker_rate_limiting(self, mock_mirror_email: MagicMock,
mock_warn: MagicMock) -> None:
fake_client = self.FakeClient() fake_client = self.FakeClient()
realm = get_realm('zulip') realm = get_realm('zulip')
clear_history(RateLimitedRealmMirror(realm)) clear_history(RateLimitedRealmMirror(realm))
@ -289,11 +291,13 @@ class WorkerTest(ZulipTestCase):
self.assertEqual(mock_mirror_email.call_count, 4) self.assertEqual(mock_mirror_email.call_count, 4)
# If RateLimiterLockingException is thrown, we rate-limit the new message: # If RateLimiterLockingException is thrown, we rate-limit the new message:
with patch('zerver.lib.email_mirror.incr_ratelimit', with patch('zerver.lib.rate_limiter.incr_ratelimit',
side_effect=RateLimiterLockingException): side_effect=RateLimiterLockingException):
fake_client.queue.append(('email_mirror', data[0])) fake_client.queue.append(('email_mirror', data[0]))
worker.start() worker.start()
self.assertEqual(mock_mirror_email.call_count, 4) self.assertEqual(mock_mirror_email.call_count, 4)
expected_warn = "Deadlock trying to incr_ratelimit for RateLimitedRealmMirror:zulip"
mock_warn.assert_called_with(expected_warn)
def test_email_sending_worker_retries(self) -> None: def test_email_sending_worker_retries(self) -> None:
"""Tests the retry_send_email_failures decorator to make sure it """Tests the retry_send_email_failures decorator to make sure it