Paginate first page of branches using Gitaly

Gitaly keyset pagination returns the same results as for offset
pagination. The first page is the most frequent request, so we
can use Gitaly to return the results when first page is requested.

Headers are added to imitate offset pagination.

changelogs/unreleased/id-use-gitaly-for-first-page-of-branches.yml

+---
+title: Paginate first page of branches using Gitaly
+merge_request: 48595
+author:
+type: performance

lib/gitlab/pagination/gitaly_keyset_pager.rb

@@ -15,6 +15,7 @@ module Gitlab
       # and supports pagination via gitaly.
       def paginate(finder)
         return paginate_via_gitaly(finder) if keyset_pagination_enabled?
+        return paginate_first_page_via_gitaly(finder) if paginate_first_page?
 
         branches = ::Kaminari.paginate_array(finder.execute)
         Gitlab::Pagination::OffsetPagination
@@ -28,12 +29,30 @@ module Gitlab
         Feature.enabled?(:branch_list_keyset_pagination, project) && params[:pagination] == 'keyset'
       end
 
+      def paginate_first_page?
+        Feature.enabled?(:branch_list_keyset_pagination, project) && (params[:page].blank? || params[:page].to_i == 1)
+      end
+
       def paginate_via_gitaly(finder)
         finder.execute(gitaly_pagination: true).tap do |records|
           apply_headers(records)
         end
       end
 
+      # When first page is requested, we paginate the data via Gitaly
+      # Headers are added to immitate offset pagination, while it is the default option
+      def paginate_first_page_via_gitaly(finder)
+        finder.execute(gitaly_pagination: true).tap do |records|
+          total = project.repository.branch_count
+          per_page = params[:per_page].presence || Kaminari.config.default_per_page
+
+          Gitlab::Pagination::OffsetHeaderBuilder.new(
+            request_context: request_context, per_page: per_page, page: 1, next_page: 2,
+            total: total, total_pages: total / per_page + 1
+          ).execute
+        end
+      end
+
       def apply_headers(records)
         if records.count == params[:per_page]
           Gitlab::Pagination::Keyset::HeaderBuilder

lib/gitlab/pagination/offset_header_builder.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Pagination
+    class OffsetHeaderBuilder
+      attr_reader :request_context, :per_page, :page, :next_page, :prev_page, :total, :total_pages
+
+      delegate :params, :header, :request, to: :request_context
+
+      def initialize(request_context:, per_page:, page:, next_page:, prev_page: nil, total:, total_pages:)
+        @request_context = request_context
+        @per_page = per_page
+        @page = page
+        @next_page = next_page
+        @prev_page = prev_page
+        @total = total
+        @total_pages = total_pages
+      end
+
+      def execute(exclude_total_headers: false, data_without_counts: false)
+        header 'X-Per-Page',    per_page.to_s
+        header 'X-Page',        page.to_s
+        header 'X-Next-Page',   next_page.to_s
+        header 'X-Prev-Page',   prev_page.to_s
+        header 'Link',          pagination_links(data_without_counts)
+
+        return if exclude_total_headers || data_without_counts
+
+        header 'X-Total',       total.to_s
+        header 'X-Total-Pages', total_pages.to_s
+      end
+
+      private
+
+      def pagination_links(data_without_counts)
+        [].tap do |links|
+          links << %(<#{page_href(page: prev_page)}>; rel="prev") if prev_page
+          links << %(<#{page_href(page: next_page)}>; rel="next") if next_page
+          links << %(<#{page_href(page: 1)}>; rel="first")
+
+          links << %(<#{page_href(page: total_pages)}>; rel="last") unless data_without_counts
+        end.join(', ')
+      end
+
+      def base_request_uri
+        @base_request_uri ||= URI.parse(request.url).tap do |uri|
+          uri.host = Gitlab.config.gitlab.host
+          uri.port = Gitlab.config.gitlab.port
+        end
+      end
+
+      def build_page_url(query_params:)
+        base_request_uri.tap do |uri|
+          uri.query = query_params
+        end.to_s
+      end
+
+      def page_href(next_page_params = {})
+        query_params = params.merge(**next_page_params, per_page: params[:per_page]).to_query
+
+        build_page_url(query_params: query_params)
+      end
+    end
+  end
+end

lib/gitlab/pagination/offset_pagination.rb

@@ -48,58 +48,26 @@ module Gitlab
       end
 
       def add_pagination_headers(paginated_data, exclude_total_headers)
-        header 'X-Per-Page',    paginated_data.limit_value.to_s
-        header 'X-Page',        paginated_data.current_page.to_s
-        header 'X-Next-Page',   paginated_data.next_page.to_s
-        header 'X-Prev-Page',   paginated_data.prev_page.to_s
-        header 'Link',          pagination_links(paginated_data)
-
-        return if exclude_total_headers || data_without_counts?(paginated_data)
-
-        header 'X-Total',       paginated_data.total_count.to_s
-        header 'X-Total-Pages', total_pages(paginated_data).to_s
-      end
-
-      def pagination_links(paginated_data)
-        [].tap do |links|
-          links << %(<#{page_href(page: paginated_data.prev_page)}>; rel="prev") if paginated_data.prev_page
-          links << %(<#{page_href(page: paginated_data.next_page)}>; rel="next") if paginated_data.next_page
-          links << %(<#{page_href(page: 1)}>; rel="first")
-
-          links << %(<#{page_href(page: total_pages(paginated_data))}>; rel="last") unless data_without_counts?(paginated_data)
-        end.join(', ')
-      end
-
-      def total_pages(paginated_data)
-        # Ensure there is in total at least 1 page
-        [paginated_data.total_pages, 1].max
+        Gitlab::Pagination::OffsetHeaderBuilder.new(
+          request_context: self, per_page: paginated_data.limit_value, page: paginated_data.current_page,
+          next_page: paginated_data.next_page, prev_page: paginated_data.prev_page,
+          total: total_count(paginated_data), total_pages: total_pages(paginated_data)
+        ).execute(exclude_total_headers: exclude_total_headers, data_without_counts: data_without_counts?(paginated_data))
       end
 
       def data_without_counts?(paginated_data)
         paginated_data.is_a?(Kaminari::PaginatableWithoutCount)
       end
 
-      def base_request_uri
-        @base_request_uri ||= URI.parse(request.url).tap do |uri|
-          uri.host = Gitlab.config.gitlab.host
-          uri.port = Gitlab.config.gitlab.port
-        end
+      def total_count(paginated_data)
+        paginated_data.total_count unless data_without_counts?(paginated_data)
       end
 
-      def build_page_url(query_params:)
-        base_request_uri.tap do |uri|
-          uri.query = query_params
-        end.to_s
-      end
-
-      def page_href(next_page_params = {})
-        query_params = params.merge(**next_page_params, per_page: per_page).to_query
-
-        build_page_url(query_params: query_params)
-      end
+      def total_pages(paginated_data)
+        return if data_without_counts?(paginated_data)
 
-      def per_page
-        @per_page ||= params[:per_page]
+        # Ensure there is in total at least 1 page
+        [paginated_data.total_pages, 1].max
       end
     end
   end

spec/lib/gitlab/pagination/gitaly_keyset_pager_spec.rb

@@ -57,17 +57,45 @@ RSpec.describe Gitlab::Pagination::GitalyKeysetPager do
     end
 
     context 'with branch_list_keyset_pagination feature on' do
+      let(:fake_request) { double(url: "#{incoming_api_projects_url}?#{query.to_query}") }
+      let(:branch1) { double 'branch', name: 'branch1' }
+      let(:branch2) { double 'branch', name: 'branch2' }
+      let(:branch3) { double 'branch', name: 'branch3' }
+
       before do
         stub_feature_flags(branch_list_keyset_pagination: project)
       end
 
       context 'without keyset pagination option' do
+        context 'when first page is requested' do
+          let(:branches) { [branch1, branch2, branch3] }
+
+          it 'keyset pagination is used with offset headers' do
+            allow(request_context).to receive(:request).and_return(fake_request)
+            allow(project.repository).to receive(:branch_count).and_return(branches.size)
+
+            expect(finder).to receive(:execute).with(gitaly_pagination: true).and_return(branches)
+            expect(request_context).to receive(:header).with('X-Per-Page', '2')
+            expect(request_context).to receive(:header).with('X-Page', '1')
+            expect(request_context).to receive(:header).with('X-Next-Page', '2')
+            expect(request_context).to receive(:header).with('X-Prev-Page', '')
+            expect(request_context).to receive(:header).with('Link', kind_of(String))
+            expect(request_context).to receive(:header).with('X-Total', '3')
+            expect(request_context).to receive(:header).with('X-Total-Pages', '2')
+
+            pager.paginate(finder)
+          end
+        end
+
+        context 'when second page is requested' do
+          let(:base_query) { { per_page: 2, page: 2 } }
+
           it_behaves_like 'offset pagination'
         end
+      end
 
       context 'with keyset pagination option' do
         let(:query) { base_query.merge(pagination: 'keyset') }
-        let(:fake_request) { double(url: "#{incoming_api_projects_url}?#{query.to_query}") }
 
         before do
           allow(request_context).to receive(:request).and_return(fake_request)
@@ -75,8 +103,6 @@ RSpec.describe Gitlab::Pagination::GitalyKeysetPager do
         end
 
         context 'when next page could be available' do
-          let(:branch1) { double 'branch', name: 'branch1' }
-          let(:branch2) { double 'branch', name: 'branch2' }
           let(:branches) { [branch1, branch2] }
 
           let(:expected_next_page_link) { %Q(<#{incoming_api_projects_url}?#{query.merge(page_token: branch2.name).to_query}>; rel="next") }
@@ -90,7 +116,6 @@ RSpec.describe Gitlab::Pagination::GitalyKeysetPager do
         end
 
         context 'when the current page is the last page' do
-          let(:branch1) { double 'branch', name: 'branch1' }
           let(:branches) { [branch1] }
 
           it 'uses keyset pagination without link headers' do

spec/lib/gitlab/pagination/offset_header_builder_spec.rb

+# frozen_string_literal: true
+
+require 'fast_spec_helper'
+
+RSpec.describe Gitlab::Pagination::OffsetHeaderBuilder do
+  let(:request) { double(url: 'http://localhost') }
+  let(:request_context) { double(header: nil, params: { per_page: 5 }, request: request) }
+
+  subject do
+    described_class.new(
+      request_context: request_context, per_page: 5, page: 2,
+      next_page: 3, prev_page: 1, total: 10, total_pages: 3
+    )
+  end
+
+  describe '#execute' do
+    let(:basic_links) do
+      %{<http://localhost?page=1&per_page=5>; rel="prev", <http://localhost?page=3&per_page=5>; rel="next", <http://localhost?page=1&per_page=5>; rel="first"}
+    end
+
+    let(:last_link) do
+      %{, <http://localhost?page=3&per_page=5>; rel="last"}
+    end
+
+    def expect_basic_headers
+      expect(request_context).to receive(:header).with('X-Per-Page', '5')
+      expect(request_context).to receive(:header).with('X-Page', '2')
+      expect(request_context).to receive(:header).with('X-Next-Page', '3')
+      expect(request_context).to receive(:header).with('X-Prev-Page', '1')
+      expect(request_context).to receive(:header).with('Link', basic_links + last_link)
+    end
+
+    it 'sets headers to request context' do
+      expect_basic_headers
+      expect(request_context).to receive(:header).with('X-Total', '10')
+      expect(request_context).to receive(:header).with('X-Total-Pages', '3')
+
+      subject.execute
+    end
+
+    context 'exclude total headers' do
+      it 'does not set total headers to request context' do
+        expect_basic_headers
+        expect(request_context).not_to receive(:header)
+
+        subject.execute(exclude_total_headers: true)
+      end
+    end
+
+    context 'pass data without counts' do
+      let(:last_link) { '' }
+
+      it 'does not set total headers to request context' do
+        expect_basic_headers
+        expect(request_context).not_to receive(:header)
+
+        subject.execute(data_without_counts: true)
+      end
+    end
+  end
+end