Add logging of failed queries in BatchCount

In order to have better insight into not performant queries we need to
log them.

changelogs/unreleased/mwaw-230438-log-usage-data-failed-queries.yml

+---
+title: Log failed BatchCount queries
+merge_request: 41552
+author:
+type: added

lib/gitlab/database/batch_count.rb

@@ -86,14 +86,16 @@ module Gitlab
         batch_start = start
 
         while batch_start <= finish
+          batch_relation = build_relation_batch(batch_start, batch_start + batch_size, mode)
           begin
-            results = merge_results(results, batch_fetch(batch_start, batch_start + batch_size, mode))
+            results = merge_results(results, batch_relation.send(@operation, *@operation_args)) # rubocop:disable GitlabSecurity/PublicSend
             batch_start += batch_size
-          rescue ActiveRecord::QueryCanceled
+          rescue ActiveRecord::QueryCanceled => error
             # retry with a safe batch size & warmer cache
             if batch_size >= 2 * MIN_REQUIRED_BATCH_SIZE
               batch_size /= 2
             else
+              log_canceled_batch_fetch(batch_start, mode, batch_relation.to_sql, error)
               return FALLBACK
             end
           end
@@ -113,13 +115,12 @@ module Gitlab
         end
       end
 
-      def batch_fetch(start, finish, mode)
-        # rubocop:disable GitlabSecurity/PublicSend
-        @relation.select(@column).public_send(mode).where(between_condition(start, finish)).send(@operation, *@operation_args)
-      end
-
       private
 
+      def build_relation_batch(start, finish, mode)
+        @relation.select(@column).public_send(mode).where(between_condition(start, finish)) # rubocop:disable GitlabSecurity/PublicSend
+      end
+
       def batch_size_for_mode_and_operation(mode, operation)
         return DEFAULT_SUM_BATCH_SIZE if operation == :sum
 
@@ -145,6 +146,20 @@ module Gitlab
         raise 'Use distinct count for optimized distinct counting' if @relation.limit(1).distinct_value.present? && mode != :distinct
         raise 'Use distinct count only with non id fields' if @column == :id && mode == :distinct
       end
+
+      def log_canceled_batch_fetch(batch_start, mode, query, error)
+        Gitlab::AppJsonLogger
+          .error(
+            event: 'batch_count',
+            relation: @relation.table_name,
+            operation: @operation,
+            operation_args: @operation_args,
+            start: batch_start,
+            mode: mode,
+            query: query,
+            message: "Query has been canceled with message: #{error.message}"
+          )
+      end
     end
   end
 end

spec/lib/gitlab/database/batch_count_spec.rb

@@ -4,7 +4,7 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::BatchCount do
   let_it_be(:fallback) { ::Gitlab::Database::BatchCounter::FALLBACK }
-  let_it_be(:small_batch_size) { ::Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE - 1 }
+  let_it_be(:small_batch_size) { calculate_batch_size(::Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE) }
   let(:model) { Issue }
   let(:column) { :author_id }
 
@@ -22,6 +22,12 @@ RSpec.describe Gitlab::Database::BatchCount do
     allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(in_transaction)
   end
 
+  def calculate_batch_size(batch_size)
+    zero_offset_modifier = -1
+
+    batch_size + zero_offset_modifier
+  end
+
   shared_examples 'disallowed configurations' do |method|
     it 'returns fallback if start is bigger than finish' do
       expect(described_class.public_send(method, *args, start: 1, finish: 0)).to eq(fallback)
@@ -45,6 +51,46 @@ RSpec.describe Gitlab::Database::BatchCount do
     end
   end
 
+  shared_examples 'when batch fetch query is canceled' do
+    let(:batch_size) { 22_000 }
+    let(:relation) { instance_double(ActiveRecord::Relation) }
+
+    it 'reduces batch size by half and retry fetch' do
+      too_big_batch_relation_mock = instance_double(ActiveRecord::Relation)
+      allow(model).to receive_message_chain(:select, public_send: relation)
+      allow(relation).to receive(:where).with("id" => 0..calculate_batch_size(batch_size)).and_return(too_big_batch_relation_mock)
+      allow(too_big_batch_relation_mock).to receive(:send).and_raise(ActiveRecord::QueryCanceled)
+
+      expect(relation).to receive(:where).with("id" => 0..calculate_batch_size(batch_size / 2)).and_return(double(send: 1))
+
+      subject.call(model, column, batch_size: batch_size, start: 0)
+    end
+
+    context 'when all retries fail' do
+      let(:batch_count_query) { 'SELECT COUNT(id) FROM relation WHERE id BETWEEN 0 and 1' }
+
+      before do
+        allow(model).to receive_message_chain(:select, :public_send, where: relation)
+        allow(relation).to receive(:send).and_raise(ActiveRecord::QueryCanceled.new('query timed out'))
+        allow(relation).to receive(:to_sql).and_return(batch_count_query)
+      end
+
+      it 'logs failing query' do
+        expect(Gitlab::AppJsonLogger).to receive(:error).with(
+          event: 'batch_count',
+          relation: model.table_name,
+          operation: operation,
+          operation_args: operation_args,
+          start: 0,
+          mode: mode,
+          query: batch_count_query,
+          message: 'Query has been canceled with message: query timed out'
+        )
+        expect(subject.call(model, column, batch_size: batch_size, start: 0)).to eq(-1)
+      end
+    end
+  end
+
   describe '#batch_count' do
     it 'counts table' do
       expect(described_class.batch_count(model)).to eq(5)
@@ -86,10 +132,11 @@ RSpec.describe Gitlab::Database::BatchCount do
 
     it "defaults the batch size to #{Gitlab::Database::BatchCounter::DEFAULT_BATCH_SIZE}" do
       min_id = model.minimum(:id)
+      relation = instance_double(ActiveRecord::Relation)
+      allow(model).to receive_message_chain(:select, public_send: relation)
+      batch_end_id = min_id + calculate_batch_size(Gitlab::Database::BatchCounter::DEFAULT_BATCH_SIZE)
 
-      expect_next_instance_of(Gitlab::Database::BatchCounter) do |batch_counter|
-        expect(batch_counter).to receive(:batch_fetch).with(min_id, Gitlab::Database::BatchCounter::DEFAULT_BATCH_SIZE + min_id, :itself).once.and_call_original
-      end
+      expect(relation).to receive(:where).with("id" => min_id..batch_end_id).and_return(double(send: 1))
 
       described_class.batch_count(model)
     end
@@ -98,6 +145,15 @@ RSpec.describe Gitlab::Database::BatchCount do
       subject { described_class.batch_count(model) }
     end
 
+    it_behaves_like 'when batch fetch query is canceled' do
+      let(:mode) { :itself }
+      let(:operation) { :count }
+      let(:operation_args) { nil }
+      let(:column) { nil }
+
+      subject { described_class.method(:batch_count) }
+    end
+
     context 'disallowed_configurations' do
       include_examples 'disallowed configurations', :batch_count do
         let(:args) { [Issue] }
@@ -169,10 +225,11 @@ RSpec.describe Gitlab::Database::BatchCount do
 
     it "defaults the batch size to #{Gitlab::Database::BatchCounter::DEFAULT_DISTINCT_BATCH_SIZE}" do
       min_id = model.minimum(:id)
+      relation = instance_double(ActiveRecord::Relation)
+      allow(model).to receive_message_chain(:select, public_send: relation)
+      batch_end_id = min_id + calculate_batch_size(Gitlab::Database::BatchCounter::DEFAULT_DISTINCT_BATCH_SIZE)
 
-      expect_next_instance_of(Gitlab::Database::BatchCounter) do |batch_counter|
-        expect(batch_counter).to receive(:batch_fetch).with(min_id, Gitlab::Database::BatchCounter::DEFAULT_DISTINCT_BATCH_SIZE + min_id, :distinct).once.and_call_original
-      end
+      expect(relation).to receive(:where).with("id" => min_id..batch_end_id).and_return(double(send: 1))
 
       described_class.batch_distinct_count(model)
     end
@@ -211,6 +268,15 @@ RSpec.describe Gitlab::Database::BatchCount do
         end
       end
     end
+
+    it_behaves_like 'when batch fetch query is canceled' do
+      let(:mode) { :distinct }
+      let(:operation) { :count }
+      let(:operation_args) { nil }
+      let(:column) { nil }
+
+      subject { described_class.method(:batch_distinct_count) }
+    end
   end
 
   describe '#batch_sum' do
@@ -245,10 +311,11 @@ RSpec.describe Gitlab::Database::BatchCount do
 
     it "defaults the batch size to #{Gitlab::Database::BatchCounter::DEFAULT_SUM_BATCH_SIZE}" do
       min_id = model.minimum(:id)
+      relation = instance_double(ActiveRecord::Relation)
+      allow(model).to receive_message_chain(:select, public_send: relation)
+      batch_end_id = min_id + calculate_batch_size(Gitlab::Database::BatchCounter::DEFAULT_SUM_BATCH_SIZE)
 
-      expect_next_instance_of(Gitlab::Database::BatchCounter) do |batch_counter|
-        expect(batch_counter).to receive(:batch_fetch).with(min_id, Gitlab::Database::BatchCounter::DEFAULT_SUM_BATCH_SIZE + min_id, :itself).once.and_call_original
-      end
+      expect(relation).to receive(:where).with("id" => min_id..batch_end_id).and_return(double(send: 1))
 
       described_class.batch_sum(model, column)
     end
@@ -262,5 +329,13 @@ RSpec.describe Gitlab::Database::BatchCount do
       let(:default_batch_size) { Gitlab::Database::BatchCounter::DEFAULT_SUM_BATCH_SIZE }
       let(:small_batch_size) { Gitlab::Database::BatchCounter::DEFAULT_SUM_BATCH_SIZE - 1 }
     end
+
+    it_behaves_like 'when batch fetch query is canceled' do
+      let(:mode) { :itself }
+      let(:operation) { :sum }
+      let(:operation_args) { [column] }
+
+      subject { described_class.method(:batch_sum) }
+    end
   end
 end