1 files changed, 20 insertions, 11 deletions

diff --git a/lib/gitlab/database/batch_counter.rb b/lib/gitlab/database/batch_counter.rb
index 522b598cd9d..abb62140503 100644
--- a/lib/gitlab/database/batch_counter.rb
+++ b/lib/gitlab/database/batch_counter.rb
@@ -6,7 +6,6 @@ module Gitlab
       FALLBACK = -1
       MIN_REQUIRED_BATCH_SIZE = 1_250
       DEFAULT_SUM_BATCH_SIZE = 1_000
-      DEFAULT_AVERAGE_BATCH_SIZE = 1_000
       MAX_ALLOWED_LOOPS = 10_000
       SLEEP_TIME_IN_SECONDS = 0.01 # 10 msec sleep
       ALLOWED_MODES = [:itself, :distinct].freeze
@@ -27,12 +26,19 @@ module Gitlab
       def unwanted_configuration?(finish, batch_size, start)
         (@operation == :count && batch_size <= MIN_REQUIRED_BATCH_SIZE) ||
           (@operation == :sum && batch_size < DEFAULT_SUM_BATCH_SIZE) ||
-          (@operation == :average && batch_size < DEFAULT_AVERAGE_BATCH_SIZE) ||
           (finish - start) / batch_size >= MAX_ALLOWED_LOOPS ||
           start >= finish
       end
 
       def count(batch_size: nil, mode: :itself, start: nil, finish: nil)
+        result = count_with_timeout(batch_size: batch_size, mode: mode, start: start, finish: finish, timeout: nil)
+
+        return FALLBACK if result[:status] != :completed
+
+        result[:count]
+      end
+
+      def count_with_timeout(batch_size: nil, mode: :itself, start: nil, finish: nil, timeout: nil, partial_results: nil)
         raise 'BatchCount can not be run inside a transaction' if transaction_open?
 
         check_mode!(mode)
@@ -44,12 +50,20 @@ module Gitlab
         finish = actual_finish(finish)
 
         raise "Batch counting expects positive values only for #{@column}" if start < 0 || finish < 0
-        return FALLBACK if unwanted_configuration?(finish, batch_size, start)
+        return { status: :bad_config } if unwanted_configuration?(finish, batch_size, start)
 
-        results = nil
+        results = partial_results
         batch_start = start
 
+        start_time = ::Gitlab::Metrics::System.monotonic_time.seconds
+
         while batch_start < finish
+
+          # Timeout elapsed, return partial result so the caller can continue later
+          if timeout && ::Gitlab::Metrics::System.monotonic_time.seconds - start_time > timeout
+            return { status: :timeout, partial_results: results, continue_from: batch_start }
+          end
+
           begin
             batch_end = [batch_start + batch_size, finish].min
             batch_relation = build_relation_batch(batch_start, batch_end, mode)
@@ -62,14 +76,14 @@ module Gitlab
               batch_size /= 2
             else
               log_canceled_batch_fetch(batch_start, mode, batch_relation.to_sql, error)
-              return FALLBACK
+              return { status: :cancelled }
             end
           end
 
           sleep(SLEEP_TIME_IN_SECONDS)
         end
 
-        results
+        { status: :completed, count: results }
       end
 
       def transaction_open?
@@ -94,7 +108,6 @@ module Gitlab
 
       def batch_size_for_mode_and_operation(mode, operation)
         return DEFAULT_SUM_BATCH_SIZE if operation == :sum
-        return DEFAULT_AVERAGE_BATCH_SIZE if operation == :average
 
         mode == :distinct ? DEFAULT_DISTINCT_BATCH_SIZE : DEFAULT_BATCH_SIZE
       end
@@ -132,10 +145,6 @@ module Gitlab
             message: "Query has been canceled with message: #{error.message}"
           )
       end
-
-      def not_group_by_query?
-        !@relation.is_a?(ActiveRecord::Relation) || @relation.group_values.blank?
-      end
     end
   end
 end