Add ApplicationRecord#safe_ensure_unique method

Port of https://dev.gitlab.org/gitlab/gitlab-ee/merge_requests/866
to CE excluding the migration and service changes as they don't
apply to CE.

2 files changed, 33 insertions, 3 deletions

diff --git a/app/models/application_record.rb b/app/models/application_record.rb
index 6976185264e..348468197bf 100644
--- a/app/models/application_record.rb
+++ b/app/models/application_record.rb
@@ -15,6 +15,19 @@ class ApplicationRecord < ActiveRecord::Base
     where(nil).pluck(self.primary_key)
   end
 
+  def self.safe_ensure_unique(retries: 0)
+    transaction(requires_new: true) do
+      yield
+    end
+  rescue ActiveRecord::RecordNotUnique
+    if retries > 0
+      retries -= 1
+      retry
+    end
+
+    false
+  end
+
   def self.safe_find_or_create_by!(*args)
     safe_find_or_create_by(*args).tap do |record|
       record.validate! unless record.persisted?
@@ -22,10 +35,8 @@ class ApplicationRecord < ActiveRecord::Base
   end
 
   def self.safe_find_or_create_by(*args)
-    transaction(requires_new: true) do
+    safe_ensure_unique(retries: 1) do
       find_or_create_by(*args)
     end
-  rescue ActiveRecord::RecordNotUnique
-    retry
   end
 end
diff --git a/spec/models/application_record_spec.rb b/spec/models/application_record_spec.rb
index fd25132ed3a..cc90a998d3f 100644
--- a/spec/models/application_record_spec.rb
+++ b/spec/models/application_record_spec.rb
@@ -11,6 +11,25 @@ describe ApplicationRecord do
     end
   end
 
+  describe '.safe_ensure_unique' do
+    let(:model) { build(:suggestion) }
+    let(:klass) { model.class }
+
+    before do
+      allow(model).to receive(:save).and_raise(ActiveRecord::RecordNotUnique)
+    end
+
+    it 'returns false when ActiveRecord::RecordNotUnique is raised' do
+      expect(model).to receive(:save).once
+      expect(klass.safe_ensure_unique { model.save }).to be_falsey
+    end
+
+    it 'retries based on retry count specified' do
+      expect(model).to receive(:save).exactly(3).times
+      expect(klass.safe_ensure_unique(retries: 2) { model.save }).to be_falsey
+    end
+  end
+
   describe '.safe_find_or_create_by' do
     it 'creates the user avoiding race conditions' do
       expect(Suggestion).to receive(:find_or_create_by).and_raise(ActiveRecord::RecordNotUnique)