1 files changed, 15 insertions, 10 deletions
diff --git a/app/workers/analytics/usage_trends/counter_job_worker.rb b/app/workers/analytics/usage_trends/counter_job_worker.rb
index b3a8f7dd3c2..e6de623f784 100644
--- a/app/workers/analytics/usage_trends/counter_job_worker.rb
+++ b/app/workers/analytics/usage_trends/counter_job_worker.rb
@@ -3,6 +3,8 @@
 module Analytics
   module UsageTrends
     class CounterJobWorker
+      TIMEOUT = 250.seconds
+
       extend ::Gitlab::Utils::Override
       include ApplicationWorker
 
@@ -15,24 +17,27 @@ module Analytics
 
       idempotent!
 
-      def perform(measurement_identifier, min_id, max_id, recorded_at)
+      def perform(measurement_identifier, min_id, max_id, recorded_at, partial_results = nil)
         query_scope = ::Analytics::UsageTrends::Measurement.identifier_query_mapping[measurement_identifier].call
 
-        count = if min_id.nil? || max_id.nil? # table is empty
-                  0
-                else
-                  counter(query_scope, min_id, max_id)
-                end
+        result = counter(query_scope, min_id, max_id, partial_results)
+
+        # If the batch counter timed out, schedule a job to continue counting later
+        if result[:status] == :timeout
+          return self.class.perform_async(measurement_identifier, result[:continue_from], max_id, recorded_at, result[:partial_results])
+        end
 
-        return if count == Gitlab::Database::BatchCounter::FALLBACK
+        return if result[:status] != :completed
 
-        UsageTrends::Measurement.insert_all([{ recorded_at: recorded_at, count: count, identifier: measurement_identifier }])
+        UsageTrends::Measurement.insert_all([{ recorded_at: recorded_at, count: result[:count], identifier: measurement_identifier }])
       end
 
       private
 
-      def counter(query_scope, min_id, max_id)
-        Gitlab::Database::BatchCount.batch_count(query_scope, start: min_id, finish: max_id)
+      def counter(query_scope, min_id, max_id, partial_results)
+        return { status: :completed, count: 0 } if min_id.nil? || max_id.nil? # table is empty
+
+        Gitlab::Database::BatchCount.batch_count_with_timeout(query_scope, start: min_id, finish: max_id, timeout: TIMEOUT, partial_results: partial_results)
       end
     end
   end