[API] Omit X-Total{,-Pages} when count > 10k

Signed-off-by: Rémy Coutable <remy@rymai.me>

5 files changed, 150 insertions, 18 deletions

diff --git a/changelogs/unreleased/52674-api-v4-projects-project_id-jobs-endpoint-hits-statement-timeout.yml b/changelogs/unreleased/52674-api-v4-projects-project_id-jobs-endpoint-hits-statement-timeout.yml
new file mode 100644
index 00000000000..f79078c1fd9
--- /dev/null
+++ b/changelogs/unreleased/52674-api-v4-projects-project_id-jobs-endpoint-hits-statement-timeout.yml
@@ -0,0 +1,5 @@
+---
+title: "[API] Omit `X-Total` and `X-Total-Pages` headers when items count is more than 10,000"
+merge_request: 23931
+author:
+type: performance
diff --git a/config/initializers/kaminari_active_record_relation_methods_with_limit.rb b/config/initializers/kaminari_active_record_relation_methods_with_limit.rb
new file mode 100644
index 00000000000..cc20b83b234
--- /dev/null
+++ b/config/initializers/kaminari_active_record_relation_methods_with_limit.rb
@@ -0,0 +1,41 @@
+module Kaminari
+  # Active Record specific page scope methods implementations
+  module ActiveRecordRelationMethodsWithLimit
+    MAX_COUNT_LIMIT = 10_000
+
+    # This is a modified version of
+    # https://github.com/kaminari/kaminari/blob/c5186f5d9b7f23299d115408e62047447fd3189d/kaminari-activerecord/lib/kaminari/activerecord/active_record_relation_methods.rb#L17-L41
+    # that limit the COUNT query to 10,000 to avoid query timeouts.
+    # rubocop: disable Gitlab/ModuleWithInstanceVariables
+    def total_count_with_limit(column_name = :all, _options = nil) #:nodoc:
+      return @total_count if defined?(@total_count) && @total_count
+
+      # There are some cases that total count can be deduced from loaded records
+      if loaded?
+        # Total count has to be 0 if loaded records are 0
+        return @total_count = 0 if (current_page == 1) && @records.empty?
+        # Total count is calculable at the last page
+        return @total_count = (current_page - 1) * limit_value + @records.length if @records.any? && (@records.length < limit_value)
+      end
+
+      # #count overrides the #select which could include generated columns referenced in #order, so skip #order here, where it's irrelevant to the result anyway
+      c = except(:offset, :limit, :order)
+      # Remove includes only if they are irrelevant
+      c = c.except(:includes) unless references_eager_loaded_tables?
+      # .group returns an OrderedHash that responds to #count
+      # The following line was modified from `c = c.count(:all)`
+      c = c.limit(MAX_COUNT_LIMIT + 1).count(column_name)
+      @total_count =
+        if c.is_a?(Hash) || c.is_a?(ActiveSupport::OrderedHash)
+          c.count
+        elsif c.respond_to? :count
+          c.count(column_name)
+        else
+          c
+        end
+    end
+    # rubocop: enable Gitlab/ModuleWithInstanceVariables
+
+    Kaminari::ActiveRecordRelationMethods.prepend(self)
+  end
+end
diff --git a/doc/api/README.md b/doc/api/README.md
index 6c5bb1c0940..7b83b0fed26 100644
--- a/doc/api/README.md
+++ b/doc/api/README.md
@@ -438,6 +438,14 @@ Additional pagination headers are also sent back.
 | `X-Next-Page`   | The index of the next page |
 | `X-Prev-Page`   | The index of the previous page |
 
+CAUTION: **Caution:**
+For performance reasons since
+[GitLab 11.8][https://gitlab.com/gitlab-org/gitlab-ce/merge_requests/23931]
+and **behind the `api_kaminari_count_with_limit`
+[feature flag](../development/feature_flags.md)**, if the number of resources is
+more than 10,000, the `X-Total` and `X-Total-Pages` headers as well as the
+`rel="last"` `Link` are not present in the response headers.
+
 ## Namespaced path encoding
 
 If using namespaced API calls, make sure that the `NAMESPACE/PROJECT_NAME` is
diff --git a/lib/api/helpers/pagination.rb b/lib/api/helpers/pagination.rb
index d311cbb5f7e..cfcce6b66ad 100644
--- a/lib/api/helpers/pagination.rb
+++ b/lib/api/helpers/pagination.rb
@@ -178,15 +178,26 @@ module API
         end
 
         def paginate(relation)
-          relation = add_default_order(relation)
-
-          relation.page(params[:page]).per(params[:per_page]).tap do |data|
+          paginate_with_limit_optimization(add_default_order(relation)).tap do |data|
             add_pagination_headers(data)
           end
         end
 
         private
 
+        def paginate_with_limit_optimization(relation)
+          pagination_data = relation.page(params[:page]).per(params[:per_page])
+          return pagination_data unless pagination_data.is_a?(ActiveRecord::Relation)
+          return pagination_data unless Feature.enabled?(:api_kaminari_count_with_limit)
+
+          limited_total_count = pagination_data.total_count_with_limit
+          if limited_total_count > Kaminari::ActiveRecordRelationMethods::MAX_COUNT_LIMIT
+            pagination_data.without_count
+          else
+            pagination_data
+          end
+        end
+
         # rubocop: disable CodeReuse/ActiveRecord
         def add_default_order(relation)
           if relation.is_a?(ActiveRecord::Relation) && relation.order_values.empty?
diff --git a/spec/lib/api/helpers/pagination_spec.rb b/spec/lib/api/helpers/pagination_spec.rb
index 0a7682d906b..2890aa4ae38 100644
--- a/spec/lib/api/helpers/pagination_spec.rb
+++ b/spec/lib/api/helpers/pagination_spec.rb
@@ -237,26 +237,89 @@ describe API::Helpers::Pagination do
             .and_return({ page: 1, per_page: 2 })
         end
 
-        it 'returns appropriate amount of resources' do
-          expect(subject.paginate(resource).count).to eq 2
+        shared_examples 'response with pagination headers' do
+          it 'adds appropriate headers' do
+            expect_header('X-Total', '3')
+            expect_header('X-Total-Pages', '2')
+            expect_header('X-Per-Page', '2')
+            expect_header('X-Page', '1')
+            expect_header('X-Next-Page', '2')
+            expect_header('X-Prev-Page', '')
+
+            expect_header('Link', anything) do |_key, val|
+              expect(val).to include('rel="first"')
+              expect(val).to include('rel="last"')
+              expect(val).to include('rel="next"')
+              expect(val).not_to include('rel="prev"')
+            end
+
+            subject.paginate(resource)
+          end
         end
 
-        it 'adds appropriate headers' do
-          expect_header('X-Total', '3')
-          expect_header('X-Total-Pages', '2')
-          expect_header('X-Per-Page', '2')
-          expect_header('X-Page', '1')
-          expect_header('X-Next-Page', '2')
-          expect_header('X-Prev-Page', '')
+        shared_examples 'paginated response' do
+          it 'returns appropriate amount of resources' do
+            expect(subject.paginate(resource).count).to eq 2
+          end
 
-          expect_header('Link', anything) do |_key, val|
-            expect(val).to include('rel="first"')
-            expect(val).to include('rel="last"')
-            expect(val).to include('rel="next"')
-            expect(val).not_to include('rel="prev"')
+          it 'executes only one SELECT COUNT query' do
+            expect { subject.paginate(resource) }.to make_queries_matching(/SELECT COUNT/, 1)
           end
+        end
 
-          subject.paginate(resource)
+        context 'when the api_kaminari_count_with_limit feature flag is unset' do
+          it_behaves_like 'paginated response'
+          it_behaves_like 'response with pagination headers'
+        end
+
+        context 'when the api_kaminari_count_with_limit feature flag is disabled' do
+          before do
+            stub_feature_flags(api_kaminari_count_with_limit: false)
+          end
+
+          it_behaves_like 'paginated response'
+          it_behaves_like 'response with pagination headers'
+        end
+
+        context 'when the api_kaminari_count_with_limit feature flag is enabled' do
+          before do
+            stub_feature_flags(api_kaminari_count_with_limit: true)
+          end
+
+          context 'when resources count is less than MAX_COUNT_LIMIT' do
+            before do
+              stub_const("::Kaminari::ActiveRecordRelationMethods::MAX_COUNT_LIMIT", 4)
+            end
+
+            it_behaves_like 'paginated response'
+            it_behaves_like 'response with pagination headers'
+          end
+
+          context 'when resources count is more than MAX_COUNT_LIMIT' do
+            before do
+              stub_const("::Kaminari::ActiveRecordRelationMethods::MAX_COUNT_LIMIT", 2)
+            end
+
+            it_behaves_like 'paginated response'
+
+            it 'does not return the X-Total and X-Total-Pages headers' do
+              expect_no_header('X-Total')
+              expect_no_header('X-Total-Pages')
+              expect_header('X-Per-Page', '2')
+              expect_header('X-Page', '1')
+              expect_header('X-Next-Page', '2')
+              expect_header('X-Prev-Page', '')
+
+              expect_header('Link', anything) do |_key, val|
+                expect(val).to include('rel="first"')
+                expect(val).not_to include('rel="last"')
+                expect(val).to include('rel="next"')
+                expect(val).not_to include('rel="prev"')
+              end
+
+              subject.paginate(resource)
+            end
+          end
         end
       end
 
@@ -348,6 +411,10 @@ describe API::Helpers::Pagination do
     expect(subject).to receive(:header).with(*args, &block)
   end
 
+  def expect_no_header(*args, &block)
+    expect(subject).not_to receive(:header).with(*args)
+  end
+
   def expect_message(method)
     expect(subject).to receive(method)
       .at_least(:once).and_return(value)