Merge branch '285352-change-rate-limiter-redis-key' into 'master'

Improve tracking of requests in rate limiter

See merge request gitlab-org/gitlab!73343

config/feature_flags/development/rate_limiter_safe_increment.yml

+---
+name: rate_limiter_safe_increment
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/73343
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/285352
+milestone: '14.5'
+type: development
+group: group::project management
+default_enabled: false

lib/gitlab/application_rate_limiter.rb

@@ -66,7 +66,6 @@ module Gitlab
       # @param key [Symbol] Key attribute registered in `.rate_limits`
       # @option scope [Array<ActiveRecord>] Array of ActiveRecord models to scope throttling to a specific request (e.g. per user per project)
       # @option threshold [Integer] Optional threshold value to override default one registered in `.rate_limits`
-      # @option interval [Integer] Optional interval value to override default one registered in `.rate_limits`
       # @option users_allowlist [Array<String>] Optional list of usernames to exclude from the limit. This param will only be functional if Scope includes a current user.
       #
       # @return [Boolean] Whether or not a request should be throttled
@@ -77,7 +76,7 @@ module Gitlab
 
         threshold_value = options[:threshold] || threshold(key)
         threshold_value > 0 &&
-          increment(key, options[:scope], options[:interval]) > threshold_value
+          increment(key, options[:scope]) > threshold_value
       end
 
       # Increments the given cache key and increments the value by 1 with the
@@ -85,12 +84,13 @@ module Gitlab
       #
       # @param key [Symbol] Key attribute registered in `.rate_limits`
       # @option scope [Array<ActiveRecord>] Array of ActiveRecord models to scope throttling to a specific request (e.g. per user per project)
-      # @option interval [Integer] Optional interval value to override default one registered in `.rate_limits`
       #
       # @return [Integer] incremented value
-      def increment(key, scope, interval = nil)
+      def increment(key, scope)
+        return safe_increment(key, scope) if Feature.enabled?(:rate_limiter_safe_increment, default_enabled: :yaml)
+
         value = 0
-        interval_value = interval || interval(key)
+        interval_value = interval(key)
 
         ::Gitlab::Redis::RateLimiting.with do |redis|
           cache_key = action_key(key, scope)
@@ -101,6 +101,32 @@ module Gitlab
         value
       end
 
+      # Increments a cache key that is based on the current time and interval.
+      # So that when time passes to the next interval, the key changes and the count starts again from 0.
+      #
+      # Based on https://github.com/rack/rack-attack/blob/886ba3a18d13c6484cd511a4dc9b76c0d14e5e96/lib/rack/attack/cache.rb#L63-L68
+      #
+      # @param key [Symbol] Key attribute registered in `.rate_limits`
+      # @option scope [Array<ActiveRecord>] Array of ActiveRecord models to scope throttling to a specific request (e.g. per user per project)
+      #
+      # @return [Integer] incremented value
+      def safe_increment(key, scope)
+        interval_value = interval(key)
+
+        period_key, time_elapsed_in_period = Time.now.to_i.divmod(interval_value)
+
+        cache_key = "#{action_key(key, scope)}:#{period_key}"
+        # We add a 1 second buffer to avoid timing issues when we're at the end of a period
+        expiry = interval_value - time_elapsed_in_period + 1
+
+        ::Gitlab::Redis::RateLimiting.with do |redis|
+          redis.pipelined do
+            redis.incr(cache_key)
+            redis.expire(cache_key, expiry)
+          end.first
+        end
+      end
+
       # Logs request using provided logger
       #
       # @param request [Http::Request] - Web request to be logged

spec/lib/gitlab/application_rate_limiter_spec.rb

@@ -3,53 +3,29 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::ApplicationRateLimiter do
-  let(:redis) { double('redis') }
-  let(:user) { create(:user) }
-  let(:project) { create(:project) }
+  let_it_be(:user) { create(:user) }
+  let_it_be(:project) { create(:project) }
+
+  subject { described_class }
+
+  describe '.throttled?' do
     let(:rate_limits) do
       {
         test_action: {
           threshold: 1,
           interval: 2.minutes
+        },
+        another_action: {
+          threshold: 2,
+          interval: 3.minutes
         }
       }
     end
 
-  let(:key) { rate_limits.keys[0] }
-
-  subject { described_class }
-
     before do
-    allow(Gitlab::Redis::RateLimiting).to receive(:with).and_yield(redis)
       allow(described_class).to receive(:rate_limits).and_return(rate_limits)
     end
 
-  shared_examples 'action rate limiter' do
-    it 'increases the throttle count and sets the expiration time' do
-      expect(redis).to receive(:incr).with(cache_key).and_return(1)
-      expect(redis).to receive(:expire).with(cache_key, 120)
-
-      expect(subject.throttled?(key, scope: scope)).to be_falsy
-    end
-
-    it 'returns true if the key is throttled' do
-      expect(redis).to receive(:incr).with(cache_key).and_return(2)
-      expect(redis).not_to receive(:expire)
-
-      expect(subject.throttled?(key, scope: scope)).to be_truthy
-    end
-
-    context 'when throttling is disabled' do
-      it 'returns false and does not set expiration time' do
-        expect(redis).not_to receive(:incr)
-        expect(redis).not_to receive(:expire)
-
-        expect(subject.throttled?(key, scope: scope, threshold: 0)).to be_falsy
-      end
-    end
-  end
-
-  describe '.throttled?' do
     context 'when the key is invalid' do
       context 'is provided as a Symbol' do
         context 'but is not defined in the rate_limits Hash' do
@@ -80,6 +56,41 @@ RSpec.describe Gitlab::ApplicationRateLimiter do
       end
     end
 
+    context 'when rate_limiter_safe_increment is disabled' do
+      let(:redis) { double('redis') }
+      let(:key) { rate_limits.keys[0] }
+
+      before do
+        allow(Gitlab::Redis::RateLimiting).to receive(:with).and_yield(redis)
+
+        stub_feature_flags(rate_limiter_safe_increment: false)
+      end
+
+      shared_examples 'action rate limiter' do
+        it 'increases the throttle count and sets the expiration time' do
+          expect(redis).to receive(:incr).with(cache_key).and_return(1)
+          expect(redis).to receive(:expire).with(cache_key, 120)
+
+          expect(subject.throttled?(key, scope: scope)).to be_falsy
+        end
+
+        it 'returns true if the key is throttled' do
+          expect(redis).to receive(:incr).with(cache_key).and_return(2)
+          expect(redis).not_to receive(:expire)
+
+          expect(subject.throttled?(key, scope: scope)).to be_truthy
+        end
+
+        context 'when throttling is disabled' do
+          it 'returns false and does not set expiration time' do
+            expect(redis).not_to receive(:incr)
+            expect(redis).not_to receive(:expire)
+
+            expect(subject.throttled?(key, scope: scope, threshold: 0)).to be_falsy
+          end
+        end
+      end
+
       context 'when the key is an array of only ActiveRecord models' do
         let(:scope) { [user, project] }
 
@@ -90,7 +101,7 @@ RSpec.describe Gitlab::ApplicationRateLimiter do
         it_behaves_like 'action rate limiter'
       end
 
-    context 'when they key a combination of ActiveRecord models and strings' do
+      context 'when the key is a combination of ActiveRecord models and strings' do
         let(:project) { create(:project, :public, :repository) }
         let(:commit) { project.repository.commit }
         let(:path) { 'app/controllers/groups_controller.rb' }
@@ -104,6 +115,60 @@ RSpec.describe Gitlab::ApplicationRateLimiter do
       end
     end
 
+    context 'when rate_limiter_safe_increment is enabled', :clean_gitlab_redis_rate_limiting do
+      before do
+        stub_feature_flags(rate_limiter_safe_increment: true)
+      end
+
+      shared_examples 'throttles based on key and scope' do
+        let(:start_time) { Time.current.beginning_of_hour }
+
+        it 'returns true when threshold is exceeded' do
+          travel_to(start_time) do
+            expect(subject.throttled?(:test_action, scope: scope)).to eq(false)
+          end
+
+          travel_to(start_time + 1.minute) do
+            expect(subject.throttled?(:test_action, scope: scope)).to eq(true)
+
+            # Assert that it does not affect other actions or scope
+            expect(subject.throttled?(:another_action, scope: scope)).to eq(false)
+            expect(subject.throttled?(:test_action, scope: [user])).to eq(false)
+          end
+        end
+
+        it 'returns false when interval has elapsed' do
+          travel_to(start_time) do
+            expect(subject.throttled?(:test_action, scope: scope)).to eq(false)
+
+            # another_action has a threshold of 3 so we simulate 2 requests
+            expect(subject.throttled?(:another_action, scope: scope)).to eq(false)
+            expect(subject.throttled?(:another_action, scope: scope)).to eq(false)
+          end
+
+          travel_to(start_time + 2.minutes) do
+            expect(subject.throttled?(:test_action, scope: scope)).to eq(false)
+
+            # Assert that another_action has its own interval that hasn't elapsed
+            expect(subject.throttled?(:another_action, scope: scope)).to eq(true)
+          end
+        end
+      end
+
+      context 'when using ActiveRecord models as scope' do
+        let(:scope) { [user, project] }
+
+        it_behaves_like 'throttles based on key and scope'
+      end
+
+      context 'when using ActiveRecord models and strings as scope' do
+        let(:scope) { [project, 'app/controllers/groups_controller.rb'] }
+
+        it_behaves_like 'throttles based on key and scope'
+      end
+    end
+  end
+
   describe '.log_request' do
     let(:file_path) { 'master/README.md' }
     let(:type) { :raw_blob_request_limit }
@@ -134,7 +199,7 @@ RSpec.describe Gitlab::ApplicationRateLimiter do
     end
 
     context 'with a current_user' do
-      let(:current_user) { create(:user) }
+      let(:current_user) { user }
 
       let(:attributes) do
         base_attributes.merge({