Split and retry background migration batches on QueryCanceled

We currently split and retry failed background migration
batches that have failed because of a timeout exception.
The list of timeout exceptions should also include
ActiveRecord:QueryCanceled since in can be raise as
PG::QueryCanceled: ERROR:  canceling statement due to statement
timeout

config/feature_flags/development/split_background_migration_on_query_canceled.yml

+---
+name: split_background_migration_on_query_canceled
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/101825
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/378846
+milestone: '15.6'
+type: development
+group: group::database
+default_enabled: false

lib/gitlab/database/background_migration/batched_job.rb

@@ -112,7 +112,7 @@ def time_efficiency
         end
 
         def can_split?(exception)
-          attempts >= MAX_ATTEMPTS && TIMEOUT_EXCEPTIONS.include?(exception&.class) && batch_size > sub_batch_size && batch_size > 1
+          attempts >= MAX_ATTEMPTS && timeout_exception?(exception&.class) && batch_size > sub_batch_size && batch_size > 1
         end
 
         def split_and_retry!
@@ -161,6 +161,15 @@ def split_and_retry!
             end
           end
         end
+
+        private
+
+        def timeout_exception?(exception_class)
+          return false unless exception_class
+
+          TIMEOUT_EXCEPTIONS.include?(exception_class) ||
+            (Feature.enabled?(:split_background_migration_on_query_canceled) && exception_class == ActiveRecord::QueryCanceled)
+        end
       end
     end
   end

spec/lib/gitlab/database/background_migration/batched_job_spec.rb

@@ -272,7 +272,21 @@
       context 'when is a timeout exception' do
         let(:exception) { ActiveRecord::StatementTimeout.new }
 
-        it { expect(subject).to be_truthy  }
+        it { expect(subject).to be_truthy }
+      end
+
+      context 'when is a QueryCanceled exception' do
+        let(:exception) { ActiveRecord::QueryCanceled.new }
+
+        it { expect(subject).to be_truthy }
+
+        context 'when split_background_migration_on_query_canceled feature flag is disabled' do
+          before do
+            stub_feature_flags(split_background_migration_on_query_canceled: false)
+          end
+
+          it { expect(subject).to be_falsey }
+        end
       end
 
       context 'when is not a timeout exception' do