Gitaly repository tree keyset pagination

Adds support for keyset pagination on the repository
tree API. This implements a new finder to work alongside
the GitalyKeysetPager.

Changelog: added

app/finders/branches_finder.rb

@@ -15,6 +15,10 @@ class BranchesFinder < GitRefsFinder
     end
   end
 
+  def total
+    repository.branch_count
+  end
+
   private
 
   def names

app/finders/repositories/tree_finder.rb

+# frozen_string_literal: true
+
+module Repositories
+  class TreeFinder < GitRefsFinder
+    attr_reader :user_project
+
+    CommitMissingError = Class.new(StandardError)
+
+    def initialize(user_project, params = {})
+      super(user_project.repository, params)
+
+      @user_project = user_project
+    end
+
+    def execute(gitaly_pagination: false)
+      raise CommitMissingError unless commit_exists?
+
+      request_params = { recursive: recursive }
+      request_params[:pagination_params] = pagination_params if gitaly_pagination
+      tree = user_project.repository.tree(commit.id, path, **request_params)
+
+      tree.sorted_entries
+    end
+
+    def total
+      # This is inefficient and we'll look at replacing this implementation
+      Gitlab::Cache.fetch_once([user_project, repository.commit, :tree_size, commit.id, path, recursive]) do
+        user_project.repository.tree(commit.id, path, recursive: recursive).entries.size
+      end
+    end
+
+    def commit_exists?
+      commit.present?
+    end
+
+    private
+
+    def commit
+      @commit ||= user_project.commit(ref)
+    end
+
+    def ref
+      params[:ref] || user_project.default_branch
+    end
+
+    def path
+      params[:path]
+    end
+
+    def recursive
+      params[:recursive]
+    end
+
+    def pagination_params
+      {
+        limit: params[:per_page] || Kaminari.config.default_per_page,
+        page_token: params[:page_token]
+      }
+    end
+  end
+end

config/feature_flags/development/repository_tree_gitaly_pagination.yml

+---
+name: repository_tree_gitaly_pagination
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/67509
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/340419
+milestone: '14.3'
+type: development
+group: group::source code
+default_enabled: false

lib/api/repositories.rb

@@ -51,18 +51,22 @@ module API
         optional :ref, type: String, desc: 'The name of a repository branch or tag, if not given the default branch is used'
         optional :path, type: String, desc: 'The path of the tree'
         optional :recursive, type: Boolean, default: false, desc: 'Used to get a recursive tree'
+
         use :pagination
+        optional :pagination, type: String, values: %w(legacy keyset), default: 'legacy', desc: 'Specify the pagination method'
+
+        given pagination: -> (value) { value == 'keyset' } do
+          optional :page_token, type: String, desc: 'Record from which to start the keyset pagination'
+        end
       end
       get ':id/repository/tree' do
-        ref = params[:ref] || user_project.default_branch
-        path = params[:path] || nil
+        tree_finder = ::Repositories::TreeFinder.new(user_project, declared_params(include_missing: false))
+
+        not_found!("Tree") unless tree_finder.commit_exists?
 
-        commit = user_project.commit(ref)
-        not_found!('Tree') unless commit
+        tree = Gitlab::Pagination::GitalyKeysetPager.new(self, user_project).paginate(tree_finder)
 
-        tree = user_project.repository.tree(commit.id, path, recursive: params[:recursive])
-        entries = ::Kaminari.paginate_array(tree.sorted_entries)
-        present paginate(entries), with: Entities::TreeObject
+        present tree, with: Entities::TreeObject
       end
 
       desc 'Get raw blob contents from the repository'

lib/gitlab/gitaly_client/commit_service.rb

@@ -127,6 +127,7 @@ module Gitlab
 
         entries = response.flat_map do |message|
           cursor = message.pagination_cursor if message.pagination_cursor
+
           message.entries.map do |gitaly_tree_entry|
             Gitlab::Git::Tree.new(
               id: gitaly_tree_entry.oid,

lib/gitlab/pagination/gitaly_keyset_pager.rb

@@ -14,23 +14,39 @@ module Gitlab
       # It is expected that the given finder will respond to `execute` method with `gitaly_pagination: true` option
       # 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?
+        return paginate_via_gitaly(finder) if keyset_pagination_enabled?(finder)
+        return paginate_first_page_via_gitaly(finder) if paginate_first_page?(finder)
 
-        branches = ::Kaminari.paginate_array(finder.execute)
+        records = ::Kaminari.paginate_array(finder.execute)
         Gitlab::Pagination::OffsetPagination
           .new(request_context)
-          .paginate(branches)
+          .paginate(records)
       end
 
       private
 
-      def keyset_pagination_enabled?
-        Feature.enabled?(:branch_list_keyset_pagination, project, default_enabled: :yaml) && params[:pagination] == 'keyset'
+      def keyset_pagination_enabled?(finder)
+        return false unless params[:pagination] == "keyset"
+
+        if finder.is_a?(BranchesFinder)
+          Feature.enabled?(:branch_list_keyset_pagination, project, default_enabled: :yaml)
+        elsif finder.is_a?(::Repositories::TreeFinder)
+          Feature.enabled?(:repository_tree_gitaly_pagination, project, default_enabled: :yaml)
+        else
+          false
+        end
       end
 
-      def paginate_first_page?
-        Feature.enabled?(:branch_list_keyset_pagination, project, default_enabled: :yaml) && (params[:page].blank? || params[:page].to_i == 1)
+      def paginate_first_page?(finder)
+        return false unless params[:page].blank? || params[:page].to_i == 1
+
+        if finder.is_a?(BranchesFinder)
+          Feature.enabled?(:branch_list_keyset_pagination, project, default_enabled: :yaml)
+        elsif finder.is_a?(::Repositories::TreeFinder)
+          Feature.enabled?(:repository_tree_gitaly_pagination, project, default_enabled: :yaml)
+        else
+          false
+        end
       end
 
       def paginate_via_gitaly(finder)
@@ -43,7 +59,7 @@ module Gitlab
       # 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
+          total = finder.total
           per_page = params[:per_page].presence || Kaminari.config.default_per_page
 
           Gitlab::Pagination::OffsetHeaderBuilder.new(

spec/finders/branches_finder_spec.rb

@@ -259,4 +259,11 @@ RSpec.describe BranchesFinder do
       end
     end
   end
+
+  describe '#total' do
+    subject { branch_finder.total }
+
+    it { is_expected.to be_an(Integer) }
+    it { is_expected.to eq(repository.branch_count) }
+  end
 end

spec/finders/repositories/tree_finder_spec.rb

+# frozen_string_literal: true
+
+require "spec_helper"
+
+RSpec.describe Repositories::TreeFinder do
+  include RepoHelpers
+
+  let_it_be(:user) { create(:user) }
+  let_it_be(:project) { create(:project, :repository, creator: user) }
+
+  let(:repository) { project.repository }
+  let(:tree_finder) { described_class.new(project, params) }
+  let(:params) { {} }
+  let(:first_page_ids) { tree_finder.execute.map(&:id) }
+  let(:second_page_token) { first_page_ids.last }
+
+  describe "#execute" do
+    subject { tree_finder.execute(gitaly_pagination: true) }
+
+    it "returns an array" do
+      is_expected.to be_an(Array)
+    end
+
+    it "includes 20 items by default" do
+      expect(subject.size).to eq(20)
+    end
+
+    it "accepts a gitaly_pagination argument" do
+      expect(repository).to receive(:tree).with(anything, anything, recursive: nil, pagination_params: { limit: 20, page_token: nil }).and_call_original
+      expect(tree_finder.execute(gitaly_pagination: true)).to be_an(Array)
+
+      expect(repository).to receive(:tree).with(anything, anything, recursive: nil).and_call_original
+      expect(tree_finder.execute(gitaly_pagination: false)).to be_an(Array)
+    end
+
+    context "commit doesn't exist" do
+      let(:params) do
+        { ref: "nonesuchref" }
+      end
+
+      it "raises an error" do
+        expect { subject }.to raise_error(described_class::CommitMissingError)
+      end
+    end
+
+    describe "pagination_params" do
+      let(:params) do
+        { per_page: 5, page_token: nil }
+      end
+
+      it "has the per_page number of items" do
+        expect(subject.size).to eq(5)
+      end
+
+      it "doesn't include any of the first page records" do
+        first_page_ids = subject.map(&:id)
+        second_page = described_class.new(project, { per_page: 5, page_token: first_page_ids.last }).execute(gitaly_pagination: true)
+
+        expect(second_page.map(&:id)).not_to include(*first_page_ids)
+      end
+    end
+  end
+
+  describe "#total", :use_clean_rails_memory_store_caching do
+    subject { tree_finder.total }
+
+    it { is_expected.to be_an(Integer) }
+
+    it "only calculates the total once" do
+      expect(repository).to receive(:tree).once.and_call_original
+
+      2.times { tree_finder.total }
+    end
+  end
+
+  describe "#commit_exists?" do
+    subject { tree_finder.commit_exists? }
+
+    context "ref exists" do
+      let(:params) do
+        { ref: project.default_branch }
+      end
+
+      it { is_expected.to be(true) }
+    end
+
+    context "ref is missing" do
+      let(:params) do
+        { ref: "nonesuchref" }
+      end
+
+      it { is_expected.to be(false) }
+    end
+  end
+end

spec/lib/gitlab/pagination/gitaly_keyset_pager_spec.rb

@@ -74,7 +74,7 @@ RSpec.describe Gitlab::Pagination::GitalyKeysetPager 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(finder).to receive(:execute).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')
@@ -99,6 +99,7 @@ RSpec.describe Gitlab::Pagination::GitalyKeysetPager do
 
         before do
           allow(request_context).to receive(:request).and_return(fake_request)
+          allow(finder).to receive(:is_a?).with(BranchesFinder) { true }
           expect(finder).to receive(:execute).with(gitaly_pagination: true).and_return(branches)
         end
 

spec/requests/api/repositories_spec.rb

@@ -22,7 +22,7 @@ RSpec.describe API::Repositories do
 
         expect(response).to have_gitlab_http_status(:ok)
         expect(response).to include_pagination_headers
-        expect(json_response).to be_an Array
+        expect(json_response).to be_an(Array)
 
         first_commit = json_response.first
         expect(first_commit['name']).to eq('bar')
@@ -73,6 +73,25 @@ RSpec.describe API::Repositories do
           end
         end
       end
+
+      context 'keyset pagination mode' do
+        let(:first_response) do
+          get api(route, current_user), params: { pagination: "keyset" }
+
+          Gitlab::Json.parse(response.body)
+        end
+
+        it 'paginates using keysets' do
+          page_token = first_response.last["id"]
+
+          get api(route, current_user), params: { pagination: "keyset", page_token: page_token }
+
+          expect(response).to have_gitlab_http_status(:ok)
+          expect(json_response).to be_an(Array)
+          expect(json_response).not_to eq(first_response)
+          expect(json_response.map { |t| t["id"] }).not_to include(page_token)
+        end
+      end
     end
 
     context 'when unauthenticated', 'and project is public' do