test_ratelimitutils.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  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 twisted.internet import defer
  16. from twisted.internet.defer import Deferred
  17. from synapse.config.homeserver import HomeServerConfig
  18. from synapse.config.ratelimiting import FederationRatelimitSettings
  19. from synapse.util.ratelimitutils import FederationRateLimiter
  20. from tests.server import ThreadedMemoryReactorClock, get_clock
  21. from tests.unittest import TestCase
  22. from tests.utils import default_config
  23. class FederationRateLimiterTestCase(TestCase):
  24. def test_ratelimit(self) -> None:
  25. """A simple test with the default values"""
  26. reactor, clock = get_clock()
  27. rc_config = build_rc_config()
  28. ratelimiter = FederationRateLimiter(clock, rc_config)
  29. with ratelimiter.ratelimit("testhost") as d1:
  30. # shouldn't block
  31. self.successResultOf(d1)
  32. def test_concurrent_limit(self) -> None:
  33. """Test what happens when we hit the concurrent limit"""
  34. reactor, clock = get_clock()
  35. rc_config = build_rc_config({"rc_federation": {"concurrent": 2}})
  36. ratelimiter = FederationRateLimiter(clock, rc_config)
  37. with ratelimiter.ratelimit("testhost") as d1:
  38. # shouldn't block
  39. self.successResultOf(d1)
  40. cm2 = ratelimiter.ratelimit("testhost")
  41. d2 = cm2.__enter__()
  42. # also shouldn't block
  43. self.successResultOf(d2)
  44. cm3 = ratelimiter.ratelimit("testhost")
  45. d3 = cm3.__enter__()
  46. # this one should block, though ...
  47. self.assertNoResult(d3)
  48. # ... until we complete an earlier request
  49. cm2.__exit__(None, None, None)
  50. reactor.advance(0.0)
  51. self.successResultOf(d3)
  52. def test_sleep_limit(self) -> None:
  53. """Test what happens when we hit the sleep limit"""
  54. reactor, clock = get_clock()
  55. rc_config = build_rc_config(
  56. {"rc_federation": {"sleep_limit": 2, "sleep_delay": 500}}
  57. )
  58. ratelimiter = FederationRateLimiter(clock, rc_config)
  59. with ratelimiter.ratelimit("testhost") as d1:
  60. # shouldn't block
  61. self.successResultOf(d1)
  62. with ratelimiter.ratelimit("testhost") as d2:
  63. # nor this
  64. self.successResultOf(d2)
  65. with ratelimiter.ratelimit("testhost") as d3:
  66. # this one should block, though ...
  67. self.assertNoResult(d3)
  68. sleep_time = _await_resolution(reactor, d3)
  69. self.assertAlmostEqual(sleep_time, 500, places=3)
  70. def test_lots_of_queued_things(self) -> None:
  71. """Tests lots of synchronous things queued up behind a slow thing.
  72. The stack should *not* explode when the slow thing completes.
  73. """
  74. reactor, clock = get_clock()
  75. rc_config = build_rc_config(
  76. {
  77. "rc_federation": {
  78. "sleep_limit": 1000000000, # never sleep
  79. "reject_limit": 1000000000, # never reject requests
  80. "concurrent": 1,
  81. }
  82. }
  83. )
  84. ratelimiter = FederationRateLimiter(clock, rc_config)
  85. with ratelimiter.ratelimit("testhost") as d:
  86. # shouldn't block
  87. self.successResultOf(d)
  88. async def task() -> None:
  89. with ratelimiter.ratelimit("testhost") as d:
  90. await d
  91. for _ in range(1, 100):
  92. defer.ensureDeferred(task())
  93. last_task = defer.ensureDeferred(task())
  94. # Upon exiting the context manager, all the synchronous things will resume.
  95. # If a stack overflow occurs, the final task will not complete.
  96. # Wait for all the things to complete.
  97. reactor.advance(0.0)
  98. self.successResultOf(last_task)
  99. def _await_resolution(reactor: ThreadedMemoryReactorClock, d: Deferred) -> float:
  100. """advance the clock until the deferred completes.
  101. Returns the number of milliseconds it took to complete.
  102. """
  103. start_time = reactor.seconds()
  104. while not d.called:
  105. reactor.advance(0.01)
  106. return (reactor.seconds() - start_time) * 1000
  107. def build_rc_config(settings: Optional[dict] = None) -> FederationRatelimitSettings:
  108. config_dict = default_config("test")
  109. config_dict.update(settings or {})
  110. config = HomeServerConfig()
  111. config.parse_config_dict(config_dict, "", "")
  112. return config.ratelimiting.rc_federation