mirror of https://github.com/zulip/zulip.git
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:
parent
40763070b7
commit
c1ceba9037
|
@ -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.redis_utils import get_redis_client
|
||||
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.rate_limiter import RateLimitedObject, RateLimiterLockingException, \
|
||||
is_ratelimited, incr_ratelimit
|
||||
from zerver.lib.rate_limiter import RateLimitedObject, rate_limit_entity
|
||||
from zerver.lib.exceptions import RateLimited
|
||||
from zerver.models import Stream, 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]]:
|
||||
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:
|
||||
# Code based on the rate_limit_user function:
|
||||
entity = RateLimitedRealmMirror(recipient_realm)
|
||||
ratelimited, time = is_ratelimited(entity)
|
||||
ratelimited = rate_limit_entity(entity)[0]
|
||||
|
||||
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()
|
||||
|
|
|
@ -237,9 +237,11 @@ class WorkerTest(ZulipTestCase):
|
|||
|
||||
self.assertEqual(mock_mirror_email.call_count, 3)
|
||||
|
||||
@patch('zerver.lib.rate_limiter.logger.warning')
|
||||
@patch('zerver.worker.queue_processors.mirror_email')
|
||||
@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()
|
||||
realm = get_realm('zulip')
|
||||
clear_history(RateLimitedRealmMirror(realm))
|
||||
|
@ -289,11 +291,13 @@ class WorkerTest(ZulipTestCase):
|
|||
self.assertEqual(mock_mirror_email.call_count, 4)
|
||||
|
||||
# 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):
|
||||
fake_client.queue.append(('email_mirror', data[0]))
|
||||
worker.start()
|
||||
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:
|
||||
"""Tests the retry_send_email_failures decorator to make sure it
|
||||
|
|
Loading…
Reference in New Issue