Detect n+1 issues involving Gitaly

5 files changed, 171 insertions, 14 deletions

diff --git a/lib/api/branches.rb b/lib/api/branches.rb
index 642c1140fcc..643c8e6fb8e 100644
--- a/lib/api/branches.rb
+++ b/lib/api/branches.rb
@@ -21,7 +21,10 @@ module API
       get ':id/repository/branches' do
         branches = ::Kaminari.paginate_array(user_project.repository.branches.sort_by(&:name))
 
-        present paginate(branches), with: Entities::RepoBranch, project: user_project
+        # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37442
+        Gitlab::GitalyClient.allow_n_plus_1_calls do
+          present paginate(branches), with: Entities::RepoBranch, project: user_project
+        end
       end
 
       resource ':id/repository/branches/:branch', requirements: BRANCH_ENDPOINT_REQUIREMENTS do
diff --git a/lib/api/entities.rb b/lib/api/entities.rb
index 30b115b1b56..71253f72533 100644
--- a/lib/api/entities.rb
+++ b/lib/api/entities.rb
@@ -244,7 +244,10 @@ module API
       end
 
       expose :merged do |repo_branch, options|
-        options[:project].repository.merged_to_root_ref?(repo_branch.name)
+        # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37442
+        Gitlab::GitalyClient.allow_n_plus_1_calls do
+          options[:project].repository.merged_to_root_ref?(repo_branch.name)
+        end
       end
 
       expose :protected do |repo_branch, options|
diff --git a/lib/gitlab/diff/file_collection/base.rb b/lib/gitlab/diff/file_collection/base.rb
index a6007ebf531..88ae65cb468 100644
--- a/lib/gitlab/diff/file_collection/base.rb
+++ b/lib/gitlab/diff/file_collection/base.rb
@@ -22,7 +22,10 @@ module Gitlab
         end
 
         def diff_files
-          @diff_files ||= @diffs.decorate! { |diff| decorate_diff!(diff) }
+          # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37445
+          Gitlab::GitalyClient.allow_n_plus_1_calls do
+            @diff_files ||= @diffs.decorate! { |diff| decorate_diff!(diff) }
+          end
         end
 
         def diff_file_with_old_path(old_path)
diff --git a/lib/gitlab/git/repository.rb b/lib/gitlab/git/repository.rb
index c499ff101b5..18210bcab4e 100644
--- a/lib/gitlab/git/repository.rb
+++ b/lib/gitlab/git/repository.rb
@@ -950,8 +950,8 @@ module Gitlab
         @gitaly_repository_client ||= Gitlab::GitalyClient::RepositoryService.new(self)
       end
 
-      def gitaly_migrate(method, &block)
-        Gitlab::GitalyClient.migrate(method, &block)
+      def gitaly_migrate(method, status: Gitlab::GitalyClient::MigrationStatus::OPT_IN, &block)
+        Gitlab::GitalyClient.migrate(method, status: status, &block)
       rescue GRPC::NotFound => e
         raise NoRepository.new(e)
       rescue GRPC::BadStatus => e
@@ -962,14 +962,17 @@ module Gitlab
 
       # Gitaly note: JV: Trying to get rid of the 'filter' option so we can implement this with 'git'.
       def branches_filter(filter: nil, sort_by: nil)
-        branches = rugged.branches.each(filter).map do |rugged_ref|
-          begin
-            target_commit = Gitlab::Git::Commit.find(self, rugged_ref.target)
-            Gitlab::Git::Branch.new(self, rugged_ref.name, rugged_ref.target, target_commit)
-          rescue Rugged::ReferenceError
-            # Omit invalid branch
-          end
-        end.compact
+        # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37464
+        branches = Gitlab::GitalyClient.allow_n_plus_1_calls do
+          rugged.branches.each(filter).map do |rugged_ref|
+            begin
+              target_commit = Gitlab::Git::Commit.find(self, rugged_ref.target)
+              Gitlab::Git::Branch.new(self, rugged_ref.name, rugged_ref.target, target_commit)
+            rescue Rugged::ReferenceError
+              # Omit invalid branch
+            end
+          end.compact
+        end
 
         sort_branches(branches, sort_by)
       end
diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index a3dc2cd0b60..cbd9ff406de 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -10,7 +10,24 @@ module Gitlab
       OPT_OUT = 3
     end
 
+    class TooManyInvocationsError < StandardError
+      attr_reader :call_site, :invocation_count, :max_call_stack
+
+      def initialize(call_site, invocation_count, max_call_stack, most_invoked_stack)
+        @call_site = call_site
+        @invocation_count = invocation_count
+        @max_call_stack = max_call_stack
+        stacks = most_invoked_stack.join('\n') if most_invoked_stack
+
+        msg = "GitalyClient##{call_site} called #{invocation_count} times from single request. Potential n+1?"
+        msg << "\nThe following call site called into Gitaly #{max_call_stack} times:\n#{stacks}\n" if stacks
+
+        super(msg)
+      end
+    end
+
     SERVER_VERSION_FILE = 'GITALY_SERVER_VERSION'.freeze
+    MAXIMUM_GITALY_CALLS = 30
 
     MUTEX = Mutex.new
     private_constant :MUTEX
@@ -53,6 +70,8 @@ module Gitlab
     # All Gitaly RPC call sites should use GitalyClient.call. This method
     # makes sure that per-request authentication headers are set.
     def self.call(storage, service, rpc, request)
+      enforce_gitaly_request_limits(:call)
+
       metadata = request_metadata(storage)
       metadata = yield(metadata) if block_given?
       stub(service, storage).__send__(rpc, request, metadata) # rubocop:disable GitlabSecurity/PublicSend
@@ -107,12 +126,100 @@ module Gitlab
     private_class_method :opt_into_all_features?
 
     def self.migrate(feature, status: MigrationStatus::OPT_IN)
+      # Enforce limits at both the `migrate` and `call` sites to ensure that
+      # problems are not hidden by a feature being disabled
+      enforce_gitaly_request_limits(:migrate)
+
       is_enabled  = feature_enabled?(feature, status: status)
       metric_name = feature.to_s
       metric_name += "_gitaly" if is_enabled
 
       Gitlab::Metrics.measure(metric_name) do
-        yield is_enabled
+        # Some migrate calls wrap other migrate calls
+        allow_n_plus_1_calls do
+          yield is_enabled
+        end
+      end
+    end
+
+    # Ensures that Gitaly is not being abuse through n+1 misuse etc
+    def self.enforce_gitaly_request_limits(call_site)
+      # Only count limits in request-response environments (not sidekiq for example)
+      return unless RequestStore.active?
+
+      # 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?
+
+      # Check if this call is nested within a allow_n_plus_1_calls
+      # block and skip check if it is
+      return if get_call_count(:gitaly_call_count_exception_block_depth) > 0
+
+      # This is the count of calls outside of a `allow_n_plus_1_calls` block
+      # It is used for enforcement but not statistics
+      permitted_call_count = increment_call_count("gitaly_#{call_site}_permitted")
+
+      count_stack
+
+      return if permitted_call_count <= MAXIMUM_GITALY_CALLS
+
+      raise TooManyInvocationsError.new(call_site, actual_call_count, max_call_count, max_stacks)
+    end
+
+    def self.allow_n_plus_1_calls
+      return yield unless RequestStore.active?
+
+      begin
+        increment_call_count(:gitaly_call_count_exception_block_depth)
+        yield
+      ensure
+        decrement_call_count(:gitaly_call_count_exception_block_depth)
+      end
+    end
+
+    def self.get_call_count(key)
+      RequestStore.store[key] || 0
+    end
+    private_class_method :get_call_count
+
+    def self.increment_call_count(key)
+      RequestStore.store[key] ||= 0
+      RequestStore.store[key] += 1
+    end
+    private_class_method :increment_call_count
+
+    def self.decrement_call_count(key)
+      RequestStore.store[key] -= 1
+    end
+    private_class_method :decrement_call_count
+
+    # Returns an estimate of the number of Gitaly calls made for this
+    # request
+    def self.get_request_count
+      return 0 unless RequestStore.active?
+
+      gitaly_migrate_count = get_call_count("gitaly_migrate_actual")
+      gitaly_call_count = get_call_count("gitaly_call_actual")
+
+      # Using the maximum of migrate and call_count will provide an
+      # indicator of how many Gitaly calls will be made, even
+      # before a feature is enabled. This provides us with a single
+      # metric, but not an exact number, but this tradeoff is acceptable
+      if gitaly_migrate_count > gitaly_call_count
+        gitaly_migrate_count
+      else
+        gitaly_call_count
+      end
+    end
+
+    def self.reset_counts
+      return unless RequestStore.active?
+
+      %w[migrate call].each do |call_site|
+        RequestStore.store["gitaly_#{call_site}_actual"] = 0
+        RequestStore.store["gitaly_#{call_site}_permitted"] = 0
       end
     end
 
@@ -124,5 +231,43 @@ module Gitlab
     def self.encode(s)
       s.dup.force_encoding(Encoding::ASCII_8BIT)
     end
+
+    # Count a stack. Used for n+1 detection
+    def self.count_stack
+      return unless RequestStore.active?
+
+      stack_string = caller.drop(1).join("\n")
+
+      RequestStore.store[:stack_counter] ||= Hash.new
+
+      count = RequestStore.store[:stack_counter][stack_string] || 0
+      RequestStore.store[:stack_counter][stack_string] = count + 1
+    end
+    private_class_method :count_stack
+
+    # Returns a count for the stack which called Gitaly the most times. Used for n+1 detection
+    def self.max_call_count
+      return 0 unless RequestStore.active?
+
+      stack_counter = RequestStore.store[:stack_counter]
+      return 0 unless stack_counter
+
+      stack_counter.values.max
+    end
+    private_class_method :max_call_count
+
+    # Returns the stacks that calls Gitaly the most times. Used for n+1 detection
+    def self.max_stacks
+      return nil unless RequestStore.active?
+
+      stack_counter = RequestStore.store[:stack_counter]
+      return nil unless stack_counter
+
+      max = max_call_count
+      return nil if max.zero?
+
+      stack_counter.select { |_, v| v == max }.keys
+    end
+    private_class_method :max_stacks
   end
 end