Use gitaly pagination to improve performance

Pagination by Kaminari is slow because it requests all branches and
paginate after receiving results from Gitaly. It's a first iteration to
implement native pagination with Gitaly.

app/finders/branches_finder.rb

@@ -5,12 +5,16 @@ class BranchesFinder < GitRefsFinder
     super(repository, params)
   end
 
-  def execute
+  def execute(gitaly_pagination: false)
+    if gitaly_pagination && names.blank? && search.blank?
+      repository.branches_sorted_by(sort, pagination_params)
+    else
       branches = repository.branches_sorted_by(sort)
       branches = by_search(branches)
       branches = by_names(branches)
       branches
     end
+  end
 
   private
 
@@ -18,6 +22,18 @@ class BranchesFinder < GitRefsFinder
     @params[:names].presence
   end
 
+  def per_page
+    @params[:per_page].presence
+  end
+
+  def page_token
+    "#{Gitlab::Git::BRANCH_REF_PREFIX}#{@params[:page_token]}" if @params[:page_token]
+  end
+
+  def pagination_params
+    { limit: per_page, page_token: page_token }
+  end
+
   def by_names(branches)
     return branches unless names
 

app/models/repository.rb

@@ -713,8 +713,8 @@ class Repository
     "#{name}-#{highest_branch_id + 1}"
   end
 
-  def branches_sorted_by(value)
-    raw_repository.local_branches(sort_by: value)
+  def branches_sorted_by(sort_by, pagination_params = nil)
+    raw_repository.local_branches(sort_by: sort_by, pagination_params: pagination_params)
   end
 
   def tags_sorted_by(value)

changelogs/unreleased/208738-improve-performance-of-branches-list-api-when-under-load.yml

+---
+title: Use native Gitaly pagination for Branch list API
+merge_request: 35819
+author:
+type: changed

lib/api/branches.rb

@@ -32,14 +32,21 @@ module API
       params do
         use :pagination
         use :filter_params
+
+        optional :page_token, type: String, desc: 'Name of branch to start the paginaition from'
       end
       get ':id/repository/branches' do
         user_project.preload_protected_branches
 
         repository = user_project.repository
 
+        if Feature.enabled?(:branch_list_keyset_pagination, user_project)
+          branches = BranchesFinder.new(repository, declared_params(include_missing: false)).execute(gitaly_pagination: true)
+        else
           branches = BranchesFinder.new(repository, declared_params(include_missing: false)).execute
           branches = paginate(::Kaminari.paginate_array(branches))
+        end
+
         merged_branch_names = repository.merged_branch_names(branches.map(&:name))
 
         present(

lib/gitlab/git/repository.rb

@@ -127,9 +127,9 @@ module Gitlab
         end
       end
 
-      def local_branches(sort_by: nil)
+      def local_branches(sort_by: nil, pagination_params: nil)
         wrapped_gitaly_errors do
-          gitaly_ref_client.local_branches(sort_by: sort_by)
+          gitaly_ref_client.local_branches(sort_by: sort_by, pagination_params: pagination_params)
         end
       end
 

lib/gitlab/gitaly_client/ref_service.rb

@@ -110,8 +110,8 @@ module Gitlab
         branch_names.count
       end
 
-      def local_branches(sort_by: nil)
-        request = Gitaly::FindLocalBranchesRequest.new(repository: @gitaly_repo)
+      def local_branches(sort_by: nil, pagination_params: nil)
+        request = Gitaly::FindLocalBranchesRequest.new(repository: @gitaly_repo, pagination_params: pagination_params)
         request.sort_by = sort_by_param(sort_by) if sort_by
         response = GitalyClient.call(@storage, :ref_service, :find_local_branches, request, timeout: GitalyClient.fast_timeout)
         consume_find_local_branches_response(response)

spec/finders/branches_finder_spec.rb

@@ -7,20 +7,28 @@ RSpec.describe BranchesFinder do
   let(:project) { create(:project, :repository) }
   let(:repository) { project.repository }
 
+  let(:branch_finder) { described_class.new(repository, params) }
+  let(:params) { {} }
+
   describe '#execute' do
+    subject { branch_finder.execute }
+
     context 'sort only' do
-      it 'sorts by name' do
-        branches_finder = described_class.new(repository, {})
+      context 'by name' do
+        let(:params) { {} }
 
-        result = branches_finder.execute
+        it 'sorts' do
+          result = subject
 
           expect(result.first.name).to eq("'test'")
         end
+      end
 
-      it 'sorts by recently_updated' do
-        branches_finder = described_class.new(repository, { sort: 'updated_desc' })
+      context 'by recently_updated' do
+        let(:params) { { sort: 'updated_desc' } }
 
-        result = branches_finder.execute
+        it 'sorts' do
+          result = subject
 
           recently_updated_branch = repository.branches.max do |a, b|
             repository.commit(a.dereferenced_target).committed_date <=> repository.commit(b.dereferenced_target).committed_date
@@ -28,122 +36,227 @@ RSpec.describe BranchesFinder do
 
           expect(result.first.name).to eq(recently_updated_branch.name)
         end
+      end
 
-      it 'sorts by last_updated' do
-        branches_finder = described_class.new(repository, { sort: 'updated_asc' })
+      context 'by last_updated' do
+        let(:params) { { sort: 'updated_asc' } }
 
-        result = branches_finder.execute
+        it 'sorts' do
+          result = subject
 
           expect(result.first.name).to eq('feature')
         end
       end
+    end
 
     context 'filter only' do
-      it 'filters branches by name' do
-        branches_finder = described_class.new(repository, { search: 'fix' })
+      context 'by name' do
+        let(:params) { { search: 'fix' } }
 
-        result = branches_finder.execute
+        it 'filters branches' do
+          result = subject
 
           expect(result.first.name).to eq('fix')
           expect(result.count).to eq(1)
         end
+      end
 
-      it 'filters branches by name ignoring letter case' do
-        branches_finder = described_class.new(repository, { search: 'FiX' })
+      context 'by name ignoring letter case' do
+        let(:params) { { search: 'FiX' } }
 
-        result = branches_finder.execute
+        it 'filters branches' do
+          result = subject
 
           expect(result.first.name).to eq('fix')
           expect(result.count).to eq(1)
         end
+      end
 
-      it 'does not find any branch with that name' do
-        branches_finder = described_class.new(repository, { search: 'random' })
+      context 'with an unknown name' do
+        let(:params) { { search: 'random' } }
 
-        result = branches_finder.execute
+        it 'does not find any branch' do
+          result = subject
 
           expect(result.count).to eq(0)
         end
+      end
 
-      it 'filters branches by provided names' do
-        branches_finder = described_class.new(repository, { names: %w[fix csv lfs does-not-exist] })
+      context 'by provided names' do
+        let(:params) { { names: %w[fix csv lfs does-not-exist] } }
 
-        result = branches_finder.execute
+        it 'filters branches' do
+          result = subject
 
           expect(result.count).to eq(3)
           expect(result.map(&:name)).to eq(%w{csv fix lfs})
         end
+      end
 
-      it 'filters branches by name that begins with' do
-        params = { search: '^feature_' }
-        branches_finder = described_class.new(repository, params)
+      context 'by name that begins with' do
+        let(:params) { { search: '^feature_' } }
 
-        result = branches_finder.execute
+        it 'filters branches' do
+          result = subject
 
           expect(result.first.name).to eq('feature_conflict')
           expect(result.count).to eq(1)
         end
+      end
 
-      it 'filters branches by name that ends with' do
-        params = { search: 'feature$' }
-        branches_finder = described_class.new(repository, params)
+      context 'by name that ends with' do
+        let(:params) { { search: 'feature$' } }
 
-        result = branches_finder.execute
+        it 'filters branches' do
+          result = subject
 
           expect(result.first.name).to eq('feature')
           expect(result.count).to eq(1)
         end
+      end
 
-      it 'filters branches by nonexistent name that begins with' do
-        params = { search: '^nope' }
-        branches_finder = described_class.new(repository, params)
+      context 'by nonexistent name that begins with' do
+        let(:params) { { search: '^nope' } }
 
-        result = branches_finder.execute
+        it 'filters branches' do
+          result = subject
 
           expect(result.count).to eq(0)
         end
+      end
 
-      it 'filters branches by nonexistent name that ends with' do
-        params = { search: 'nope$' }
-        branches_finder = described_class.new(repository, params)
+      context 'by nonexistent name that ends with' do
+        let(:params) { { search: 'nope$' } }
 
-        result = branches_finder.execute
+        it 'filters branches' do
+          result = subject
 
           expect(result.count).to eq(0)
         end
       end
+    end
 
     context 'filter and sort' do
-      it 'filters branches by name and sorts by recently_updated' do
-        params = { sort: 'updated_desc', search: 'feat' }
-        branches_finder = described_class.new(repository, params)
+      context 'by name and sorts by recently_updated' do
+        let(:params) { { sort: 'updated_desc', search: 'feat' } }
 
-        result = branches_finder.execute
+        it 'filters branches' do
+          result = subject
 
           expect(result.first.name).to eq('feature_conflict')
           expect(result.count).to eq(2)
         end
+      end
 
-      it 'filters branches by name and sorts by recently_updated, with exact matches first' do
-        params = { sort: 'updated_desc', search: 'feature' }
-        branches_finder = described_class.new(repository, params)
+      context 'by name and sorts by recently_updated, with exact matches first' do
+        let(:params) { { sort: 'updated_desc', search: 'feature' } }
 
-        result = branches_finder.execute
+        it 'filters branches' do
+          result = subject
 
           expect(result.first.name).to eq('feature')
           expect(result.second.name).to eq('feature_conflict')
           expect(result.count).to eq(2)
         end
+      end
 
-      it 'filters branches by name and sorts by last_updated' do
-        params = { sort: 'updated_asc', search: 'feature' }
-        branches_finder = described_class.new(repository, params)
+      context 'by name and sorts by last_updated' do
+        let(:params) { { sort: 'updated_asc', search: 'feature' } }
 
-        result = branches_finder.execute
+        it 'filters branches' do
+          result = subject
 
           expect(result.first.name).to eq('feature')
           expect(result.count).to eq(2)
         end
       end
     end
+
+    context 'with gitaly pagination' do
+      subject { branch_finder.execute(gitaly_pagination: true) }
+
+      context 'by page_token and per_page' do
+        let(:params) { { page_token: 'feature', per_page: 2 } }
+
+        it 'filters branches' do
+          result = subject
+
+          expect(result.map(&:name)).to eq(%w(feature_conflict fix))
+        end
+      end
+
+      context 'by next page_token and per_page' do
+        let(:params) { { page_token: 'fix', per_page: 2 } }
+
+        it 'filters branches' do
+          result = subject
+
+          expect(result.map(&:name)).to eq(%w(flatten-dir gitattributes))
+        end
+      end
+
+      context 'by per_page only' do
+        let(:params) { { per_page: 2 } }
+
+        it 'filters branches' do
+          result = subject
+
+          expect(result.map(&:name)).to eq(["'test'", '2-mb-file'])
+        end
+      end
+
+      context 'by page_token only' do
+        let(:params) { { page_token: 'feature' } }
+
+        it 'returns nothing' do
+          result = subject
+
+          expect(result.count).to eq(0)
+        end
+      end
+
+      context 'pagination and sort' do
+        context 'by per_page' do
+          let(:params) { { sort: 'updated_asc', per_page: 5 } }
+
+          it 'filters branches' do
+            result = subject
+
+            expect(result.map(&:name)).to eq(%w(feature improve/awesome merge-test markdown feature_conflict))
+          end
+        end
+
+        context 'by page_token and per_page' do
+          let(:params) { { sort: 'updated_asc', page_token: 'improve/awesome', per_page: 2 } }
+
+          it 'filters branches' do
+            result = subject
+
+            expect(result.map(&:name)).to eq(%w(merge-test markdown))
+          end
+        end
+      end
+
+      context 'pagination and names' do
+        let(:params) { { page_token: 'fix', per_page: 2, names: %w[fix csv lfs does-not-exist] } }
+
+        it 'falls back to default execute and ignore paginations' do
+          result = subject
+
+          expect(result.count).to eq(3)
+          expect(result.map(&:name)).to eq(%w{csv fix lfs})
+        end
+      end
+
+      context 'pagination and search' do
+        let(:params) { { page_token: 'feature', per_page: 2, search: '^f' } }
+
+        it 'falls back to default execute and ignore paginations' do
+          result = subject
+
+          expect(result.map(&:name)).to eq(%w(feature feature_conflict fix flatten-dir))
+        end
+      end
+    end
+  end
 end

spec/requests/api/branches_spec.rb

@@ -17,6 +17,7 @@ RSpec.describe API::Branches do
   before do
     project.add_maintainer(user)
     project.repository.add_branch(user, 'ends-with.txt', branch_sha)
+    stub_feature_flags(branch_list_keyset_pagination: false)
   end
 
   describe "GET /projects/:id/repository/branches" do
@@ -29,6 +30,16 @@ RSpec.describe API::Branches do
         end
       end
 
+      def check_merge_status(json_response)
+        merged, unmerged = json_response.partition { |branch| branch['merged'] }
+        merged_branches = merged.map { |branch| branch['name'] }
+        unmerged_branches = unmerged.map { |branch| branch['name'] }
+        expect(Set.new(merged_branches)).to eq(project.repository.merged_branch_names(merged_branches + unmerged_branches))
+        expect(project.repository.merged_branch_names(unmerged_branches)).to be_empty
+      end
+
+      context 'with branch_list_keyset_pagination feature off' do
+        context 'with legacy pagination params' do
           it 'returns the repository branches' do
             get api(route, current_user), params: { per_page: 100 }
 
@@ -39,12 +50,58 @@ RSpec.describe API::Branches do
             expect(branch_names).to match_array(project.repository.branch_names)
           end
 
-      def check_merge_status(json_response)
-        merged, unmerged = json_response.partition { |branch| branch['merged'] }
-        merged_branches = merged.map { |branch| branch['name'] }
-        unmerged_branches = unmerged.map { |branch| branch['name'] }
-        expect(Set.new(merged_branches)).to eq(project.repository.merged_branch_names(merged_branches + unmerged_branches))
-        expect(project.repository.merged_branch_names(unmerged_branches)).to be_empty
+          it 'determines only a limited number of merged branch names' do
+            expect(API::Entities::Branch).to receive(:represent).with(anything, has_up_to_merged_branch_names_count(2)).and_call_original
+
+            get api(route, current_user), params: { per_page: 2 }
+
+            expect(response).to have_gitlab_http_status(:ok)
+            expect(json_response.count).to eq 2
+
+            check_merge_status(json_response)
+          end
+
+          it 'merge status matches reality on paginated input' do
+            expected_first_branch_name = project.repository.branches_sorted_by('name')[20].name
+
+            get api(route, current_user), params: { per_page: 20, page: 2 }
+
+            expect(response).to have_gitlab_http_status(:ok)
+            expect(json_response.count).to eq 20
+            expect(json_response.first['name']).to eq(expected_first_branch_name)
+
+            check_merge_status(json_response)
+          end
+        end
+
+        context 'with gitaly pagination params ' do
+          it 'merge status matches reality on paginated input' do
+            expected_first_branch_name = project.repository.branches_sorted_by('name').first.name
+
+            get api(route, current_user), params: { per_page: 20, page_token: 'feature' }
+
+            expect(response).to have_gitlab_http_status(:ok)
+            expect(json_response.count).to eq 20
+            expect(json_response.first['name']).to eq(expected_first_branch_name)
+
+            check_merge_status(json_response)
+          end
+        end
+      end
+
+      context 'with branch_list_keyset_pagination feature on' do
+        before do
+          stub_feature_flags(branch_list_keyset_pagination: true)
+        end
+
+        context 'with gitaly pagination params ' do
+          it 'returns the repository branches' do
+            get api(route, current_user), params: { per_page: 100 }
+
+            expect(response).to have_gitlab_http_status(:ok)
+            expect(response).to match_response_schema('public_api/v4/branches')
+            branch_names = json_response.map { |x| x['name'] }
+            expect(branch_names).to match_array(project.repository.branch_names)
           end
 
           it 'determines only a limited number of merged branch names' do
@@ -53,17 +110,36 @@ RSpec.describe API::Branches do
             get api(route, current_user), params: { per_page: 2 }
 
             expect(response).to have_gitlab_http_status(:ok)
+            expect(json_response.count).to eq 2
 
             check_merge_status(json_response)
           end
 
           it 'merge status matches reality on paginated input' do
+            expected_first_branch_name = project.repository.branches_sorted_by('name').drop_while { |b| b.name <= 'feature' }.first.name
+
+            get api(route, current_user), params: { per_page: 20, page_token: 'feature' }
+
+            expect(response).to have_gitlab_http_status(:ok)
+            expect(json_response.count).to eq 20
+            expect(json_response.first['name']).to eq(expected_first_branch_name)
+
+            check_merge_status(json_response)
+          end
+        end
+
+        context 'with legacy pagination params' do
+          it 'ignores legacy pagination params' do
+            expected_first_branch_name = project.repository.branches_sorted_by('name').first.name
             get api(route, current_user), params: { per_page: 20, page: 2 }
 
             expect(response).to have_gitlab_http_status(:ok)
+            expect(json_response.first['name']).to eq(expected_first_branch_name)
 
             check_merge_status(json_response)
           end
+        end
+      end
 
       context 'when repository is disabled' do
         include_context 'disabled repository'