PERF: improve speed of rate limiter

Also

- adds a global rate limiter option
- cleans up usage in tests
- fixes freeze_time so it handles clock_gettime

lib/rate_limiter.rb

@@ -20,9 +20,10 @@ class RateLimiter
 
   # We don't observe rate limits in test mode
   def self.disabled?
-    @disabled || Rails.env.test?
+    @disabled
   end
 
+  # Only used in test, only clears current namespace, does not clear globals
   def self.clear_all!
     $redis.delete_prefixed(RateLimiter.key_prefix)
   end
@@ -31,65 +32,102 @@ class RateLimiter
     "#{RateLimiter.key_prefix}:#{@user && @user.id}:#{type}"
   end
 
-  def initialize(user, type, max, secs)
+  def initialize(user, type, max, secs, global: false)
     @user = user
     @type = type
     @key = build_key(type)
     @max = max
     @secs = secs
+    @global = false
   end
 
   def clear!
-    $redis.del(@key)
+    redis.del(prefixed_key)
   end
 
   def can_perform?
     rate_unlimited? || is_under_limit?
   end
 
+  # reloader friendly
+  unless defined? PERFORM_LUA
+    PERFORM_LUA = <<~LUA
+      local now = tonumber(ARGV[1])
+      local secs = tonumber(ARGV[2])
+      local max = tonumber(ARGV[3])
+
+      local key = KEYS[1]
+
+
+      if ((tonumber(redis.call("LLEN", key)) < max) or
+          (now - tonumber(redis.call("LRANGE", key, -1, -1)[1])) > secs) then
+        redis.call("LPUSH", key, now)
+        redis.call("LTRIM", key, 0, max - 1)
+        redis.call("EXPIRE", key, secs * 2)
+
+        return 1
+      else
+        return 0
+      end
+    LUA
+
+    PERFORM_LUA_SHA = Digest::SHA1.hexdigest(PERFORM_LUA)
+  end
+
   def performed!
     return if rate_unlimited?
-
-    if is_under_limit?
-      # simple ring buffer.
-      $redis.lpush(@key, Time.now.to_i)
-      $redis.ltrim(@key, 0, @max - 1)
-
-      # let's ensure we expire this key at some point, otherwise we have leaks
-      $redis.expire(@key, @secs * 2)
-    else
+    now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    if eval_lua(PERFORM_LUA, PERFORM_LUA_SHA, [prefixed_key], [now, @secs, @max]) == 0
       raise RateLimiter::LimitExceeded.new(seconds_to_wait, @type)
     end
+  rescue Redis::CommandError => e
+    if e.message =~ /READONLY/
+      # TODO,switch to in-memory rate limiter
+    else
+      raise
+    end
   end
 
   def rollback!
     return if RateLimiter.disabled?
-    $redis.lpop(@key)
+    redis.lpop(prefixed_key)
   end
 
   def remaining
     return @max if @user && @user.staff?
 
-    arr = $redis.lrange(@key, 0, @max) || []
-    t0 = Time.now.to_i
+    arr = redis.lrange(prefixed_key, 0, @max) || []
+    t0 = Process.clock_gettime(Process::CLOCK_MONOTONIC)
     arr.reject! { |a| (t0 - a.to_i) > @secs }
     @max - arr.size
   end
 
   private
 
+  def prefixed_key
+    if @global
+      "GLOBAL::#{key}"
+    else
+      $redis.namespace_key(key)
+    end
+  end
+
+  def redis
+    $redis.without_namespace
+  end
+
   def seconds_to_wait
     @secs - age_of_oldest
   end
 
   def age_of_oldest
     # age of oldest event in buffer, in seconds
-    Time.now.to_i - $redis.lrange(@key, -1, -1).first.to_i
+    Process.clock_gettime(Process::CLOCK_MONOTONIC) - redis.lrange(prefixed_key, -1, -1).first.to_i
   end
 
   def is_under_limit?
     # number of events in buffer less than max allowed? OR
-    ($redis.llen(@key) < @max) ||
+    (redis.llen(prefixed_key) < @max) ||
     # age bigger than silding window size?
     (age_of_oldest > @secs)
   end
@@ -97,4 +135,14 @@ class RateLimiter
   def rate_unlimited?
     !!(RateLimiter.disabled? || (@user && @user.staff?))
   end
+
+  def eval_lua(lua, sha, keys, args)
+    redis.evalsha(sha, keys, args)
+  rescue Redis::CommandError => e
+    if e.to_s =~ /^NOSCRIPT/
+      redis.eval(lua, keys, args)
+    else
+      raise
+    end
+  end
 end