Use Gitaly keyset pagination to optimize branches page

Gitaly keyset pagination returns a limited number of
branches. We can use it instead of fetching all
branches and then paginating them manually.

Projects::BranchesByModeService is introduced.
It uses Gitaly page-token pagination in order to
optimally fetch branches.
The drawback of the page-token pagination is that
it doesn't provide an option of going to the previous
page of the collection.
That's why we need to fall back to offset pagination
when previous page is requested.

app/controllers/projects/branches_controller.rb

@@ -177,10 +177,8 @@ class Projects::BranchesController < Projects::ApplicationController
   def fetch_branches_by_mode
     return fetch_branches_for_overview if @mode == 'overview'
 
-    # active/stale/all view mode
-    @branches = BranchesFinder.new(@repository, params.merge(sort: @sort)).execute
-    @branches = @branches.select { |b| b.state.to_s == @mode } if %w[active stale].include?(@mode)
-    @branches = Kaminari.paginate_array(@branches).page(params[:page])
+    @branches, @prev_path, @next_path =
+      Projects::BranchesByModeService.new(@project, params.merge(sort: @sort, mode: @mode)).execute
   end
 
   def fetch_branches_for_overview

app/services/projects/branches_by_mode_service.rb

+# frozen_string_literal: true
+
+# Projects::BranchesByModeService uses Gitaly page-token pagination
+# in order to optimally fetch branches.
+# The drawback of the page-token pagination is that it doesn't provide
+# an option of going to the previous page of the collection.
+# That's why we need to fall back to offset pagination when previous page
+# is requested.
+class Projects::BranchesByModeService
+  include Gitlab::Routing
+
+  attr_reader :project, :params
+
+  def initialize(project, params = {})
+    @project = project
+    @params = params
+  end
+
+  def execute
+    return fetch_branches_via_gitaly_pagination if use_gitaly_pagination?
+
+    fetch_branches_via_offset_pagination
+  end
+
+  private
+
+  def mode
+    params[:mode]
+  end
+
+  def by_mode(branches)
+    return branches unless %w[active stale].include?(mode)
+
+    branches.select { |b| b.state.to_s == mode }
+  end
+
+  def use_gitaly_pagination?
+    return false if params[:page].present? || params[:search].present?
+
+    Feature.enabled?(:branch_list_keyset_pagination, project, default_enabled: true)
+  end
+
+  def fetch_branches_via_offset_pagination
+    branches = BranchesFinder.new(project.repository, params).execute
+    branches = Kaminari.paginate_array(by_mode(branches)).page(params[:page])
+
+    branches_with_links(branches, last_page: branches.last_page?)
+  end
+
+  def fetch_branches_via_gitaly_pagination
+    per_page = Kaminari.config.default_per_page
+    options = params.merge(per_page: per_page + 1, page_token: params[:page_token])
+
+    branches = BranchesFinder.new(project.repository, options).execute(gitaly_pagination: true)
+
+    # Branch is stale if it hasn't been updated for 3 months
+    # This logic is specified in Gitlab Rails and isn't specified in Gitaly
+    # To display stale branches we fetch branches sorted as most-stale-at-the-top
+    # If the result contains active branches we filter them out and define that no more stale branches left
+    # Same logic applies to fetching active branches
+    branches = by_mode(branches)
+    last_page = branches.size <= per_page
+
+    branches = branches.take(per_page) # rubocop:disable CodeReuse/ActiveRecord
+
+    branches_with_links(branches, last_page: last_page)
+  end
+
+  def branches_with_links(branches, last_page:)
+    # To fall back to offset pagination we need to track current page via offset param
+    # And increase it whenever we go to the next page
+    previous_offset = params[:offset].to_i
+
+    previous_path, next_path = nil, nil
+
+    return [branches, previous_path, next_path] if branches.blank?
+
+    unless last_page
+      next_path = project_branches_filtered_path(project, state: mode, page_token: branches.last.name, sort: params[:sort], offset: previous_offset + 1)
+    end
+
+    if previous_offset > 0
+      previous_path = project_branches_filtered_path(project, state: mode, sort: params[:sort], page: previous_offset, offset: previous_offset - 1)
+    end
+
+    [branches, previous_path, next_path]
+  end
+end

app/views/projects/branches/index.html.haml

@@ -61,7 +61,7 @@
     - @branches.each do |branch|
       = render "projects/branches/branch", branch: branch, merged: @merged_branch_names.include?(branch.name), commit_status: @branch_pipeline_statuses[branch.name], show_commit_status: @branch_pipeline_statuses.any?
   - if Feature.enabled?(:branches_pagination_without_count, @project, default_enabled: true)
-    = paginate_without_count @branches
+    = render('kaminari/gitlab/without_count', previous_path: @prev_path, next_path: @next_path)
   - else
     = paginate @branches, theme: 'gitlab'
 - else

changelogs/unreleased/id-optimize-branches-index.yml

+---
+title: Use Gitaly keyset pagination to optimize branches page
+merge_request: 53409
+author:
+type: performance

spec/services/projects/branches_by_mode_service_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Projects::BranchesByModeService do
+  let_it_be(:user) { create(:user) }
+  let_it_be(:project) { create(:project, :repository) }
+
+  let(:finder) { described_class.new(project, params) }
+  let(:params) { { mode: 'all' } }
+
+  subject { finder.execute }
+
+  describe '#execute' do
+    context 'page is passed' do
+      let(:params) { { page: 4, mode: 'all', offset: 3 } }
+
+      it 'uses offset pagination' do
+        expect(finder).to receive(:fetch_branches_via_offset_pagination).and_call_original
+
+        branches, prev_page, next_page = subject
+
+        expect(branches.size).to eq(10)
+        expect(next_page).to be_nil
+        expect(prev_page).to eq("/#{project.full_path}/-/branches/all?offset=2&page=3")
+      end
+
+      context 'but the page does not contain any branches' do
+        let(:params) { { page: 10, mode: 'all' } }
+
+        it 'uses offset pagination' do
+          expect(finder).to receive(:fetch_branches_via_offset_pagination).and_call_original
+
+          branches, prev_page, next_page = subject
+
+          expect(branches).to eq([])
+          expect(next_page).to be_nil
+          expect(prev_page).to be_nil
+        end
+      end
+    end
+
+    context 'search is passed' do
+      let(:params) { { search: 'feature' } }
+
+      it 'uses offset pagination' do
+        expect(finder).to receive(:fetch_branches_via_offset_pagination).and_call_original
+
+        branches, prev_page, next_page = subject
+
+        expect(branches.map(&:name)).to match_array(%w(feature feature_conflict))
+        expect(next_page).to be_nil
+        expect(prev_page).to be_nil
+      end
+    end
+
+    context 'branch_list_keyset_pagination is disabled' do
+      it 'uses offset pagination' do
+        stub_feature_flags(branch_list_keyset_pagination: false)
+
+        expect(finder).to receive(:fetch_branches_via_offset_pagination).and_call_original
+
+        branches, prev_page, next_page = subject
+
+        expect(branches.size).to eq(20)
+        expect(next_page).to eq("/#{project.full_path}/-/branches/all?offset=1&page_token=conflict-resolvable")
+        expect(prev_page).to be_nil
+      end
+    end
+
+    context 'uses gitaly pagination' do
+      before do
+        expect(finder).to receive(:fetch_branches_via_gitaly_pagination).and_call_original
+      end
+
+      it 'returns branches for the first page' do
+        branches, prev_page, next_page = subject
+
+        expect(branches.size).to eq(20)
+        expect(next_page).to eq("/#{project.full_path}/-/branches/all?offset=1&page_token=conflict-resolvable")
+        expect(prev_page).to be_nil
+      end
+
+      context 'when second page is requested' do
+        let(:params) { { page_token: 'conflict-resolvable', mode: 'all', sort: 'name_asc', offset: 1 } }
+
+        it 'returns branches for the first page' do
+          branches, prev_page, next_page = subject
+
+          expect(branches.size).to eq(20)
+          expect(next_page).to eq("/#{project.full_path}/-/branches/all?offset=2&page_token=improve%2Fawesome&sort=name_asc")
+          expect(prev_page).to eq("/#{project.full_path}/-/branches/all?offset=0&page=1&sort=name_asc")
+        end
+      end
+
+      context 'when last page is requested' do
+        let(:params) { { page_token: 'signed-commits', mode: 'all', sort: 'name_asc', offset: 4 } }
+
+        it 'returns branches after the specified branch' do
+          branches, prev_page, next_page = subject
+
+          expect(branches.size).to eq(14)
+          expect(next_page).to be_nil
+          expect(prev_page).to eq("/#{project.full_path}/-/branches/all?offset=3&page=4&sort=name_asc")
+        end
+      end
+    end
+
+    context 'filter by mode' do
+      let(:stale) { double(state: 'stale') }
+      let(:active) { double(state: 'active') }
+
+      before do
+        allow_next_instance_of(BranchesFinder) do |instance|
+          allow(instance).to receive(:execute).and_return([stale, active])
+        end
+      end
+
+      context 'stale' do
+        let(:params) { { mode: 'stale' } }
+
+        it 'returns stale branches' do
+          is_expected.to eq([[stale], nil, nil])
+        end
+      end
+
+      context 'active' do
+        let(:params) { { mode: 'active' } }
+
+        it 'returns active branches' do
+          is_expected.to eq([[active], nil, nil])
+        end
+      end
+    end
+  end
+end