Rewrite Redis RateLimiterBackend with Lua scripts

Author: mikeroll

dramatiq/rate_limits/backends/redis.py

@@ -17,10 +17,17 @@
 
 from __future__ import annotations
 
+from pathlib import Path
+
 import redis
 
 from ..backend import RateLimiterBackend
 
+_SCRIPTS = {
+    path.stem: path.read_text()
+    for path in (Path(__file__).parent / "redis").glob("*.lua")
+}
+
 
 class RedisBackend(RateLimiterBackend):
     """A rate limiter backend for Redis_.
@@ -41,68 +48,26 @@ def __init__(self, *, client=None, url=None, **parameters):
             parameters["connection_pool"] = redis.ConnectionPool.from_url(url)
 
         self.client = client or redis.Redis(**parameters)
+        self.scripts = {
+            name: self.client.register_script(text) for name, text in _SCRIPTS.items()
+        }
 
     def add(self, key, value, ttl):
         return bool(self.client.set(key, value, px=ttl, nx=True))
 
     def incr(self, key, amount, maximum, ttl):
-        with self.client.pipeline() as pipe:
-            while True:
-                try:
-                    pipe.watch(key)
-                    value = int(pipe.get(key) or b"0")
-                    value += amount
-                    if value > maximum:
-                        return False
-
-                    pipe.multi()
-                    pipe.set(key, value, px=ttl)
-                    pipe.execute()
-                    return True
-                except redis.WatchError:
-                    continue
+        incr_up_to = self.scripts["incr_up_to"]
+        return incr_up_to([key], [amount, maximum, ttl]) == 1
 
     def decr(self, key, amount, minimum, ttl):
-        with self.client.pipeline() as pipe:
-            while True:
-                try:
-                    pipe.watch(key)
-                    value = int(pipe.get(key) or b"0")
-                    value -= amount
-                    if value < minimum:
-                        return False
-
-                    pipe.multi()
-                    pipe.set(key, value, px=ttl)
-                    pipe.execute()
-                    return True
-                except redis.WatchError:
-                    continue
+        decr_down_to = self.scripts["decr_down_to"]
+        return decr_down_to([key], [amount, minimum, ttl]) == 1
 
     def incr_and_sum(self, key, keys, amount, maximum, ttl):
-        with self.client.pipeline() as pipe:
-            while True:
-                try:
-                    # TODO: Drop non-callable keys in Dramatiq v2.
-                    key_list = keys() if callable(keys) else keys
-                    pipe.watch(key, *key_list)
-                    value = int(pipe.get(key) or b"0")
-                    value += amount
-                    if value > maximum:
-                        return False
-
-                    # Fetch keys again to account for net/server latency.
-                    values = pipe.mget(keys() if callable(keys) else keys)
-                    total = amount + sum(int(n) for n in values if n)
-                    if total > maximum:
-                        return False
-
-                    pipe.multi()
-                    pipe.set(key, value, px=ttl)
-                    pipe.execute()
-                    return True
-                except redis.WatchError:
-                    continue
+        # TODO: Drop non-callable keys in Dramatiq v2.
+        keys_list = keys() if callable(keys) else keys
+        incr_up_to_with_sum_check = self.scripts["incr_up_to_with_sum_check"]
+        return incr_up_to_with_sum_check([key, *keys_list], [amount, maximum, ttl]) == 1
 
     def wait(self, key, timeout):
         assert timeout is None or timeout >= 1000, "wait timeouts must be >= 1000"

dramatiq/rate_limits/backends/redis/decr_down_to.lua

@@ -0,0 +1,19 @@
+-- decr_down_to(
+--   keys=[key]
+--   args=[amount, minimum, ttl]
+-- )
+--
+-- Decrement `key` by `amount`, unless the resulting value is less than `minimum`.
+
+local key = KEYS[1]
+local amount = tonumber(ARGV[1])
+local minimum = tonumber(ARGV[2])
+local ttl = tonumber(ARGV[3])
+
+local value = redis.call('GET', key) or 0
+if value - amount < minimum then
+    return false
+end
+
+redis.call('SET', key, value - amount, 'PX', ttl)
+return true

dramatiq/rate_limits/backends/redis/incr_up_to.lua

@@ -0,0 +1,19 @@
+-- incr_up_to(
+--   keys=[key]
+--   args=[amount, maximum, ttl]
+-- )
+--
+-- Increment `key` by `amount`, unless the resulting value is greater than `maximum`.
+
+local key = KEYS[1]
+local amount = tonumber(ARGV[1])
+local maximum = tonumber(ARGV[2])
+local ttl = tonumber(ARGV[3])
+
+local value = redis.call('GET', key) or 0
+if value + amount > maximum then
+    return false
+end
+
+redis.call('SET', key, value + amount, 'PX', ttl)
+return true

dramatiq/rate_limits/backends/redis/incr_up_to_with_sum_check.lua

@@ -0,0 +1,43 @@
+-- incr_up_to_with_sum_check(
+--   keys=[key, *keys]
+--   args=[amount, maximum, ttl]
+-- )
+--
+-- Atomically increment `key` by `amount`, unless:
+--   - the incremented value is greater than `maximum`, or
+--   - the incremented sum of `keys` is greater than `maximum`.
+
+-- split the key list into the first - `key` and the rest - `keys`
+local key = KEYS[1]
+local keys = {}
+for i, k in ipairs(KEYS) do
+    if i > 1 then
+        keys[i - 1] = KEYS[i]
+    end
+end
+
+local amount = tonumber(ARGV[1])
+local maximum = tonumber(ARGV[2])
+local ttl = tonumber(ARGV[3])
+
+-- check if `key` can be incremented, bail if not
+local value = redis.call('GET', key) or 0
+if value + amount > maximum then
+    return false
+end
+
+-- check if sum of `keys` can be incremented, bail if not
+local values = redis.call('MGET', unpack(keys))
+local sum = 0
+for _, v in ipairs(values) do
+    if v then
+        sum = sum + tonumber(v)
+    end
+end
+if sum + amount > maximum then
+    return false
+end
+
+-- increment `key` if we got this far
+redis.call('SET', key, value + amount, 'PX', ttl)
+return true