Merge pull request #11 from otovo/deterministic-lua-scripts

Remove call to TIME in Lua scripts

limiters/token_bucket.lua

@@ -20,13 +20,14 @@ redis.replicate_commands()
 local capacity = tonumber(ARGV[1])
 local refill_amount = tonumber(ARGV[2])
 local time_between_slots = tonumber(ARGV[3]) * 1000 -- ms
+local seconds = tonumber(ARGV[4])
+local microseconds = tonumber(ARGV[5])
 
 -- Keys
 local data_key = KEYS[1]
 
 -- Get current time (ms timestamp)
-local redis_time = redis.call('TIME') -- Array of [seconds, microseconds]
-local now = tonumber(redis_time[1]) * 1000 + (tonumber(redis_time[2]) / 1000)
+local now = tonumber(seconds) * 1000 + (tonumber(microseconds) / 1000)
 
 -- Instantiate default bucket values
 -- These are only used if a bucket doesn't already exist

limiters/token_bucket.py

@@ -13,6 +13,19 @@
 logger = logging.getLogger(__name__)
 
 
+def create_redis_time_tuple() -> tuple[int, int]:
+    """
+    Create a tuple of two integers representing the current time in seconds and microseconds.
+
+    This mimmicks the TIME command in Redis, which returns the current time in seconds and microseconds.
+    See: https://redis.io/commands/time/
+    """
+    now = time.time()
+    seconds_part = int(now)
+    microseconds_part = int((now - seconds_part) * 1_000_000)
+    return seconds_part, microseconds_part
+
+
 class TokenBucketBase(BaseModel):
     name: str
     capacity: int = Field(gt=0)
@@ -60,10 +73,12 @@ def __enter__(self) -> float:
         Call the token bucket Lua script, receive a datetime for
         when to wake up, then sleep up until that point in time.
         """
+
         # Retrieve timestamp for when to wake up from Redis
+        seconds, microseconds = create_redis_time_tuple()
         timestamp: int = self.script(
             keys=[self.key],
-            args=[self.capacity, self.refill_amount, self.refill_frequency],
+            args=[self.capacity, self.refill_amount, self.refill_frequency, seconds, microseconds],
         )
 
         # Estimate sleep time
@@ -91,10 +106,12 @@ async def __aenter__(self) -> None:
         Call the token bucket Lua script, receive a datetime for
         when to wake up, then sleep up until that point in time.
         """
+
         # Retrieve timestamp for when to wake up from Redis
+        seconds, microseconds = create_redis_time_tuple()
         timestamp = await self.script(
             keys=[self.key],
-            args=[self.capacity, self.refill_amount, self.refill_frequency],
+            args=[self.capacity, self.refill_amount, self.refill_frequency, seconds, microseconds],
         )
 
         # Estimate sleep time