Optimize queries and pagination in `GroupDescendantsFinder`

1 files changed, 61 insertions, 0 deletions

diff --git a/lib/gitlab/multi_collection_paginator.rb b/lib/gitlab/multi_collection_paginator.rb
new file mode 100644
index 00000000000..eb3c9002710
--- /dev/null
+++ b/lib/gitlab/multi_collection_paginator.rb
@@ -0,0 +1,61 @@
+module Gitlab
+  class MultiCollectionPaginator
+    attr_reader :first_collection, :second_collection, :per_page
+
+    def initialize(*collections, per_page: nil)
+      raise ArgumentError.new('Only 2 collections are supported') if collections.size != 2
+
+      @per_page = per_page || Kaminari.config.default_per_page
+      @first_collection, @second_collection = collections
+    end
+
+    def paginate(page)
+      page = page.to_i
+      paginated_first_collection(page) + paginated_second_collection(page)
+    end
+
+    def total_count
+      @total_count ||= first_collection.size + second_collection.size
+    end
+
+    private
+
+    def paginated_first_collection(page)
+      @first_collection_pages ||= Hash.new do |hash, page|
+        hash[page] = first_collection.page(page).per(per_page)
+      end
+
+      @first_collection_pages[page]
+    end
+
+    def paginated_second_collection(page)
+      @second_collection_pages ||= Hash.new do |hash, page|
+        second_collection_page = page - first_collection_page_count
+
+        offset = if second_collection_page < 1 || first_collection_page_count.zero?
+                   0
+                 else
+                   per_page - first_collection_last_page_size
+                 end
+        hash[page] = second_collection.page(second_collection_page)
+                       .per(per_page - paginated_first_collection(page).size)
+                       .padding(offset)
+      end
+
+      @second_collection_pages[page]
+    end
+
+    def first_collection_page_count
+      return @first_collection_page_count if defined?(@first_collection_page_count)
+
+      first_collection_page = paginated_first_collection(0)
+      @first_collection_page_count = first_collection_page.total_pages
+    end
+
+    def first_collection_last_page_size
+      return @first_collection_last_page_size if defined?(@first_collection_last_page_size)
+
+      @first_collection_last_page_size = paginated_first_collection(first_collection_page_count).count
+    end
+  end
+end