Detect n+1 issues involving Gitaly

.gitlab-ci.yml

@@ -128,7 +128,7 @@ stages:
     - export CACHE_CLASSES=true
     - cp ${KNAPSACK_SPINACH_SUITE_REPORT_PATH} ${KNAPSACK_REPORT_PATH}
     - scripts/gitaly-test-spawn
-    - knapsack spinach "-r rerun" || retry '[[ -e tmp/spinach-rerun.txt ]] && bundle exec spinach -r rerun $(cat tmp/spinach-rerun.txt)'
+    - knapsack spinach "-r rerun" -b || retry '[[ -e tmp/spinach-rerun.txt ]] && bundle exec spinach -b -r rerun $(cat tmp/spinach-rerun.txt)'
   artifacts:
     expire_in: 31d
     when: always

app/controllers/projects/branches_controller.rb

@@ -15,11 +15,15 @@ class Projects::BranchesController < Projects::ApplicationController
     respond_to do |format|
       format.html do
         @refs_pipelines = @project.pipelines.latest_successful_for_refs(@branches.map(&:name))
-
+        # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37429
+        Gitlab::GitalyClient.allow_n_plus_1_calls do
           @max_commits = @branches.reduce(0) do |memo, branch|
             diverging_commit_counts = repository.diverging_commit_counts(branch)
             [memo, diverging_commit_counts[:behind], diverging_commit_counts[:ahead]].max
           end
+
+          render
+        end
       end
       format.json do
         render json: @branches.map(&:name)

app/controllers/projects/commit_controller.rb

@@ -20,7 +20,12 @@ class Projects::CommitController < Projects::ApplicationController
     apply_diff_view_cookie!
 
     respond_to do |format|
-      format.html
+      format.html do
+        # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37599
+        Gitlab::GitalyClient.allow_n_plus_1_calls do
+          render
+        end
+      end
       format.diff  { render text: @commit.to_diff }
       format.patch { render text: @commit.to_patch }
     end

app/controllers/projects/compare_controller.rb

@@ -17,6 +17,10 @@ class Projects::CompareController < Projects::ApplicationController
 
   def show
     apply_diff_view_cookie!
+    # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37430
+    Gitlab::GitalyClient.allow_n_plus_1_calls do
+      render
+    end
   end
 
   def diff_for_path

app/controllers/projects/merge_requests/diffs_controller.rb

@@ -10,8 +10,11 @@ class Projects::MergeRequests::DiffsController < Projects::MergeRequests::Applic
   def show
     @environment = @merge_request.environments_for(current_user).last
 
+    # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37431
+    Gitlab::GitalyClient.allow_n_plus_1_calls do
       render json: { html: view_to_html_string("projects/merge_requests/diffs/_diffs") }
     end
+  end
 
   def diff_for_path
     render_diff_for_path(@diffs)

app/controllers/projects/merge_requests_controller.rb

@@ -56,6 +56,9 @@ class Projects::MergeRequestsController < Projects::MergeRequests::ApplicationCo
     close_merge_request_without_source_project
     check_if_can_be_merged
 
+    # Return if the response has already been rendered
+    return if response_body
+
     respond_to do |format|
       format.html do
         # Build a note object for comment form
@@ -70,6 +73,11 @@ class Projects::MergeRequestsController < Projects::MergeRequests::ApplicationCo
         labels
 
         set_pipeline_variables
+
+        # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37432
+        Gitlab::GitalyClient.allow_n_plus_1_calls do
+          render
+        end
       end
 
       format.json do

app/controllers/projects/network_controller.rb

@@ -8,6 +8,8 @@ class Projects::NetworkController < Projects::ApplicationController
   before_action :assign_commit
 
   def show
+    # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37602
+    Gitlab::GitalyClient.allow_n_plus_1_calls do
       @url = project_network_path(@project, @ref, @options.merge(format: :json))
       @commit_url = project_commit_path(@project, 'ae45ca32').gsub("ae45ca32", "%s")
 
@@ -22,6 +24,9 @@ class Projects::NetworkController < Projects::ApplicationController
           @graph = Network::Graph.new(project, @ref, @commit, @options[:filter_ref])
         end
       end
+
+      render
+    end
   end
 
   def assign_commit

app/controllers/projects/refs_controller.rb

@@ -51,7 +51,9 @@ class Projects::RefsController < Projects::ApplicationController
     contents.push(*tree.blobs)
     contents.push(*tree.submodules)
 
-    @logs = contents[@offset, @limit].to_a.map do |content|
+    # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37433
+    @logs = Gitlab::GitalyClient.allow_n_plus_1_calls do
+      contents[@offset, @limit].to_a.map do |content|
         file = @path ? File.join(@path, content.name) : content.name
         last_commit = @repo.last_commit_for_path(@commit.id, file)
         {
@@ -59,6 +61,7 @@ class Projects::RefsController < Projects::ApplicationController
           commit: last_commit
         }
       end
+    end
 
     offset = (@offset + @limit)
     if contents.size > offset

app/controllers/root_controller.rb

@@ -13,8 +13,11 @@ class RootController < Dashboard::ProjectsController
   before_action :redirect_logged_user, if: -> { current_user.present? }
 
   def index
+    # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37434
+    Gitlab::GitalyClient.allow_n_plus_1_calls do
       super
     end
+  end
 
   private
 

app/models/merge_request.rb

@@ -415,9 +415,12 @@ class MergeRequest < ActiveRecord::Base
   end
 
   def create_merge_request_diff
+    # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37435
+    Gitlab::GitalyClient.allow_n_plus_1_calls do
       merge_request_diffs.create
       reload_merge_request_diff
     end
+  end
 
   def reload_merge_request_diff
     merge_request_diff(true)

app/models/network/graph.rb

@@ -61,9 +61,12 @@ module Network
         @reserved[i] = []
       end
 
+      # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37436
+      Gitlab::GitalyClient.allow_n_plus_1_calls do
         commits_sort_by_ref.each do |commit|
           place_chain(commit)
         end
+      end
 
       # find parent spaces for not overlap lines
       @commits.each do |c|

app/models/repository.rb

@@ -983,6 +983,7 @@ class Repository
   def empty_repo?
     !exists? || !has_visible_content?
   end
+  cache_method :empty_repo?, memoize_only: true
 
   def search_files_by_content(query, ref)
     return [] if empty_repo? || query.blank?

app/services/delete_merged_branches_service.rb

@@ -6,6 +6,8 @@ class DeleteMergedBranchesService < BaseService
   def execute
     raise Gitlab::Access::AccessDeniedError unless can?(current_user, :push_code, project)
 
+    # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37438
+    Gitlab::GitalyClient.allow_n_plus_1_calls do
       branches = project.repository.branch_names
       branches = branches.select { |branch| project.repository.merged_to_root_ref?(branch) }
       # Prevent deletion of branches relevant to open merge requests
@@ -17,6 +19,7 @@ class DeleteMergedBranchesService < BaseService
         DeleteBranchService.new(project, current_user).execute(branch)
       end
     end
+  end
 
   private
 

app/services/merge_requests/create_service.rb

@@ -13,8 +13,11 @@ module MergeRequests
       merge_request.source_branch = params[:source_branch]
       merge_request.merge_params['force_remove_source_branch'] = params.delete(:force_remove_source_branch)
 
+      # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37439
+      Gitlab::GitalyClient.allow_n_plus_1_calls do
         create(merge_request)
       end
+    end
 
     def before_create(merge_request)
       # current_user (defined in BaseService) is not available within run_after_commit block

app/services/notes/create_service.rb

@@ -4,7 +4,13 @@ module Notes
       merge_request_diff_head_sha = params.delete(:merge_request_diff_head_sha)
 
       note = Notes::BuildService.new(project, current_user, params).execute
-      return note unless note.valid?
+
+      # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37440
+      note_valid = Gitlab::GitalyClient.allow_n_plus_1_calls do
+        note.valid?
+      end
+
+      return note unless note_valid
 
       # We execute commands (extracted from `params[:note]`) on the noteable
       # **before** we save the note because if the note consists of commands

config/initializers/lograge.rb

@@ -12,13 +12,18 @@ unless Sidekiq.server?
     config.lograge.logger = ActiveSupport::Logger.new(filename)
     # Add request parameters to log output
     config.lograge.custom_options = lambda do |event|
-      {
+      payload = {
         time: event.time.utc.iso8601(3),
         params: event.payload[:params].except(*%w(controller action format)),
         remote_ip: event.payload[:remote_ip],
         user_id: event.payload[:user_id],
         username: event.payload[:username]
       }
+
+      gitaly_calls = Gitlab::GitalyClient.get_request_count
+      payload[:gitaly_calls] = gitaly_calls if gitaly_calls > 0
+
+      payload
     end
   end
 end

lib/api/branches.rb

@@ -21,8 +21,11 @@ module API
       get ':id/repository/branches' do
         branches = ::Kaminari.paginate_array(user_project.repository.branches.sort_by(&:name))
 
+        # 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
         desc 'Get a single branch' do

lib/api/entities.rb

@@ -244,8 +244,11 @@ module API
       end
 
       expose :merged do |repo_branch, options|
+        # 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|
         ::ProtectedBranch.protected?(options[:project], repo_branch.name)

lib/gitlab/diff/file_collection/base.rb

@@ -22,8 +22,11 @@ module Gitlab
         end
 
         def diff_files
+          # 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_files.find { |diff_file| diff_file.old_path == old_path }

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,7 +962,9 @@ 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|
+        # 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)
@@ -970,6 +972,7 @@ module Gitlab
               # Omit invalid branch
             end
           end.compact
+        end
 
         sort_branches(branches, sort_by)
       end

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,14 +126,102 @@ 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
+        # 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
 
     def self.expected_server_version
       path = Rails.root.join(SERVER_VERSION_FILE)
@@ -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

spec/lib/gitlab/gitaly_client_spec.rb

@@ -38,6 +38,130 @@ describe Gitlab::GitalyClient, skip_gitaly_mock: true do
     end
   end
 
+  describe 'allow_n_plus_1_calls' do
+    context 'when RequestStore is enabled', :request_store do
+      it 'returns the result of the allow_n_plus_1_calls block' do
+        expect(described_class.allow_n_plus_1_calls { "result" }).to eq("result")
+      end
+    end
+
+    context 'when RequestStore is not active' do
+      it 'returns the result of the allow_n_plus_1_calls block' do
+        expect(described_class.allow_n_plus_1_calls { "something" }).to eq("something")
+      end
+    end
+  end
+
+  describe 'enforce_gitaly_request_limits?' do
+    def call_gitaly(count = 1)
+      (1..count).each do
+        described_class.enforce_gitaly_request_limits(:test)
+      end
+    end
+
+    context 'when RequestStore is enabled', :request_store do
+      it 'allows up the maximum number of allowed calls' do
+        expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS) }.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)
+        end
+
+        it 'fails on the next call' do
+          expect { call_gitaly(1) }.to raise_error(Gitlab::GitalyClient::TooManyInvocationsError)
+        end
+      end
+
+      it 'allows the maximum number of calls to be exceeded within an allow_n_plus_1_calls block' do
+        expect do
+          described_class.allow_n_plus_1_calls do
+            call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS + 1)
+          end
+        end.not_to raise_error
+      end
+
+      context 'when the maximum number of calls has been reached within an allow_n_plus_1_calls block' do
+        before do
+          described_class.allow_n_plus_1_calls do
+            call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS)
+          end
+        end
+
+        it 'allows up to the maximum number of calls outside of an allow_n_plus_1_calls block' do
+          expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS) }.not_to raise_error
+        end
+
+        it 'does not allow the maximum number of calls to be exceeded outside of an allow_n_plus_1_calls block' do
+          expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS + 1) }.to raise_error(Gitlab::GitalyClient::TooManyInvocationsError)
+        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
+      end
+
+      it 'does not fail when the maximum number of calls is exceeded within an allow_n_plus_1_calls block' do
+        expect do
+          described_class.allow_n_plus_1_calls do
+            call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS + 1)
+          end
+        end.not_to raise_error
+      end
+    end
+  end
+
+  describe 'get_request_count' do
+    context 'when RequestStore is enabled', :request_store do
+      context 'when enforce_gitaly_request_limits is called outside of allow_n_plus_1_calls blocks' do
+        before do
+          described_class.enforce_gitaly_request_limits(:call)
+        end
+
+        it 'counts gitaly calls' do
+          expect(described_class.get_request_count).to eq(1)
+        end
+      end
+
+      context 'when enforce_gitaly_request_limits is called inside and outside of allow_n_plus_1_calls blocks' do
+        before do
+          described_class.enforce_gitaly_request_limits(:call)
+          described_class.allow_n_plus_1_calls do
+            described_class.enforce_gitaly_request_limits(:call)
+          end
+        end
+
+        it 'counts gitaly calls' do
+          expect(described_class.get_request_count).to eq(2)
+        end
+      end
+
+      context 'when reset_counts is called' do
+        before do
+          described_class.enforce_gitaly_request_limits(:call)
+          described_class.reset_counts
+        end
+
+        it 'resets counts' do
+          expect(described_class.get_request_count).to eq(0)
+        end
+      end
+    end
+
+    context 'when RequestStore is not active' do
+      before do
+        described_class.enforce_gitaly_request_limits(:call)
+      end
+
+      it 'returns zero' do
+        expect(described_class.get_request_count).to eq(0)
+      end
+    end
+  end
+
   describe 'feature_enabled?' do
     let(:feature_name) { 'my_feature' }
     let(:real_feature_name) { "gitaly_#{feature_name}" }

spec/serializers/pipeline_serializer_spec.rb

@@ -103,10 +103,16 @@ describe PipelineSerializer do
       let(:project) { create(:project) }
 
       before do
+        # Since RequestStore.active? is true we have to allow the
+        # gitaly calls in this block
+        # Issue: https://gitlab.com/gitlab-org/gitlab-ce/issues/37772
+        Gitlab::GitalyClient.allow_n_plus_1_calls do
           Ci::Pipeline::AVAILABLE_STATUSES.each do |status|
             create_pipeline(status)
           end
         end
+        Gitlab::GitalyClient.reset_counts
+      end
 
       shared_examples 'no N+1 queries' do
         it 'verifies number of queries', :request_store do