diff --git a/zerver/lib/email_mirror.py b/zerver/lib/email_mirror.py index 381a462a40..bfe8952daa 100644 --- a/zerver/lib/email_mirror.py +++ b/zerver/lib/email_mirror.py @@ -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() diff --git a/zerver/tests/test_queue_worker.py b/zerver/tests/test_queue_worker.py index ba349c49b4..c486e480cd 100644 --- a/zerver/tests/test_queue_worker.py +++ b/zerver/tests/test_queue_worker.py @@ -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