From 29ce3cb8ced81e82489b93a6d2cf858f099bf43a Mon Sep 17 00:00:00 2001
From: Robert May <rmay@gitlab.com>
Date: Wed, 4 Aug 2021 14:06:32 +0100
Subject: [PATCH 1/9] 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/repositories/tree_finder.rb       | 60 ++++++++++++++
 .../repository_tree_gitaly_pagination.yml     |  8 ++
 lib/api/repositories.rb                       | 14 ++--
 lib/gitlab/pagination/gitaly_keyset_pager.rb  | 34 +++++---
 spec/finders/repositories/tree_finder_spec.rb | 80 +++++++++++++++++++
 spec/requests/api/repositories_spec.rb        | 21 ++++-
 6 files changed, 200 insertions(+), 17 deletions(-)
 create mode 100644 app/finders/repositories/tree_finder.rb
 create mode 100644 config/feature_flags/development/repository_tree_gitaly_pagination.yml
 create mode 100644 spec/finders/repositories/tree_finder_spec.rb

diff --git a/app/finders/repositories/tree_finder.rb b/app/finders/repositories/tree_finder.rb
new file mode 100644
index 00000000000000..6dfd6249da0029
--- /dev/null
+++ b/app/finders/repositories/tree_finder.rb
@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+class Repositories::TreeFinder < GitRefsFinder
+  attr_reader :user_project
+
+  CommitMissingError = Class.new(StandardError)
+
+  def initialize(user_project, params = {})
+    @user_project = user_project
+
+    super(user_project.repository, params)
+  end
+
+  def execute(gitaly_pagination: true)
+    raise CommitMissingError unless commit_exists?
+
+    tree =
+      if gitaly_pagination
+        user_project.repository.tree(commit.id, path, recursive: recursive, pagination_params: pagination_params)
+      else
+        user_project.repository.tree(commit.id, path, recursive: recursive)
+      end
+
+    tree.sorted_entries
+  end
+
+  def commit_exists?
+    commit.present?
+  end
+
+  private
+
+  def ref
+    @params[:ref] || @user_project.default_branch
+  end
+
+  def path
+    @params[:path] || nil
+  end
+
+  def recursive
+    @params[:recursive]
+  end
+
+  def commit
+    @user_project.commit(ref)
+  end
+
+  def per_page
+    @params[:per_page] || 20
+  end
+
+  def page_token
+    @params[:page_token]
+  end
+
+  def pagination_params
+    { limit: per_page, page_token: page_token }
+  end
+end
diff --git a/config/feature_flags/development/repository_tree_gitaly_pagination.yml b/config/feature_flags/development/repository_tree_gitaly_pagination.yml
new file mode 100644
index 00000000000000..5dcea66862b0ed
--- /dev/null
+++ b/config/feature_flags/development/repository_tree_gitaly_pagination.yml
@@ -0,0 +1,8 @@
+---
+name: repository_tree_gitaly_pagination
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/67509
+rollout_issue_url: 
+milestone: '14.2'
+type: development
+group: group::source code
+default_enabled: false
diff --git a/lib/api/repositories.rb b/lib/api/repositories.rb
index 20320d1b7ae88f..611bbb088b9de3 100644
--- a/lib/api/repositories.rb
+++ b/lib/api/repositories.rb
@@ -51,18 +51,18 @@ def assign_blob_vars!(limit:)
         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 :page_token, type: String, desc: 'Record from which to start the pagination'
       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'
diff --git a/lib/gitlab/pagination/gitaly_keyset_pager.rb b/lib/gitlab/pagination/gitaly_keyset_pager.rb
index b05891066acae4..17946e1b7d9e4e 100644
--- a/lib/gitlab/pagination/gitaly_keyset_pager.rb
+++ b/lib/gitlab/pagination/gitaly_keyset_pager.rb
@@ -14,23 +14,39 @@ def initialize(request_context, project)
       # 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 @@ def paginate_via_gitaly(finder)
       # 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 = project.repository.branch_count # FIXME: needs to be generic
           per_page = params[:per_page].presence || Kaminari.config.default_per_page
 
           Gitlab::Pagination::OffsetHeaderBuilder.new(
diff --git a/spec/finders/repositories/tree_finder_spec.rb b/spec/finders/repositories/tree_finder_spec.rb
new file mode 100644
index 00000000000000..a9a735baee5168
--- /dev/null
+++ b/spec/finders/repositories/tree_finder_spec.rb
@@ -0,0 +1,80 @@
+# 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 }
+
+    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 for compatibility" do
+      expect(tree_finder.execute(gitaly_pagination: true)).to be_an(Array)
+      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 = tree_finder.execute.map(&:id)
+        second_page = described_class.new(project, { per_page: 5, page_token: first_page_ids.last }).execute
+
+        expect(second_page.map(&:id)).not_to include(*first_page_ids)
+      end
+    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
diff --git a/spec/requests/api/repositories_spec.rb b/spec/requests/api/repositories_spec.rb
index 1b90d6fa47e87f..a576e1ab1ee236 100644
--- a/spec/requests/api/repositories_spec.rb
+++ b/spec/requests/api/repositories_spec.rb
@@ -22,7 +22,7 @@
 
         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 @@
           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
-- 
GitLab


From e9ac69abf2d6e7b4ba56c605f2e5ba5c63b3d9b4 Mon Sep 17 00:00:00 2001
From: Robert May <rmay@gitlab.com>
Date: Thu, 26 Aug 2021 14:53:23 +0100
Subject: [PATCH 2/9] Tweaks and new total system

---
 app/finders/branches_finder.rb                |  4 ++++
 app/finders/repositories/tree_finder.rb       |  7 ++++++-
 lib/gitlab/gitaly_client/commit_service.rb    |  1 +
 lib/gitlab/pagination/gitaly_keyset_pager.rb  |  2 +-
 spec/finders/branches_finder_spec.rb          |  7 +++++++
 spec/finders/repositories/tree_finder_spec.rb | 11 +++++++----
 6 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/app/finders/branches_finder.rb b/app/finders/branches_finder.rb
index 157c454183a6fb..a62d47071d4972 100644
--- a/app/finders/branches_finder.rb
+++ b/app/finders/branches_finder.rb
@@ -15,6 +15,10 @@ def execute(gitaly_pagination: false)
     end
   end
 
+  def total
+    repository.branch_count
+  end
+
   private
 
   def names
diff --git a/app/finders/repositories/tree_finder.rb b/app/finders/repositories/tree_finder.rb
index 6dfd6249da0029..2f9fa5809e53c9 100644
--- a/app/finders/repositories/tree_finder.rb
+++ b/app/finders/repositories/tree_finder.rb
@@ -11,7 +11,7 @@ def initialize(user_project, params = {})
     super(user_project.repository, params)
   end
 
-  def execute(gitaly_pagination: true)
+  def execute(gitaly_pagination: false)
     raise CommitMissingError unless commit_exists?
 
     tree =
@@ -24,6 +24,11 @@ def execute(gitaly_pagination: true)
     tree.sorted_entries
   end
 
+  def total
+    # TODO: No idea how to calculate this
+    100
+  end
+
   def commit_exists?
     commit.present?
   end
diff --git a/lib/gitlab/gitaly_client/commit_service.rb b/lib/gitlab/gitaly_client/commit_service.rb
index ddf04bd8e79d79..75588ad980ca6a 100644
--- a/lib/gitlab/gitaly_client/commit_service.rb
+++ b/lib/gitlab/gitaly_client/commit_service.rb
@@ -127,6 +127,7 @@ def tree_entries(repository, revision, path, recursive, pagination_params)
 
         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,
diff --git a/lib/gitlab/pagination/gitaly_keyset_pager.rb b/lib/gitlab/pagination/gitaly_keyset_pager.rb
index 17946e1b7d9e4e..a16bf7a379ca00 100644
--- a/lib/gitlab/pagination/gitaly_keyset_pager.rb
+++ b/lib/gitlab/pagination/gitaly_keyset_pager.rb
@@ -59,7 +59,7 @@ def paginate_via_gitaly(finder)
       # 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 # FIXME: needs to be generic
+          total = finder.total
           per_page = params[:per_page].presence || Kaminari.config.default_per_page
 
           Gitlab::Pagination::OffsetHeaderBuilder.new(
diff --git a/spec/finders/branches_finder_spec.rb b/spec/finders/branches_finder_spec.rb
index a62dd3842dbe9b..f9d525c33a409e 100644
--- a/spec/finders/branches_finder_spec.rb
+++ b/spec/finders/branches_finder_spec.rb
@@ -259,4 +259,11 @@
       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
diff --git a/spec/finders/repositories/tree_finder_spec.rb b/spec/finders/repositories/tree_finder_spec.rb
index a9a735baee5168..32ac32c9f66816 100644
--- a/spec/finders/repositories/tree_finder_spec.rb
+++ b/spec/finders/repositories/tree_finder_spec.rb
@@ -15,7 +15,7 @@
   let(:second_page_token) { first_page_ids.last }
 
   describe "#execute" do
-    subject { tree_finder.execute }
+    subject { tree_finder.execute(gitaly_pagination: true) }
 
     it "returns an array" do
       is_expected.to be_an(Array)
@@ -25,8 +25,11 @@
       expect(subject.size).to eq(20)
     end
 
-    it "accepts a gitaly_pagination argument for compatibility" do
+    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
 
@@ -50,8 +53,8 @@
       end
 
       it "doesn't include any of the first page records" do
-        first_page_ids = tree_finder.execute.map(&:id)
-        second_page = described_class.new(project, { per_page: 5, page_token: first_page_ids.last }).execute
+        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
-- 
GitLab


From 51bdaa15d6b557d6bb8ce4351d7eeffdc59d03b1 Mon Sep 17 00:00:00 2001
From: Robert May <rmay@gitlab.com>
Date: Fri, 27 Aug 2021 14:27:20 +0100
Subject: [PATCH 3/9] Implement temporary total calculation

---
 app/finders/repositories/tree_finder.rb       |  6 ++++--
 spec/finders/repositories/tree_finder_spec.rb | 12 ++++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/app/finders/repositories/tree_finder.rb b/app/finders/repositories/tree_finder.rb
index 2f9fa5809e53c9..132b991486b32c 100644
--- a/app/finders/repositories/tree_finder.rb
+++ b/app/finders/repositories/tree_finder.rb
@@ -25,8 +25,10 @@ def execute(gitaly_pagination: false)
   end
 
   def total
-    # TODO: No idea how to calculate this
-    100
+    # This is inefficient and we'll look at replacing this implementation
+    Gitlab::Cache.fetch_once([user_project, repository.commit, :tree_size, @params]) do
+      user_project.repository.tree(commit.id, path, recursive: recursive).sorted_entries.size
+    end
   end
 
   def commit_exists?
diff --git a/spec/finders/repositories/tree_finder_spec.rb b/spec/finders/repositories/tree_finder_spec.rb
index 32ac32c9f66816..0d70d5f92d3423 100644
--- a/spec/finders/repositories/tree_finder_spec.rb
+++ b/spec/finders/repositories/tree_finder_spec.rb
@@ -61,6 +61,18 @@
     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? }
 
-- 
GitLab


From 8c07709f289a47d287d3538f1f940bb525d623eb Mon Sep 17 00:00:00 2001
From: Robert May <rmay@gitlab.com>
Date: Fri, 27 Aug 2021 14:32:18 +0100
Subject: [PATCH 4/9] Style change for class name

---
 app/finders/repositories/tree_finder.rb | 94 +++++++++++++------------
 1 file changed, 48 insertions(+), 46 deletions(-)

diff --git a/app/finders/repositories/tree_finder.rb b/app/finders/repositories/tree_finder.rb
index 132b991486b32c..9631648977f361 100644
--- a/app/finders/repositories/tree_finder.rb
+++ b/app/finders/repositories/tree_finder.rb
@@ -1,67 +1,69 @@
 # frozen_string_literal: true
 
-class Repositories::TreeFinder < GitRefsFinder
-  attr_reader :user_project
+module Repositories
+  class TreeFinder < GitRefsFinder
+    attr_reader :user_project
 
-  CommitMissingError = Class.new(StandardError)
+    CommitMissingError = Class.new(StandardError)
 
-  def initialize(user_project, params = {})
-    @user_project = user_project
+    def initialize(user_project, params = {})
+      @user_project = user_project
 
-    super(user_project.repository, params)
-  end
+      super(user_project.repository, params)
+    end
 
-  def execute(gitaly_pagination: false)
-    raise CommitMissingError unless commit_exists?
+    def execute(gitaly_pagination: false)
+      raise CommitMissingError unless commit_exists?
 
-    tree =
-      if gitaly_pagination
-        user_project.repository.tree(commit.id, path, recursive: recursive, pagination_params: pagination_params)
-      else
-        user_project.repository.tree(commit.id, path, recursive: recursive)
-      end
+      tree =
+        if gitaly_pagination
+          user_project.repository.tree(commit.id, path, recursive: recursive, pagination_params: pagination_params)
+        else
+          user_project.repository.tree(commit.id, path, recursive: recursive)
+        end
 
-    tree.sorted_entries
-  end
+      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, @params]) do
-      user_project.repository.tree(commit.id, path, recursive: recursive).sorted_entries.size
+    def total
+      # This is inefficient and we'll look at replacing this implementation
+      Gitlab::Cache.fetch_once([user_project, repository.commit, :tree_size, @params]) do
+        user_project.repository.tree(commit.id, path, recursive: recursive).sorted_entries.size
+      end
     end
-  end
 
-  def commit_exists?
-    commit.present?
-  end
+    def commit_exists?
+      commit.present?
+    end
 
-  private
+    private
 
-  def ref
-    @params[:ref] || @user_project.default_branch
-  end
+    def ref
+      @params[:ref] || @user_project.default_branch
+    end
 
-  def path
-    @params[:path] || nil
-  end
+    def path
+      @params[:path] || nil
+    end
 
-  def recursive
-    @params[:recursive]
-  end
+    def recursive
+      @params[:recursive]
+    end
 
-  def commit
-    @user_project.commit(ref)
-  end
+    def commit
+      @user_project.commit(ref)
+    end
 
-  def per_page
-    @params[:per_page] || 20
-  end
+    def per_page
+      @params[:per_page] || 20
+    end
 
-  def page_token
-    @params[:page_token]
-  end
+    def page_token
+      @params[:page_token]
+    end
 
-  def pagination_params
-    { limit: per_page, page_token: page_token }
+    def pagination_params
+      { limit: per_page, page_token: page_token }
+    end
   end
 end
-- 
GitLab


From cd2b787317202c939d8379125272b105272c3ccb Mon Sep 17 00:00:00 2001
From: Robert May <rmay@gitlab.com>
Date: Fri, 27 Aug 2021 14:49:32 +0100
Subject: [PATCH 5/9] Add definition for pagination switch param

---
 lib/api/repositories.rb | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/lib/api/repositories.rb b/lib/api/repositories.rb
index 611bbb088b9de3..3c9255e311712b 100644
--- a/lib/api/repositories.rb
+++ b/lib/api/repositories.rb
@@ -53,7 +53,11 @@ def assign_blob_vars!(limit:)
         optional :recursive, type: Boolean, default: false, desc: 'Used to get a recursive tree'
 
         use :pagination
-        optional :page_token, type: String, desc: 'Record from which to start the 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
         tree_finder = ::Repositories::TreeFinder.new(user_project, declared_params(include_missing: false))
-- 
GitLab


From 34606b3cd0d12b9cdd39d71274ba91860402a029 Mon Sep 17 00:00:00 2001
From: Robert May <rmay@gitlab.com>
Date: Fri, 27 Aug 2021 15:10:03 +0100
Subject: [PATCH 6/9] Don't use full params in cache key

---
 app/finders/repositories/tree_finder.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/app/finders/repositories/tree_finder.rb b/app/finders/repositories/tree_finder.rb
index 9631648977f361..0d948024eb6ce3 100644
--- a/app/finders/repositories/tree_finder.rb
+++ b/app/finders/repositories/tree_finder.rb
@@ -27,7 +27,7 @@ def execute(gitaly_pagination: false)
 
     def total
       # This is inefficient and we'll look at replacing this implementation
-      Gitlab::Cache.fetch_once([user_project, repository.commit, :tree_size, @params]) do
+      Gitlab::Cache.fetch_once([user_project, repository.commit, :tree_size, commit.id, path, recursive]) do
         user_project.repository.tree(commit.id, path, recursive: recursive).sorted_entries.size
       end
     end
-- 
GitLab


From d6789ef6a948362d0f6fdcec253a4c5480d208b3 Mon Sep 17 00:00:00 2001
From: Vasilii Iakliushin <viakliushin@gitlab.com>
Date: Wed, 8 Sep 2021 15:23:25 +0200
Subject: [PATCH 7/9] Fix tests and minor optimizations

---
 app/finders/repositories/tree_finder.rb       | 29 ++++++++-----------
 .../repository_tree_gitaly_pagination.yml     |  4 +--
 .../pagination/gitaly_keyset_pager_spec.rb    |  3 +-
 3 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/app/finders/repositories/tree_finder.rb b/app/finders/repositories/tree_finder.rb
index 0d948024eb6ce3..77c7c9359dc518 100644
--- a/app/finders/repositories/tree_finder.rb
+++ b/app/finders/repositories/tree_finder.rb
@@ -28,7 +28,7 @@ def execute(gitaly_pagination: false)
     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).sorted_entries.size
+        user_project.repository.tree(commit.id, path, recursive: recursive).entries.size
       end
     end
 
@@ -38,32 +38,27 @@ def commit_exists?
 
     private
 
+    def commit
+      @commit ||= user_project.commit(ref)
+    end
+
     def ref
-      @params[:ref] || @user_project.default_branch
+      params[:ref] || user_project.default_branch
     end
 
     def path
-      @params[:path] || nil
+      params[:path] || nil
     end
 
     def recursive
-      @params[:recursive]
-    end
-
-    def commit
-      @user_project.commit(ref)
-    end
-
-    def per_page
-      @params[:per_page] || 20
-    end
-
-    def page_token
-      @params[:page_token]
+      params[:recursive]
     end
 
     def pagination_params
-      { limit: per_page, page_token: page_token }
+      {
+        limit: params[:per_page] || 20,
+        page_token: params[:page_token]
+      }
     end
   end
 end
diff --git a/config/feature_flags/development/repository_tree_gitaly_pagination.yml b/config/feature_flags/development/repository_tree_gitaly_pagination.yml
index 5dcea66862b0ed..afae937b62ef19 100644
--- a/config/feature_flags/development/repository_tree_gitaly_pagination.yml
+++ b/config/feature_flags/development/repository_tree_gitaly_pagination.yml
@@ -1,8 +1,8 @@
 ---
 name: repository_tree_gitaly_pagination
 introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/67509
-rollout_issue_url: 
-milestone: '14.2'
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/340419
+milestone: '14.3'
 type: development
 group: group::source code
 default_enabled: false
diff --git a/spec/lib/gitlab/pagination/gitaly_keyset_pager_spec.rb b/spec/lib/gitlab/pagination/gitaly_keyset_pager_spec.rb
index 8a26e153385c8d..dcb8138bdded29 100644
--- a/spec/lib/gitlab/pagination/gitaly_keyset_pager_spec.rb
+++ b/spec/lib/gitlab/pagination/gitaly_keyset_pager_spec.rb
@@ -74,7 +74,7 @@
             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 @@
 
         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
 
-- 
GitLab


From 317ae10097ea07e2bf8002f171ad79302e8b0767 Mon Sep 17 00:00:00 2001
From: Robert May <rmay@gitlab.com>
Date: Tue, 14 Sep 2021 13:48:04 +0100
Subject: [PATCH 8/9] Small tweaks

---
 app/finders/repositories/tree_finder.rb | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/app/finders/repositories/tree_finder.rb b/app/finders/repositories/tree_finder.rb
index 77c7c9359dc518..716e3647a84d22 100644
--- a/app/finders/repositories/tree_finder.rb
+++ b/app/finders/repositories/tree_finder.rb
@@ -7,9 +7,9 @@ class TreeFinder < GitRefsFinder
     CommitMissingError = Class.new(StandardError)
 
     def initialize(user_project, params = {})
-      @user_project = user_project
-
       super(user_project.repository, params)
+
+      @user_project = user_project
     end
 
     def execute(gitaly_pagination: false)
@@ -47,7 +47,7 @@ def ref
     end
 
     def path
-      params[:path] || nil
+      params[:path]
     end
 
     def recursive
@@ -56,7 +56,7 @@ def recursive
 
     def pagination_params
       {
-        limit: params[:per_page] || 20,
+        limit: params[:per_page] || Kaminari.config.default_per_page,
         page_token: params[:page_token]
       }
     end
-- 
GitLab


From 079cb2cfff38304543e3715e7e37d6e17bf420a6 Mon Sep 17 00:00:00 2001
From: Robert May <rmay@gitlab.com>
Date: Tue, 14 Sep 2021 15:12:56 +0100
Subject: [PATCH 9/9] Rework execute method

---
 app/finders/repositories/tree_finder.rb | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/app/finders/repositories/tree_finder.rb b/app/finders/repositories/tree_finder.rb
index 716e3647a84d22..2ea5a8856ecd78 100644
--- a/app/finders/repositories/tree_finder.rb
+++ b/app/finders/repositories/tree_finder.rb
@@ -15,12 +15,9 @@ def initialize(user_project, params = {})
     def execute(gitaly_pagination: false)
       raise CommitMissingError unless commit_exists?
 
-      tree =
-        if gitaly_pagination
-          user_project.repository.tree(commit.id, path, recursive: recursive, pagination_params: pagination_params)
-        else
-          user_project.repository.tree(commit.id, path, recursive: recursive)
-        end
+      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
-- 
GitLab