Revert "Revert "PERF: improve speed of rate limiter""

This reverts commit 2373d85239.

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

spec/components/auth/default_current_user_provider_spec.rb

@@ -172,38 +172,47 @@ describe Auth::DefaultCurrentUserProvider do
 
   end
 
-  it "can only try 10 bad cookies a minute" do
-    user = Fabricate(:user)
-    token = UserAuthToken.generate!(user_id: user.id)
+  context "rate limiting" do
 
-    provider('/').log_on_user(user, {}, {})
-
-    RateLimiter.stubs(:disabled?).returns(false)
-
-    RateLimiter.new(nil, "cookie_auth_10.0.0.1", 10, 60).clear!
-    RateLimiter.new(nil, "cookie_auth_10.0.0.2", 10, 60).clear!
-
-    ip = "10.0.0.1"
-    env = { "HTTP_COOKIE" => "_t=#{SecureRandom.hex}", "REMOTE_ADDR" => ip }
-
-    10.times do
-      provider('/', env).current_user
+    before do
+      RateLimiter.enable
     end
 
-    expect {
-      provider('/', env).current_user
-    }.to raise_error(Discourse::InvalidAccess)
+    after do
+      RateLimiter.disable
+    end
 
-    expect {
-      env["HTTP_COOKIE"] = "_t=#{token.unhashed_auth_token}"
-      provider("/", env).current_user
-    }.to raise_error(Discourse::InvalidAccess)
+    it "can only try 10 bad cookies a minute" do
+      user = Fabricate(:user)
+      token = UserAuthToken.generate!(user_id: user.id)
 
-    env["REMOTE_ADDR"] = "10.0.0.2"
+      provider('/').log_on_user(user, {}, {})
 
-    expect {
-      provider('/', env).current_user
-    }.not_to raise_error
+      RateLimiter.new(nil, "cookie_auth_10.0.0.1", 10, 60).clear!
+      RateLimiter.new(nil, "cookie_auth_10.0.0.2", 10, 60).clear!
+
+      ip = "10.0.0.1"
+      env = { "HTTP_COOKIE" => "_t=#{SecureRandom.hex}", "REMOTE_ADDR" => ip }
+
+      10.times do
+        provider('/', env).current_user
+      end
+
+      expect {
+        provider('/', env).current_user
+      }.to raise_error(Discourse::InvalidAccess)
+
+      expect {
+        env["HTTP_COOKIE"] = "_t=#{token.unhashed_auth_token}"
+        provider("/", env).current_user
+      }.to raise_error(Discourse::InvalidAccess)
+
+      env["REMOTE_ADDR"] = "10.0.0.2"
+
+      expect {
+        provider('/', env).current_user
+      }.not_to raise_error
+    end
   end
 
   it "correctly removes invalid cookies" do
@@ -295,44 +304,53 @@ describe Auth::DefaultCurrentUserProvider do
 
     end
 
-    it "rate limits api usage" do
+    context "rate limiting" do
 
-      RateLimiter.stubs(:disabled?).returns(false)
-      limiter1 = RateLimiter.new(nil, "user_api_day_#{api_key.key}", 10, 60)
-      limiter2 = RateLimiter.new(nil, "user_api_min_#{api_key.key}", 10, 60)
-      limiter1.clear!
-      limiter2.clear!
-
-      SiteSetting.max_user_api_reqs_per_day = 3
-      SiteSetting.max_user_api_reqs_per_minute = 4
-
-      params = {
-        "REQUEST_METHOD" => "GET",
-        "HTTP_USER_API_KEY" => api_key.key,
-      }
-
-      3.times do
-        provider("/", params).current_user
+      before do
+        RateLimiter.enable
       end
 
-      expect {
-        provider("/", params).current_user
-      }.to raise_error(RateLimiter::LimitExceeded)
-
-      SiteSetting.max_user_api_reqs_per_day = 4
-      SiteSetting.max_user_api_reqs_per_minute = 3
-
-      limiter1.clear!
-      limiter2.clear!
-
-      3.times do
-        provider("/", params).current_user
+      after do
+        RateLimiter.disable
       end
 
-      expect {
-        provider("/", params).current_user
-      }.to raise_error(RateLimiter::LimitExceeded)
+      it "rate limits api usage" do
+        limiter1 = RateLimiter.new(nil, "user_api_day_#{api_key.key}", 10, 60)
+        limiter2 = RateLimiter.new(nil, "user_api_min_#{api_key.key}", 10, 60)
+        limiter1.clear!
+        limiter2.clear!
 
+        SiteSetting.max_user_api_reqs_per_day = 3
+        SiteSetting.max_user_api_reqs_per_minute = 4
+
+        params = {
+          "REQUEST_METHOD" => "GET",
+          "HTTP_USER_API_KEY" => api_key.key,
+        }
+
+        3.times do
+          provider("/", params).current_user
+        end
+
+        expect {
+          provider("/", params).current_user
+        }.to raise_error(RateLimiter::LimitExceeded)
+
+        SiteSetting.max_user_api_reqs_per_day = 4
+        SiteSetting.max_user_api_reqs_per_minute = 3
+
+        limiter1.clear!
+        limiter2.clear!
+
+        3.times do
+          provider("/", params).current_user
+        end
+
+        expect {
+          provider("/", params).current_user
+        }.to raise_error(RateLimiter::LimitExceeded)
+
+      end
     end
   end
 end

spec/components/rate_limiter_spec.rb

@@ -8,11 +8,14 @@ describe RateLimiter do
 
   context 'disabled' do
     before do
-      RateLimiter.stubs(:disabled?).returns(true)
       rate_limiter.performed!
       rate_limiter.performed!
     end
 
+    it "should be disabled" do
+      expect(RateLimiter.disabled?).to eq(true)
+    end
+
     it "returns true for can_perform?" do
       expect(rate_limiter.can_perform?).to eq(true)
     end
@@ -25,10 +28,41 @@ describe RateLimiter do
 
   context 'enabled' do
     before do
-      RateLimiter.stubs(:disabled?).returns(false)
+      RateLimiter.enable
       rate_limiter.clear!
     end
 
+    after do
+      RateLimiter.disable
+    end
+
+    context 'global rate limiter' do
+
+      it 'can operate in global mode' do
+        limiter = RateLimiter.new(nil, "test", 2, 10, global: true)
+        limiter.clear!
+
+        limiter.performed!
+        limiter.performed!
+        expect { limiter.performed! }.to raise_error(RateLimiter::LimitExceeded)
+      end
+
+    end
+
+    context 'handles readonly' do
+      before do
+        $redis.without_namespace.slaveof '10.0.0.1', '99999'
+      end
+
+      after do
+        $redis.without_namespace.slaveof 'no', 'one'
+      end
+
+      it 'does not explode' do
+        expect { rate_limiter.performed! }.not_to raise_error
+      end
+    end
+
     context 'never done' do
       it "should perform right away" do
         expect(rate_limiter.can_perform?).to eq(true)

spec/controllers/session_controller_spec.rb

@@ -756,7 +756,7 @@ describe SessionController do
     context 'rate limited' do
       it 'rate limits login' do
         SiteSetting.max_logins_per_ip_per_hour = 2
-        RateLimiter.stubs(:disabled?).returns(false)
+        RateLimiter.enable
         RateLimiter.clear_all!
 
         2.times do

spec/models/post_action_spec.rb

@@ -20,7 +20,7 @@ describe PostAction do
 
     it "limits redo/undo" do
 
-      RateLimiter.stubs(:disabled?).returns(false)
+      RateLimiter.enable
 
       PostAction.act(eviltrout, post, PostActionType.types[:like])
       PostAction.remove_act(eviltrout, post, PostActionType.types[:like])

spec/models/topic_spec.rb

@@ -579,7 +579,7 @@ describe Topic do
 
     it "rate limits topic invitations" do
       SiteSetting.max_topic_invitations_per_day = 2
-      RateLimiter.stubs(:disabled?).returns(false)
+      RateLimiter.enable
       RateLimiter.clear_all!
 
       start = Time.now.tomorrow.beginning_of_day
@@ -1694,7 +1694,7 @@ describe Topic do
       SiteSetting.max_replies_in_first_day = 1
       SiteSetting.stubs(:client_settings_json).returns(SiteSetting.client_settings_json_uncached)
       RateLimiter.stubs(:rate_limit_create_topic).returns(100)
-      RateLimiter.stubs(:disabled?).returns(false)
+      RateLimiter.enable
       RateLimiter.clear_all!
     end
 

spec/rails_helper.rb

@@ -106,7 +106,8 @@ RSpec.configure do |config|
     #   perf benefit seems low (shaves 20 secs off a 4 minute test suite)
     #
     # $redis = DiscourseMockRedis.new
-    #
+
+    RateLimiter.disable
     PostActionNotifier.disable
     SearchIndexer.disable
     UserActionCreator.disable
@@ -202,6 +203,7 @@ def freeze_time(now = Time.now)
   Time.stubs(:now).returns(time)
   Date.stubs(:today).returns(datetime.to_date)
   TrackTimeStub.stubs(:stubbed).returns(true)
+  Process.stubs(:clock_gettime).with(Process::CLOCK_MONOTONIC).returns(datetime.to_f)
 
   if block_given?
     begin
@@ -217,6 +219,7 @@ def unfreeze_time
   Time.unstub(:now)
   Date.unstub(:today)
   TrackTimeStub.unstub(:stubbed)
+  Process.unstub(:clock_gettime)
 end
 
 def file_from_fixtures(filename, directory = "images")