Performance improvements and refactor of Ratelimiter

tests/api/test_ratelimiting.py

@@ -5,35 +5,35 @@
 
 class TestRatelimiter(unittest.TestCase):
     def test_allowed(self):
-        limiter = Ratelimiter()
+        limiter = Ratelimiter(rate_hz=0.1, burst_count=1)
         allowed, time_allowed = limiter.can_do_action(
-            key="test_id", time_now_s=0, rate_hz=0.1, burst_count=1
+            key="test_id", time_now_s=0
         )
         self.assertTrue(allowed)
         self.assertEquals(10.0, time_allowed)
 
         allowed, time_allowed = limiter.can_do_action(
-            key="test_id", time_now_s=5, rate_hz=0.1, burst_count=1
+            key="test_id", time_now_s=5
         )
         self.assertFalse(allowed)
         self.assertEquals(10.0, time_allowed)
 
         allowed, time_allowed = limiter.can_do_action(
-            key="test_id", time_now_s=10, rate_hz=0.1, burst_count=1
+            key="test_id", time_now_s=10
         )
         self.assertTrue(allowed)
         self.assertEquals(20.0, time_allowed)
 
     def test_pruning(self):
-        limiter = Ratelimiter()
-        allowed, time_allowed = limiter.can_do_action(
-            key="test_id_1", time_now_s=0, rate_hz=0.1, burst_count=1
+        limiter = Ratelimiter(rate_hz=0.1, burst_count=1)
+        _, _ = limiter.can_do_action(
+            key="test_id_1", time_now_s=0
         )
 
-        self.assertIn("test_id_1", limiter.message_counts)
+        self.assertIn("test_id_1", limiter.actions)
 
-        allowed, time_allowed = limiter.can_do_action(
-            key="test_id_2", time_now_s=10, rate_hz=0.1, burst_count=1
+        _, _ = limiter.can_do_action(
+            key="test_id_2", time_now_s=10
         )
 
-        self.assertNotIn("test_id_1", limiter.message_counts)
+        self.assertNotIn("test_id_1", limiter.actions)

tests/handlers/test_profile.py

@@ -55,11 +55,17 @@ def register_query_handler(query_type, handler):
             federation_client=self.mock_federation,
             federation_server=Mock(),
             federation_registry=self.mock_registry,
-            ratelimiter=NonCallableMock(spec_set=["can_do_action"]),
+            request_ratelimiter=NonCallableMock(
+                spec_set=["can_do_action", "ratelimit", "rate_hz", "burst_count"]
+            ),
+            login_ratelimiter=NonCallableMock(spec_set=["can_do_action", "ratelimit"]),
         )
 
-        self.ratelimiter = hs.get_ratelimiter()
-        self.ratelimiter.can_do_action.return_value = (True, 0)
+        self.request_ratelimiter = hs.get_request_ratelimiter()
+        self.request_ratelimiter.can_do_action.return_value = (True, 0)
+
+        self.login_ratelimiter = hs.get_login_ratelimiter()
+        self.login_ratelimiter.can_do_action.return_value = (True, 0)
 
         self.store = hs.get_datastore()
 

tests/replication/slave/storage/_base.py

@@ -23,10 +23,16 @@ def make_homeserver(self, reactor, clock):
 
         hs = self.setup_test_homeserver(
             federation_client=Mock(),
-            ratelimiter=NonCallableMock(spec_set=["can_do_action"]),
+            request_ratelimiter=NonCallableMock(spec_set=["can_do_action", "ratelimit"]),
+            login_ratelimiter=NonCallableMock(spec_set=["can_do_action", "ratelimit"]),
         )
 
-        hs.get_ratelimiter().can_do_action.return_value = (True, 0)
+        # Prevent ratelimiting
+        self.request_ratelimiter = hs.get_request_ratelimiter()
+        self.request_ratelimiter.can_do_action.return_value = (True, 0)
+
+        self.login_ratelimiter = hs.get_login_ratelimiter()
+        self.login_ratelimiter.can_do_action.return_value = (True, 0)
 
         return hs
 

tests/rest/client/v1/test_events.py

@@ -41,10 +41,17 @@ def make_homeserver(self, reactor, clock):
         config["auto_join_rooms"] = []
 
         hs = self.setup_test_homeserver(
-            config=config, ratelimiter=NonCallableMock(spec_set=["can_do_action"])
+            config=config,
+            request_ratelimiter=NonCallableMock(
+                spec_set=["can_do_action", "ratelimit", "rate_hz", "burst_count"]
+            ),
+            login_ratelimiter = NonCallableMock(spec_set=["can_do_action", "ratelimit"]),
         )
-        self.ratelimiter = hs.get_ratelimiter()
-        self.ratelimiter.can_do_action.return_value = (True, 0)
+        self.request_ratelimiter = hs.get_request_ratelimiter()
+        self.request_ratelimiter.can_do_action.return_value = (True, 0)
+
+        self.login_ratelimiter = hs.get_login_ratelimiter()
+        self.login_ratelimiter.can_do_action.return_value = (True, 0)
 
         hs.get_handlers().federation_handler = Mock()
 

tests/rest/client/v1/test_login.py

@@ -36,8 +36,8 @@ def make_homeserver(self, reactor, clock):
         return self.hs
 
     def test_POST_ratelimiting_per_address(self):
-        self.hs.config.rc_login_address.burst_count = 5
-        self.hs.config.rc_login_address.per_second = 0.17
+        self.hs.get_login_ratelimiter().burst_count = 5
+        self.hs.get_login_ratelimiter().rate_hz = 0.17
 
         # Create different users so we're sure not to be bothered by the per-user
         # ratelimiter.
@@ -78,8 +78,8 @@ def test_POST_ratelimiting_per_address(self):
         self.assertEquals(channel.result["code"], b"200", channel.result)
 
     def test_POST_ratelimiting_per_account(self):
-        self.hs.config.rc_login_account.burst_count = 5
-        self.hs.config.rc_login_account.per_second = 0.17
+        self.hs.get_login_ratelimiter().burst_count = 5
+        self.hs.get_login_ratelimiter().rate_hz = 0.17
 
         self.register_user("kermit", "monkey")
 
@@ -117,8 +117,8 @@ def test_POST_ratelimiting_per_account(self):
         self.assertEquals(channel.result["code"], b"200", channel.result)
 
     def test_POST_ratelimiting_per_account_failed_attempts(self):
-        self.hs.config.rc_login_failed_attempts.burst_count = 5
-        self.hs.config.rc_login_failed_attempts.per_second = 0.17
+        self.hs.get_login_failed_attempts_ratelimiter().burst_count = 5
+        self.hs.get_login_failed_attempts_ratelimiter().rate_hz = 0.17
 
         self.register_user("kermit", "monkey")
 

tests/rest/client/v1/test_rooms.py

@@ -49,10 +49,17 @@ def make_homeserver(self, reactor, clock):
             "red",
             http_client=None,
             federation_client=Mock(),
-            ratelimiter=NonCallableMock(spec_set=["can_do_action"]),
+            request_ratelimiter=NonCallableMock(
+                spec_set=["can_do_action", "ratelimit", "rate_hz", "burst_count"]
+            ),
+            login_ratelimiter=NonCallableMock(spec_set=["can_do_action", "ratelimit"]),
         )
-        self.ratelimiter = self.hs.get_ratelimiter()
-        self.ratelimiter.can_do_action.return_value = (True, 0)
+        self.request_ratelimiter = self.hs.get_request_ratelimiter()
+        self.request_ratelimiter.can_do_action.return_value = (True, 0)
+        self.request_ratelimiter.rate_hz = Mock()
+
+        self.login_ratelimiter = self.hs.get_login_ratelimiter()
+        self.login_ratelimiter.can_do_action.return_value = (True, 0)
 
         self.hs.get_federation_handler = Mock(return_value=Mock())
 

tests/rest/client/v1/test_typing.py

@@ -42,13 +42,19 @@ def make_homeserver(self, reactor, clock):
             "red",
             http_client=None,
             federation_client=Mock(),
-            ratelimiter=NonCallableMock(spec_set=["can_do_action"]),
+            request_ratelimiter=NonCallableMock(
+                spec_set=["can_do_action", "ratelimit", "rate_hz", "burst_count"]
+            ),
+            login_ratelimiter=NonCallableMock(spec_set=["can_do_action", "ratelimit"]),
         )
 
         self.event_source = hs.get_event_sources().sources["typing"]
 
-        self.ratelimiter = hs.get_ratelimiter()
-        self.ratelimiter.can_do_action.return_value = (True, 0)
+        self.request_ratelimiter = hs.get_request_ratelimiter()
+        self.request_ratelimiter.can_do_action.return_value = (True, 0)
+
+        self.login_ratelimiter = hs.get_login_ratelimiter()
+        self.login_ratelimiter.can_do_action.return_value = (True, 0)
 
         hs.get_handlers().federation_handler = Mock()
 

tests/rest/client/v2_alpha/test_register.py

@@ -147,8 +147,8 @@ def test_POST_disabled_guest_registration(self):
         self.assertEquals(channel.json_body["error"], "Guest access is disabled")
 
     def test_POST_ratelimiting_guest(self):
-        self.hs.config.rc_registration.burst_count = 5
-        self.hs.config.rc_registration.per_second = 0.17
+        self.hs.get_registration_ratelimiter().burst_count = 5
+        self.hs.get_registration_ratelimiter().rate_hz = 0.17
 
         for i in range(0, 6):
             url = self.url + b"?kind=guest"
@@ -169,8 +169,8 @@ def test_POST_ratelimiting_guest(self):
         self.assertEquals(channel.result["code"], b"200", channel.result)
 
     def test_POST_ratelimiting(self):
-        self.hs.config.rc_registration.burst_count = 5
-        self.hs.config.rc_registration.per_second = 0.17
+        self.hs.get_registration_ratelimiter().burst_count = 5
+        self.hs.get_registration_ratelimiter().rate_hz = 0.17
 
         for i in range(0, 6):
             params = {