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.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()
|
||||||
|
|
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue