Incorporate Gitaly's WikiService.WikiGetAllPages RPCgitaly-wiki-get-all-pages

3 files changed, 62 insertions, 18 deletions

diff --git a/lib/gitlab/git/wiki.rb b/lib/gitlab/git/wiki.rb
index 5d36f16abd4..022d1f249a9 100644
--- a/lib/gitlab/git/wiki.rb
+++ b/lib/gitlab/git/wiki.rb
@@ -59,7 +59,13 @@ module Gitlab
       end
 
       def pages
-        gollum_wiki.pages.map { |gollum_page| new_page(gollum_page) }
+        @repository.gitaly_migrate(:wiki_get_all_pages) do |is_enabled|
+          if is_enabled
+            gitaly_get_all_pages
+          else
+            gollum_get_all_pages
+          end
+        end
       end
 
       def page(title:, version: nil, dir: nil)
@@ -179,6 +185,10 @@ module Gitlab
         Gitlab::Git::WikiFile.new(gollum_file)
       end
 
+      def gollum_get_all_pages
+        gollum_wiki.pages.map { |gollum_page| new_page(gollum_page) }
+      end
+
       def gitaly_write_page(name, format, content, commit_details)
         gitaly_wiki_client.write_page(name, format, content, commit_details)
       end
@@ -204,6 +214,12 @@ module Gitlab
 
         Gitlab::Git::WikiFile.new(wiki_file)
       end
+
+      def gitaly_get_all_pages
+        gitaly_wiki_client.get_all_pages.map do |wiki_page, version|
+          Gitlab::Git::WikiPage.new(wiki_page, version)
+        end
+      end
     end
   end
 end
diff --git a/lib/gitlab/gitaly_client/wiki_page.rb b/lib/gitlab/gitaly_client/wiki_page.rb
index 8226278d5f6..98d96fe6211 100644
--- a/lib/gitlab/gitaly_client/wiki_page.rb
+++ b/lib/gitlab/gitaly_client/wiki_page.rb
@@ -11,6 +11,10 @@ module Gitlab
         FIELDS.each do |field|
           instance_variable_set("@#{field}", params[field])
         end
+
+        # All gRPC strings in a response are frozen, so we get an unfrozen
+        # version here so appending to `raw_data` doesn't blow up.
+        @raw_data = @raw_data.dup
       end
 
       def historical?
diff --git a/lib/gitlab/gitaly_client/wiki_service.rb b/lib/gitlab/gitaly_client/wiki_service.rb
index 1a668338f57..8f05f40365e 100644
--- a/lib/gitlab/gitaly_client/wiki_service.rb
+++ b/lib/gitlab/gitaly_client/wiki_service.rb
@@ -81,28 +81,23 @@ module Gitlab
         )
 
         response = GitalyClient.call(@repository.storage, :wiki_service, :wiki_find_page, request)
-        wiki_page = version = nil
 
-        response.each do |message|
-          page = message.page
-          next unless page
+        wiki_page_from_iterator(response)
+      end
 
-          if wiki_page
-            wiki_page.raw_data << page.raw_data
-          else
-            wiki_page = GitalyClient::WikiPage.new(page.to_h)
-            # All gRPC strings in a response are frozen, so we get
-            # an unfrozen version here so appending in the else clause below doesn't blow up.
-            wiki_page.raw_data = wiki_page.raw_data.dup
+      def get_all_pages
+        request = Gitaly::WikiGetAllPagesRequest.new(repository: @gitaly_repo)
+        response = GitalyClient.call(@repository.storage, :wiki_service, :wiki_get_all_pages, request)
+        pages = []
 
-            version = Gitlab::Git::WikiPageVersion.new(
-              Gitlab::Git::Commit.decorate(@repository, page.version.commit),
-              page.version.format
-            )
-          end
+        loop do
+          page, version = wiki_page_from_iterator(response) { |message| message.end_of_page }
+
+          break unless page && version
+          pages << [page, version]
         end
 
-        [wiki_page, version]
+        pages
       end
 
       def find_file(name, revision)
@@ -133,6 +128,35 @@ module Gitlab
 
       private
 
+      # If a block is given and the yielded value is true, iteration will be
+      # stopped early at that point; else the iterator is consumed entirely.
+      # The iterator is traversed with `next` to allow resuming the iteration.
+      def wiki_page_from_iterator(iterator)
+        wiki_page = version = nil
+
+        while message = iterator.next
+          break if block_given? && yield(message)
+
+          page = message.page
+          next unless page
+
+          if wiki_page
+            wiki_page.raw_data << page.raw_data
+          else
+            wiki_page = GitalyClient::WikiPage.new(page.to_h)
+
+            version = Gitlab::Git::WikiPageVersion.new(
+              Gitlab::Git::Commit.decorate(@repository, page.version.commit),
+              page.version.format
+            )
+          end
+        end
+
+        [wiki_page, version]
+      rescue StopIteration
+        [wiki_page, version]
+      end
+
       def gitaly_commit_details(commit_details)
         Gitaly::WikiCommitDetails.new(
           name: GitalyClient.encode(commit_details.name),