Merge branch 'id-use-gitaly-for-first-page-of-branches' into 'master'

Paginate first page of branches using Gitaly

See merge request gitlab-org/gitlab!48595

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
-        it_behaves_like 'offset pagination'
+        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