Add retry logic for post import action

  - migration script to add `source` to import failure table. Initially
    table `import_failures` is designed to save import failures for
    `build relation` action. Now we also want to save exceptions from
    `merge_requests.set_latest_merge_request_diff_ids!` action. So add a
    new column to record the `source` of the failure.
  - add retry logic for set_latest_merge_request_diff_ids
  - change with_try and log_import_failure signature to keyword args

db/migrate/20200124053531_add_source_to_import_failures.rb

+# frozen_string_literal: true
+
+class AddSourceToImportFailures < ActiveRecord::Migration[5.2]
+  DOWNTIME = false
+
+  def change
+    add_column :import_failures, :source, :string, limit: 128
+  end
+end

db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema.define(version: 2020_01_23_155929) do
+ActiveRecord::Schema.define(version: 2020_01_24_053531) do
 
   # These are extensions that must be enabled in order to support this database
   enable_extension "pg_trgm"
@@ -2046,6 +2046,7 @@ ActiveRecord::Schema.define(version: 2020_01_23_155929) do
     t.string "exception_message", limit: 255
     t.integer "retry_count"
     t.integer "group_id"
+    t.string "source", limit: 128
     t.index ["correlation_id_value"], name: "index_import_failures_on_correlation_id_value"
     t.index ["group_id"], name: "index_import_failures_on_group_id_not_null", where: "(group_id IS NOT NULL)"
     t.index ["project_id"], name: "index_import_failures_on_project_id_not_null", where: "(project_id IS NOT NULL)"

lib/gitlab/import_export/import_failure_service.rb

@@ -12,9 +12,14 @@ module Gitlab
         @association = importable.association(:import_failures)
       end
 
-      def with_retry(relation_key, relation_index)
+      def with_retry(action:, relation_key: nil, relation_index: nil)
         on_retry = -> (exception, retry_count, *_args) do
-          log_import_failure(relation_key, relation_index, exception, retry_count)
+          log_import_failure(
+            source: action,
+            relation_key: relation_key,
+            relation_index: relation_index,
+            exception: exception,
+            retry_count: retry_count)
         end
 
         Retriable.with_context(:relation_import, on_retry: on_retry) do
@@ -22,8 +27,9 @@ module Gitlab
         end
       end
 
-      def log_import_failure(relation_key, relation_index, exception, retry_count = 0)
+      def log_import_failure(source:, relation_key: nil, relation_index: nil, exception:, retry_count: 0)
         extra = {
+          source: source,
           relation_key: relation_key,
           relation_index: relation_index,
           retry_count: retry_count

lib/gitlab/import_export/project_tree_restorer.rb

@@ -21,7 +21,9 @@ module Gitlab
         RelationRenameService.rename(@tree_hash)
 
         if relation_tree_restorer.restore
-          @project.merge_requests.set_latest_merge_request_diff_ids!
+          import_failure_service.with_retry(action: 'set_latest_merge_request_diff_ids!') do
+            @project.merge_requests.set_latest_merge_request_diff_ids!
+          end
 
           true
         else
@@ -72,6 +74,10 @@ module Gitlab
       def reader
         @reader ||= Gitlab::ImportExport::Reader.new(shared: @shared)
       end
+
+      def import_failure_service
+        @import_failure_service ||= ImportFailureService.new(@project)
+      end
     end
   end
 end

lib/gitlab/import_export/relation_tree_restorer.rb

@@ -73,13 +73,17 @@ module Gitlab
 
         relation_object.assign_attributes(importable_class_sym => @importable)
 
-        import_failure_service.with_retry(relation_key, relation_index) do
+        import_failure_service.with_retry(action: 'relation_object.save!', relation_key: relation_key, relation_index: relation_index) do
           relation_object.save!
         end
 
         save_id_mapping(relation_key, data_hash, relation_object)
       rescue => e
-        import_failure_service.log_import_failure(relation_key, relation_index, e)
+        import_failure_service.log_import_failure(
+          source: 'process_relation_item!',
+          relation_key: relation_key,
+          relation_index: relation_index,
+          exception: e)
       end
 
       def import_failure_service

spec/lib/gitlab/import_export/import_failure_service_spec.rb

@@ -6,6 +6,7 @@ describe Gitlab::ImportExport::ImportFailureService do
   let(:importable) { create(:project, :builds_enabled, :issues_disabled, name: 'project', path: 'project') }
   let(:label) { create(:label) }
   let(:subject) { described_class.new(importable) }
+  let(:action) { "save_relation" }
   let(:relation_key) { "labels" }
   let(:relation_index) { 0 }
 
@@ -15,7 +16,12 @@ describe Gitlab::ImportExport::ImportFailureService do
     let(:correlation_id) { 'my-correlation-id' }
     let(:retry_count) { 2 }
     let(:log_import_failure) do
-      subject.log_import_failure(relation_key, relation_index, exception, retry_count)
+      subject.log_import_failure(
+        source: action,
+        relation_key: relation_key,
+        relation_index: relation_index,
+        exception: exception,
+        retry_count: retry_count)
     end
 
     before do
@@ -44,7 +50,7 @@ describe Gitlab::ImportExport::ImportFailureService do
 
   describe '#with_retry' do
     let(:perform_retry) do
-      subject.with_retry(relation_key, relation_index) do
+      subject.with_retry(action: action, relation_key: relation_key, relation_index: relation_index) do
         label.save!
       end
     end
@@ -60,7 +66,12 @@ describe Gitlab::ImportExport::ImportFailureService do
           end
 
           it 'retries and logs import failure once with correct params' do
-            expect(subject).to receive(:log_import_failure).with(relation_key, relation_index, instance_of(exception), 1).once
+            expect(subject).to receive(:log_import_failure).with(
+              source: action,
+              relation_key: relation_key,
+              relation_index: relation_index,
+              exception: instance_of(exception),
+              retry_count: 1).once
 
             perform_retry
           end
@@ -85,7 +96,11 @@ describe Gitlab::ImportExport::ImportFailureService do
             maximum_retry_count.times do |index|
               retry_count = index + 1
 
-              expect(subject).to receive(:log_import_failure).with(relation_key, relation_index, instance_of(exception), retry_count)
+              expect(subject).to receive(:log_import_failure).with(
+                source: action, relation_key: relation_key,
+                relation_index: relation_index,
+                exception: instance_of(exception),
+                retry_count: retry_count)
             end
 
             expect { perform_retry }.to raise_exception(exception)

spec/lib/gitlab/import_export/project_tree_restorer_spec.rb

@@ -498,6 +498,58 @@ describe Gitlab::ImportExport::ProjectTreeRestorer do
       end
     end
 
+    context 'when post import action throw non-retriable exception' do
+      let(:exception) { StandardError.new('post_import_error') }
+
+      before do
+        setup_import_export_config('light')
+        expect(project)
+          .to receive(:merge_requests)
+          .and_raise(exception)
+      end
+
+      it 'report post import error' do
+        expect(restored_project_json).to eq(false)
+        expect(shared.errors).to include('post_import_error')
+      end
+    end
+
+    context 'when post import action throw retriable exception one time' do
+      let(:exception) { GRPC::DeadlineExceeded.new }
+
+      before do
+        setup_import_export_config('light')
+        expect(project)
+          .to receive(:merge_requests)
+          .and_raise(exception)
+        expect(project)
+          .to receive(:merge_requests)
+          .and_call_original
+        expect(restored_project_json).to eq(true)
+      end
+
+      it_behaves_like 'restores project successfully',
+                      issues: 1,
+                      labels: 2,
+                      label_with_priorities: 'A project label',
+                      milestones: 1,
+                      first_issue_labels: 1,
+                      services: 1,
+                      import_failures: 1
+
+      it 'records the failures in the database' do
+        import_failure = ImportFailure.last
+
+        expect(import_failure.project_id).to eq(project.id)
+        expect(import_failure.relation_key).to be_nil
+        expect(import_failure.relation_index).to be_nil
+        expect(import_failure.exception_class).to eq('GRPC::DeadlineExceeded')
+        expect(import_failure.exception_message).to be_present
+        expect(import_failure.correlation_id_value).not_to be_empty
+        expect(import_failure.created_at).to be_present
+      end
+    end
+
     context 'when the project has overridden params in import data' do
       before do
         setup_import_export_config('light')

spec/support/shared_examples/lib/gitlab/import_export/import_failure_service_shared_examples.rb

@@ -3,6 +3,7 @@
 RSpec.shared_examples 'log import failure' do |importable_column|
   it 'tracks error' do
     extra = {
+              source: action,
               relation_key: relation_key,
               relation_index: relation_index,
               retry_count: retry_count
@@ -11,7 +12,12 @@ RSpec.shared_examples 'log import failure' do |importable_column|
 
     expect(Gitlab::ErrorTracking).to receive(:track_exception).with(exception, extra)
 
-    subject.log_import_failure(relation_key, relation_index, exception, retry_count)
+    subject.log_import_failure(
+      source: action,
+      relation_key: relation_key,
+      relation_index: relation_index,
+      exception: exception,
+      retry_count: retry_count)
   end
 
   it 'saves data to ImportFailure' do
@@ -21,6 +27,7 @@ RSpec.shared_examples 'log import failure' do |importable_column|
 
     aggregate_failures do
       expect(import_failure[importable_column]).to eq(importable.id)
+      expect(import_failure.source).to eq(action)
       expect(import_failure.relation_key).to eq(relation_key)
       expect(import_failure.relation_index).to eq(relation_index)
       expect(import_failure.exception_class).to eq('StandardError')