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,11 +5,15 @@ class BranchesFinder < GitRefsFinder
     super(repository, params)
   end
 
-  def execute
-    branches = repository.branches_sorted_by(sort)
-    branches = by_search(branches)
-    branches = by_names(branches)
-    branches
+  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
 
-        branches = BranchesFinder.new(repository, declared_params(include_missing: false)).execute
-        branches = paginate(::Kaminari.paginate_array(branches))
+        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/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,16 +30,6 @@ RSpec.describe API::Branches do
         end
       end
 
-      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')
-        expect(response).to include_pagination_headers
-        branch_names = json_response.map { |x| x['name'] }
-        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'] }
@@ -47,22 +38,107 @@ RSpec.describe API::Branches do
         expect(project.repository.merged_branch_names(unmerged_branches)).to be_empty
       end
 
-      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
+      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 }
 
-        get api(route, current_user), params: { per_page: 2 }
+            expect(response).to have_gitlab_http_status(:ok)
+            expect(response).to match_response_schema('public_api/v4/branches')
+            expect(response).to include_pagination_headers
+            branch_names = json_response.map { |x| x['name'] }
+            expect(branch_names).to match_array(project.repository.branch_names)
+          end
 
-        expect(response).to have_gitlab_http_status(:ok)
+          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
 
-        check_merge_status(json_response)
+          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
 
-      it 'merge status matches reality on paginated input' do
-        get api(route, current_user), params: { per_page: 20, page: 2 }
+      context 'with branch_list_keyset_pagination feature on' do
+        before do
+          stub_feature_flags(branch_list_keyset_pagination: true)
+        end
 
-        expect(response).to have_gitlab_http_status(:ok)
+        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
+            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
 
-        check_merge_status(json_response)
+          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