test_ratelimitutils.py 3.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. # Copyright 2019 The Matrix.org Foundation C.I.C.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from typing import Optional
  15. from synapse.config.homeserver import HomeServerConfig
  16. from synapse.util.ratelimitutils import FederationRateLimiter
  17. from tests.server import get_clock
  18. from tests.unittest import TestCase
  19. from tests.utils import default_config
  20. class FederationRateLimiterTestCase(TestCase):
  21. def test_ratelimit(self):
  22. """A simple test with the default values"""
  23. reactor, clock = get_clock()
  24. rc_config = build_rc_config()
  25. ratelimiter = FederationRateLimiter(clock, rc_config)
  26. with ratelimiter.ratelimit("testhost") as d1:
  27. # shouldn't block
  28. self.successResultOf(d1)
  29. def test_concurrent_limit(self):
  30. """Test what happens when we hit the concurrent limit"""
  31. reactor, clock = get_clock()
  32. rc_config = build_rc_config({"rc_federation": {"concurrent": 2}})
  33. ratelimiter = FederationRateLimiter(clock, rc_config)
  34. with ratelimiter.ratelimit("testhost") as d1:
  35. # shouldn't block
  36. self.successResultOf(d1)
  37. cm2 = ratelimiter.ratelimit("testhost")
  38. d2 = cm2.__enter__()
  39. # also shouldn't block
  40. self.successResultOf(d2)
  41. cm3 = ratelimiter.ratelimit("testhost")
  42. d3 = cm3.__enter__()
  43. # this one should block, though ...
  44. self.assertNoResult(d3)
  45. # ... until we complete an earlier request
  46. cm2.__exit__(None, None, None)
  47. self.successResultOf(d3)
  48. def test_sleep_limit(self):
  49. """Test what happens when we hit the sleep limit"""
  50. reactor, clock = get_clock()
  51. rc_config = build_rc_config(
  52. {"rc_federation": {"sleep_limit": 2, "sleep_delay": 500}}
  53. )
  54. ratelimiter = FederationRateLimiter(clock, rc_config)
  55. with ratelimiter.ratelimit("testhost") as d1:
  56. # shouldn't block
  57. self.successResultOf(d1)
  58. with ratelimiter.ratelimit("testhost") as d2:
  59. # nor this
  60. self.successResultOf(d2)
  61. with ratelimiter.ratelimit("testhost") as d3:
  62. # this one should block, though ...
  63. self.assertNoResult(d3)
  64. sleep_time = _await_resolution(reactor, d3)
  65. self.assertAlmostEqual(sleep_time, 500, places=3)
  66. def _await_resolution(reactor, d):
  67. """advance the clock until the deferred completes.
  68. Returns the number of milliseconds it took to complete.
  69. """
  70. start_time = reactor.seconds()
  71. while not d.called:
  72. reactor.advance(0.01)
  73. return (reactor.seconds() - start_time) * 1000
  74. def build_rc_config(settings: Optional[dict] = None):
  75. config_dict = default_config("test")
  76. config_dict.update(settings or {})
  77. config = HomeServerConfig()
  78. config.parse_config_dict(config_dict, "", "")
  79. return config.ratelimiting.rc_federation