Reduce subtransactions in safe_find_or_create_by

A subtraction will only be started when attempting to create a record

app/models/application_record.rb

@@ -63,11 +63,27 @@ class ApplicationRecord < ActiveRecord::Base
   end
 
   def self.safe_find_or_create_by(*args, &block)
+    return optimized_safe_find_or_create_by(*args, &block) if Feature.enabled?(:optimize_safe_find_or_create_by, default_enabled: :yaml)
+
     safe_ensure_unique(retries: 1) do
       find_or_create_by(*args, &block)
     end
   end
 
+  def self.optimized_safe_find_or_create_by(*args, &block)
+    record = find_by(*args)
+    return record if record.present?
+
+    # We need to use `all.create` to make this implementation follow `find_or_create_by` which delegates this in
+    # https://github.com/rails/rails/blob/v6.1.3.2/activerecord/lib/active_record/querying.rb#L22
+    #
+    # When calling this method on an association, just calling `self.create` would call `ActiveRecord::Persistence.create`
+    # and that skips some code that adds the newly created record to the association.
+    transaction(requires_new: true) { all.create(*args, &block) }
+  rescue ActiveRecord::RecordNotUnique
+    find_by(*args)
+  end
+
   def create_or_load_association(association_name)
     association(association_name).create unless association(association_name).loaded?
   rescue ActiveRecord::RecordNotUnique, PG::UniqueViolation

config/feature_flags/development/optimize_safe_find_or_create_by.yml

+---
+name: optimize_safe_find_or_create_by
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/68458
+rollout_issue_url: 
+milestone: '14.3'
+type: development
+group: group::database
+default_enabled: false

spec/models/application_record_spec.rb

@@ -39,13 +39,14 @@ RSpec.describe ApplicationRecord do
 
     let(:suggestion_attributes) { attributes_for(:suggestion).merge!(note_id: note.id) }
 
-    describe '.safe_find_or_create_by' do
+    shared_examples '.safe_find_or_create_by' do
       it 'creates the suggestion avoiding race conditions' do
-        expect(Suggestion).to receive(:find_or_create_by).and_raise(ActiveRecord::RecordNotUnique)
-        allow(Suggestion).to receive(:find_or_create_by).and_call_original
+        existing_suggestion = double(:Suggestion)
 
-        expect { Suggestion.safe_find_or_create_by(suggestion_attributes) }
-          .to change { Suggestion.count }.by(1)
+        expect(Suggestion).to receive(:find_by).and_return(nil, existing_suggestion)
+        expect(Suggestion).to receive(:create).and_raise(ActiveRecord::RecordNotUnique)
+
+        expect(Suggestion.safe_find_or_create_by(suggestion_attributes)).to eq(existing_suggestion)
       end
 
       it 'passes a block to find_or_create_by' do
@@ -62,10 +63,8 @@ RSpec.describe ApplicationRecord do
       end
     end
 
-    describe '.safe_find_or_create_by!' do
+    shared_examples '.safe_find_or_create_by!' do
       it 'creates a record using safe_find_or_create_by' do
-        expect(Suggestion).to receive(:find_or_create_by).and_call_original
-
         expect(Suggestion.safe_find_or_create_by!(suggestion_attributes))
           .to be_a(Suggestion)
       end
@@ -89,6 +88,24 @@ RSpec.describe ApplicationRecord do
           .to raise_error(ActiveRecord::RecordNotFound)
       end
     end
+
+    context 'when optimized_safe_find_or_create_by is enabled' do
+      before do
+        stub_feature_flags(optimized_safe_find_or_create_by: true)
+      end
+
+      it_behaves_like '.safe_find_or_create_by'
+      it_behaves_like '.safe_find_or_create_by!'
+    end
+
+    context 'when optimized_safe_find_or_create_by is disabled' do
+      before do
+        stub_feature_flags(optimized_safe_find_or_create_by: false)
+      end
+
+      it_behaves_like '.safe_find_or_create_by'
+      it_behaves_like '.safe_find_or_create_by!'
+    end
   end
 
   describe '.underscore' do

spec/models/gpg_signature_spec.rb

@@ -54,8 +54,10 @@ RSpec.describe GpgSignature do
     end
 
     it 'does not raise an error in case of a race condition' do
-      expect(described_class).to receive(:find_or_create_by).and_raise(ActiveRecord::RecordNotUnique)
-      allow(described_class).to receive(:find_or_create_by).and_call_original
+      expect(described_class).to receive(:find_by).and_return(nil, double(described_class, persisted?: true))
+
+      expect(described_class).to receive(:create).and_raise(ActiveRecord::RecordNotUnique)
+      allow(described_class).to receive(:create).and_call_original
 
       described_class.safe_create!(attributes)
     end