Merge branch 'ml-add-gitaly-enforce-request-limit-feature-flag' into 'master'

Add feature flag to enforce gitaly request limits

See merge request gitlab-org/gitlab-ce!25931

2 files changed, 51 insertions, 3 deletions

diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index 48c113a8b14..0a371889af2 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -257,8 +257,7 @@ module Gitlab
       # This is this actual number of times this call was made. Used for information purposes only
       actual_call_count = increment_call_count("gitaly_#{call_site}_actual")
 
-      # Do no enforce limits in production
-      return if Rails.env.production? || ENV["GITALY_DISABLE_REQUEST_LIMITS"]
+      return unless enforce_gitaly_request_limits?
 
       # Check if this call is nested within a allow_n_plus_1_calls
       # block and skip check if it is
@@ -275,6 +274,19 @@ module Gitlab
       raise TooManyInvocationsError.new(call_site, actual_call_count, max_call_count, max_stacks)
     end
 
+    def self.enforce_gitaly_request_limits?
+      # We typically don't want to enforce request limits in production
+      # However, we have some production-like test environments, i.e., ones
+      # where `Rails.env.production?` returns `true`. We do want to be able to
+      # check if the limit is being exceeded while testing in those environments
+      # In that case we can use a feature flag to indicate that we do want to
+      # enforce request limits.
+      return true if feature_enabled?('enforce_requests_limits')
+
+      !(Rails.env.production? || ENV["GITALY_DISABLE_REQUEST_LIMITS"])
+    end
+    private_class_method :enforce_gitaly_request_limits?
+
     def self.allow_n_plus_1_calls
       return yield unless Gitlab::SafeRequestStore.active?
 
diff --git a/spec/lib/gitlab/gitaly_client_spec.rb b/spec/lib/gitlab/gitaly_client_spec.rb
index cf12baf1a93..f1acb1d9bc4 100644
--- a/spec/lib/gitlab/gitaly_client_spec.rb
+++ b/spec/lib/gitlab/gitaly_client_spec.rb
@@ -149,11 +149,21 @@ describe Gitlab::GitalyClient do
       end
     end
 
-    context 'when RequestStore is enabled', :request_store do
+    context 'when RequestStore is enabled and the maximum number of calls is not enforced by a feature flag', :request_store do
+      before do
+        stub_feature_flags(gitaly_enforce_requests_limits: false)
+      end
+
       it 'allows up the maximum number of allowed calls' do
         expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS) }.not_to raise_error
       end
 
+      it 'allows the maximum number of calls to be exceeded if GITALY_DISABLE_REQUEST_LIMITS is set' do
+        stub_env('GITALY_DISABLE_REQUEST_LIMITS', 'true')
+
+        expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS + 1) }.not_to raise_error
+      end
+
       context 'when the maximum number of calls has been reached' do
         before do
           call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS)
@@ -189,6 +199,32 @@ describe Gitlab::GitalyClient do
       end
     end
 
+    context 'in production and when RequestStore is enabled', :request_store do
+      before do
+        allow(Rails.env).to receive(:production?).and_return(true)
+      end
+
+      context 'when the maximum number of calls is enforced by a feature flag' do
+        before do
+          stub_feature_flags(gitaly_enforce_requests_limits: true)
+        end
+
+        it 'does not allow the maximum number of calls to be exceeded' do
+          expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS + 1) }.to raise_error(Gitlab::GitalyClient::TooManyInvocationsError)
+        end
+      end
+
+      context 'when the maximum number of calls is not enforced by a feature flag' do
+        before do
+          stub_feature_flags(gitaly_enforce_requests_limits: false)
+        end
+
+        it 'allows the maximum number of calls to be exceeded' do
+          expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS + 1) }.not_to raise_error
+        end
+      end
+    end
+
     context 'when RequestStore is not active' do
       it 'does not raise errors when the maximum number of allowed calls is exceeded' do
         expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS + 2) }.not_to raise_error