Avoid unnecessary locks on internal_ids

This avoid an unnecessary FOR UPDATE lock on internal id records when
the iid value has been generated internally anyways.

The lock is still required to track any values that come in through
other means, i.e. set by a user (through the API).

The broader picture is in
https://gitlab.com/gitlab-org/gitlab/issues/30515.

app/models/concerns/atomic_internal_id.rb

@@ -44,9 +44,15 @@ module AtomicInternalId
         scope_attrs = { scope_value.class.table_name.singularize.to_sym => scope_value }
         usage = self.class.table_name.to_sym
 
-        if value.present?
+        if value.present? && (@iid_needs_tracking || Feature.enabled?(:iid_always_track, default_enabled: true))
+          # The value was set externally, e.g. by the user
+          # We update the InternalId record to keep track of the greatest value.
           InternalId.track_greatest(self, scope_attrs, usage, value, init)
-        else
+
+          @iid_needs_tracking = false
+        elsif !value.present?
+          # We don't have a value yet and use a InternalId record to generate
+          # the next value.
           value = InternalId.generate_next(self, scope_attrs, usage, init)
           write_attribute(column, value)
         end
@@ -54,6 +60,13 @@ module AtomicInternalId
         value
       end
 
+      define_method("#{column}=") do |value|
+        super(value).tap do |v|
+          # Indicate the iid was set from externally
+          @iid_needs_tracking = true
+        end
+      end
+
       define_method("reset_#{scope}_#{column}") do
         if value = read_attribute(column)
           scope_value = association(scope).reader

spec/models/concerns/atomic_internal_id_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe AtomicInternalId do
+  let(:milestone) { build(:milestone) }
+  let(:iid) { double('iid', to_i: 42) }
+  let(:external_iid) { 100 }
+  let(:scope_attrs) { { project: milestone.project } }
+  let(:usage) { :milestones }
+
+  describe '#ensure_project_iid!' do
+    subject { milestone.ensure_project_iid! }
+
+    it 'generates a new value if non is present' do
+      expect(InternalId).to receive(:generate_next).with(milestone, scope_attrs, usage, anything).and_return(iid)
+
+      expect { subject }.to change { milestone.iid }.from(nil).to(iid.to_i)
+    end
+
+    it 'tracks the present value if not generated by InternalId.generate_next' do
+      milestone.iid = external_iid
+
+      expect(InternalId).to receive(:track_greatest).once.with(milestone, scope_attrs, usage, external_iid, anything)
+      expect(InternalId).not_to receive(:generate_next)
+
+      subject
+    end
+
+    it 'generates a new value if first set with iid= but later set to nil' do
+      expect(InternalId).to receive(:generate_next).with(milestone, scope_attrs, usage, anything).and_return(iid)
+
+      milestone.iid = external_iid
+      milestone.iid = nil
+
+      expect { subject }.to change { milestone.iid }.from(nil).to(iid.to_i)
+    end
+
+    context 'with iid_always_track disabled' do
+      before do
+        stub_feature_flags(iid_always_track: false)
+      end
+
+      it 'does not track the present value if generated by InternalId.generate_next' do
+        milestone.ensure_project_iid!
+
+        expect(InternalId).not_to receive(:track_greatest)
+
+        subject
+      end
+    end
+  end
+end