Repository: zulip/zulip

Issue:
Optimize rate_limiter performance for get_events queries
See https://chat.zulip.org/#narrow/stream/3-backend/topic/profiling.20get_events/near/816860 for profiling details, but basically, currently a get_events request spends 1.4ms/request talking to redis for our rate limiter, which is somewhere between 15% and 50% of the total request runtime (my measurement technique is susceptible to issues like the first request on a code path being extra expensive).  Since get_events is our most scalability-critical endpoint, this is a big deal.

We should do some rethinking of the redis internals for our rate limiter.  I have a few ideas:
* Writing an alternative rate-limiter implementation for `get_events `specifically that's entirely in-process and would be basically instant.  Since the Tornado system has a relatively strong constraint that a given user always connect to the same server, this might be fairly cheap to implement and would bring that 1.4ms to probably 50us or less. (And gate it on `RUNNING_INSIDE_TORNADO`).  
* Look at rewriting our redis transactions to be more efficient for the highest-traffic cases (E.g. user is not close to limit, or user is way over limit).  E.g. maybe `incr_rateimit` should automatically return the `api_calls_left` result rather than requiring 2 transactions.
* Looking at https://github.com/popravich/python-redis-benchmark, there may be some alternative async IO redis clients we could consider migrating to, and possibly some that are just faster.  Given how little code we have interacting with redis directly, this might be an easy port to do; I'm not sure whether or not it would help.  (And unlike the in-process hack approach, this would have side benefits to non-Tornado endpoints).


Discussion:
Hello @zulip/server-production members, this issue was labeled with the "area: production" label, so you may want to check it out!

<!-- areaLabelAddition -->
@zulipbot claim

Resolving patch:
diff --git a/zerver/lib/rate_limiter.py b/zerver/lib/rate_limiter.py
--- a/zerver/lib/rate_limiter.py
+++ b/zerver/lib/rate_limiter.py
@@ -32,12 +32,14 @@ class RateLimitedObject(ABC):
     def __init__(self, backend: Optional['Type[RateLimiterBackend]']=None) -> None:
         if backend is not None:
             self.backend = backend  # type: Type[RateLimiterBackend]
+        elif settings.RUNNING_INSIDE_TORNADO:
+            self.backend = TornadoInMemoryRateLimiterBackend
         else:
             self.backend = RedisRateLimiterBackend
 
     def rate_limit(self) -> Tuple[bool, float]:
         # Returns (ratelimited, secs_to_freedom)
-        return self.backend.rate_limit_entity(self.key(), self.rules(),
+        return self.backend.rate_limit_entity(self.key(), self.get_rules(),
                                               self.max_api_calls(),
                                               self.max_api_window())
 
@@ -74,11 +76,11 @@ def clear_history(self) -> None:
 
     def max_api_calls(self) -> int:
         "Returns the API rate limit for the highest limit"
-        return self.rules()[-1][1]
+        return self.get_rules()[-1][1]
 
     def max_api_window(self) -> int:
         "Returns the API time window for the highest limit"
-        return self.rules()[-1][0]
+        return self.get_rules()[-1][0]
 
     def api_calls_left(self) -> Tuple[int, float]:
         """Returns how many API calls in this range this client has, as well as when
@@ -87,6 +89,16 @@ def api_calls_left(self) -> Tuple[int, float]:
         max_calls = self.max_api_calls()
         return self.backend.get_api_calls_left(self.key(), max_window, max_calls)
 
+    def get_rules(self) -> List[Tuple[int, int]]:
+        """
+        This is a simple wrapper meant to protect against having to deal with
+        an empty list of rules, as it would require fiddling with that special case
+        all around this system. "9999 max request per seconds" should be a good proxy
+        for "no rules".
+        """
+        rules_list = self.rules()
+        return rules_list or [(1, 9999), ]
+
     @abstractmethod
     def key(self) -> str:
         pass
@@ -139,6 +151,7 @@ class RateLimiterBackend(ABC):
     @abstractmethod
     def block_access(cls, entity_key: str, seconds: int) -> None:
         "Manually blocks an entity for the desired number of seconds"
+        pass
 
     @classmethod
     @abstractmethod
@@ -163,6 +176,120 @@ def rate_limit_entity(cls, entity_key: str, rules: List[Tuple[int, int]],
         # Returns (ratelimited, secs_to_freedom)
         pass
 
+class TornadoInMemoryRateLimiterBackend(RateLimiterBackend):
+    # reset_times[rule][key] is the time at which the event
+    # request from the rate-limited key will be accepted.
+    reset_times = {}  # type: Dict[Tuple[int, int], Dict[str, float]]
+
+    # last_gc_time is the last time when the garbage was
+    # collected from reset_times for rule (time_window, max_count).
+    last_gc_time = {}  # type: Dict[Tuple[int, int], float]
+
+    # timestamps_blocked_until[key] contains the timestamp
+    # up to which the key has been blocked manually.
+    timestamps_blocked_until = {}  # type: Dict[str, float]
+
+    @classmethod
+    def _garbage_collect_for_rule(cls, now: float, time_window: int, max_count: int) -> None:
+        keys_to_delete = []
+        reset_times_for_rule = cls.reset_times.get((time_window, max_count), None)
+        if reset_times_for_rule is None:
+            return
+
+        keys_to_delete = [entity_key for entity_key in reset_times_for_rule
+                          if reset_times_for_rule[entity_key] < now]
+
+        for entity_key in keys_to_delete:
+            del reset_times_for_rule[entity_key]
+
+        if not reset_times_for_rule:
+            del cls.reset_times[(time_window, max_count)]
+
+    @classmethod
+    def need_to_limit(cls, entity_key: str, time_window: int,
+                      max_count: int) -> Tuple[bool, float]:
+        '''
+        Returns a tuple of `(rate_limited, time_till_free)`.
+        For simplicity, we have loosened the semantics here from
+        - each key may make atmost `count * (t / window)` request within any t
+          time interval.
+        to
+        - each key may make atmost `count * [(t / window) + 1]` request within
+          any t time interval.
+        Thus, we only need to store reset_times for each key which will be less
+        memory-intensive. This also has the advantage that you can only ever
+        lock yourself out completely for `window / count` seconds instead of
+        `window` seconds.
+        '''
+        now = time.time()
+
+        # Remove all timestamps from `reset_times` that are too old.
+        if cls.last_gc_time.get((time_window, max_count), 0) <= now - time_window / max_count:
+            cls.last_gc_time[(time_window, max_count)] = now
+            cls._garbage_collect_for_rule(now, time_window, max_count)
+
+        reset_times_for_rule = cls.reset_times.setdefault((time_window, max_count), {})
+        new_reset = max(reset_times_for_rule.get(entity_key, now), now) \
+            + time_window / max_count
+
+        if new_reset > now + time_window:
+            # Compute for how long the bucket will remain filled.
+            time_till_free = new_reset - time_window - now
+            return True, time_till_free
+
+        reset_times_for_rule[entity_key] = new_reset
+        return False, 0.0
+
+    @classmethod
+    def get_api_calls_left(cls, entity_key: str, range_seconds: int,
+                           max_calls: int) -> Tuple[int, float]:
+        now = time.time()
+        if (range_seconds, max_calls) in cls.reset_times and \
+                entity_key in cls.reset_times[(range_seconds, max_calls)]:
+            reset_time = cls.reset_times[(range_seconds, max_calls)][entity_key]
+        else:
+            return max_calls, 0
+
+        calls_remaining = (now + range_seconds - reset_time) * max_calls // range_seconds
+        return int(calls_remaining), reset_time - now
+
+    @classmethod
+    def block_access(cls, entity_key: str, seconds: int) -> None:
+        now = time.time()
+        cls.timestamps_blocked_until[entity_key] = now + seconds
+
+    @classmethod
+    def unblock_access(cls, entity_key: str) -> None:
+        del cls.timestamps_blocked_until[entity_key]
+
+    @classmethod
+    def clear_history(cls, entity_key: str) -> None:
+        for rule, reset_times_for_rule in cls.reset_times.items():
+            reset_times_for_rule.pop(entity_key, None)
+        cls.timestamps_blocked_until.pop(entity_key, None)
+
+    @classmethod
+    def rate_limit_entity(cls, entity_key: str, rules: List[Tuple[int, int]],
+                          max_api_calls: int, max_api_window: int) -> Tuple[bool, float]:
+        now = time.time()
+        if entity_key in cls.timestamps_blocked_until:
+            # Check whether the key is manually blocked.
+            if now < cls.timestamps_blocked_until[entity_key]:
+                blocking_ttl = cls.timestamps_blocked_until[entity_key] - now
+                return True, blocking_ttl
+            else:
+                del cls.timestamps_blocked_until[entity_key]
+
+        assert rules
+        for time_window, max_count in rules:
+            ratelimited, time_till_free = cls.need_to_limit(entity_key, time_window, max_count)
+
+            if ratelimited:
+                statsd.incr("ratelimiter.limited.%s" % (entity_key,))
+                break
+
+        return ratelimited, time_till_free
+
 class RedisRateLimiterBackend(RateLimiterBackend):
     @classmethod
     def get_keys(cls, entity_key: str) -> List[str]:
@@ -220,11 +347,9 @@ def get_api_calls_left(cls, entity_key: str, range_seconds: int,
     @classmethod
     def is_ratelimited(cls, entity_key: str, rules: List[Tuple[int, int]]) -> Tuple[bool, float]:
         "Returns a tuple of (rate_limited, time_till_free)"
+        assert rules
         list_key, set_key, blocking_key = cls.get_keys(entity_key)
 
-        if len(rules) == 0:
-            return False, 0.0
-
         # Go through the rules from shortest to longest,
         # seeing if this user has violated any of them. First
         # get the timestamps for each nth items
@@ -244,7 +369,7 @@ def is_ratelimited(cls, entity_key: str, rules: List[Tuple[int, int]]) -> Tuple[
 
         if key_blocked is not None:
             # We are manually blocked. Report for how much longer we will be
-            if blocking_ttl_b is None:
+            if blocking_ttl_b is None:  # nocoverage # defensive code, this should never happen
                 blocking_ttl = 0.5
             else:
                 blocking_ttl = int(blocking_ttl_b)
@@ -265,16 +390,11 @@ def is_ratelimited(cls, entity_key: str, rules: List[Tuple[int, int]]) -> Tuple[
         return False, 0.0
 
     @classmethod
-    def incr_ratelimit(cls, entity_key: str, rules: List[Tuple[int, int]],
-                       max_api_calls: int, max_api_window: int) -> None:
+    def incr_ratelimit(cls, entity_key: str, max_api_calls: int, max_api_window: int) -> None:
         """Increases the rate-limit for the specified entity"""
         list_key, set_key, _ = cls.get_keys(entity_key)
         now = time.time()
 
-        # If we have no rules, we don't store anything
-        if len(rules) == 0:
-            return
-
         # Start redis transaction
         with client.pipeline() as pipe:
             count = 0
@@ -316,7 +436,7 @@ def incr_ratelimit(cls, entity_key: str, rules: List[Tuple[int, int]],
 
                     # If no exception was raised in the execution, there were no transaction conflicts
                     break
-                except redis.WatchError:
+                except redis.WatchError:  # nocoverage # Ideally we'd have a test for this.
                     if count > 10:
                         raise RateLimiterLockingException()
                     count += 1
@@ -333,7 +453,7 @@ def rate_limit_entity(cls, entity_key: str, rules: List[Tuple[int, int]],
 
         else:
             try:
-                cls.incr_ratelimit(entity_key, rules, max_api_calls, max_api_window)
+                cls.incr_ratelimit(entity_key, max_api_calls, max_api_window)
             except RateLimiterLockingException:
                 logger.warning("Deadlock trying to incr_ratelimit for %s" % (entity_key,))
                 # rate-limit users who are hitting the API so hard we can't update our stats.


Test patch:
diff --git a/zerver/tests/test_rate_limiter.py b/zerver/tests/test_rate_limiter.py
--- a/zerver/tests/test_rate_limiter.py
+++ b/zerver/tests/test_rate_limiter.py
@@ -5,6 +5,7 @@
     RateLimitedUser,
     RateLimiterBackend,
     RedisRateLimiterBackend,
+    TornadoInMemoryRateLimiterBackend,
 )
 
 from zerver.lib.test_classes import ZulipTestCase
@@ -68,7 +69,7 @@ def verify_api_calls_left(self, obj: RateLimitedTestObject) -> None:
         self.assertEqual(expected_time_till_reset, time_till_reset)
 
     def expected_api_calls_left(self, obj: RateLimitedTestObject, now: float) -> Tuple[int, float]:
-        longest_rule = obj.rules()[-1]
+        longest_rule = obj.get_rules()[-1]
         max_window, max_calls = longest_rule
         history = self.requests_record.get(obj.key())
         if history is None:
@@ -82,7 +83,7 @@ def api_calls_left_from_history(self, history: List[float], max_window: int,
         """
         This depends on the algorithm used in the backend, and should be defined by the test class.
         """
-        raise NotImplementedError  # nocoverage
+        raise NotImplementedError()
 
     def test_hit_ratelimits(self) -> None:
         obj = self.create_object('test', [(2, 3), ])
@@ -165,13 +166,45 @@ def test_block_access(self) -> None:
         obj.block_access(1)
         self.make_request(obj, expect_ratelimited=True, verify_api_calls_left=False)
 
-class RateLimitedUserTest(ZulipTestCase):
+class TornadoInMemoryRateLimiterBackendTest(RateLimiterBackendBase):
+    __unittest_skip__ = False
+    backend = TornadoInMemoryRateLimiterBackend
+
+    def api_calls_left_from_history(self, history: List[float], max_window: int,
+                                    max_calls: int, now: float) -> Tuple[int, float]:
+        reset_time = 0.0
+        for timestamp in history:
+            reset_time = max(reset_time, timestamp) + (max_window / max_calls)
+
+        calls_left = (now + max_window - reset_time) * max_calls // max_window
+        calls_left = int(calls_left)
+
+        return calls_left, reset_time - now
+
+    def test_used_in_tornado(self) -> None:
+        user_profile = self.example_user("hamlet")
+        with self.settings(RUNNING_INSIDE_TORNADO=True):
+            obj = RateLimitedUser(user_profile)
+        self.assertEqual(obj.backend, TornadoInMemoryRateLimiterBackend)
+
+    def test_block_access(self) -> None:
+        obj = self.create_object('test', [(2, 5), ])
+        start_time = time.time()
+
+        obj.block_access(1)
+        with mock.patch('time.time', return_value=(start_time)):
+            self.make_request(obj, expect_ratelimited=True, verify_api_calls_left=False)
+
+        with mock.patch('time.time', return_value=(start_time + 1.01)):
+            self.make_request(obj, expect_ratelimited=False, verify_api_calls_left=False)
+
+class RateLimitedObjectsTest(ZulipTestCase):
     def test_user_rate_limits(self) -> None:
         user_profile = self.example_user("hamlet")
         user_profile.rate_limits = "1:3,2:4"
         obj = RateLimitedUser(user_profile)
 
-        self.assertEqual(obj.rules(), [(1, 3), (2, 4)])
+        self.assertEqual(obj.get_rules(), [(1, 3), (2, 4)])
 
     def test_add_remove_rule(self) -> None:
         user_profile = self.example_user("hamlet")
@@ -180,9 +213,13 @@ def test_add_remove_rule(self) -> None:
         add_ratelimit_rule(10, 100, domain='some_new_domain')
         obj = RateLimitedUser(user_profile)
 
-        self.assertEqual(obj.rules(), [(1, 2), ])
+        self.assertEqual(obj.get_rules(), [(1, 2), ])
         obj.domain = 'some_new_domain'
-        self.assertEqual(obj.rules(), [(4, 5), (10, 100)])
+        self.assertEqual(obj.get_rules(), [(4, 5), (10, 100)])
 
         remove_ratelimit_rule(10, 100, domain='some_new_domain')
-        self.assertEqual(obj.rules(), [(4, 5), ])
+        self.assertEqual(obj.get_rules(), [(4, 5), ])
+
+    def test_empty_rules_edge_case(self) -> None:
+        obj = RateLimitedTestObject("test", rules=[], backend=RedisRateLimiterBackend)
+        self.assertEqual(obj.get_rules(), [(1, 9999), ])