Fix regression for models with multiple iid scopes

For models with multiple iid scopes (Milestone), the flag
`@internal_id_needs_tracking` should only be reset if the tracking was
actually done.

Most likely introduced by the refactoring in
https://gitlab.com/gitlab-org/gitlab/merge_requests/18003

app/models/concerns/atomic_internal_id.rb

@@ -60,12 +60,11 @@ module AtomicInternalId
         iid_always_track = Feature.enabled?(:iid_always_track, default_enabled: true)
         return unless @internal_id_needs_tracking || iid_always_track
 
-        @internal_id_needs_tracking = false
-
         scope_value = internal_id_read_scope(scope)
-        value = read_attribute(column)
         return unless scope_value
 
+        value = read_attribute(column)
+
         if value.present?
           # The value was set externally, e.g. by the user
           # We update the InternalId record to keep track of the greatest value.
@@ -75,6 +74,8 @@ module AtomicInternalId
             internal_id_scope_usage,
             value,
             init)
+
+          @internal_id_needs_tracking = false
         end
       end
 

spec/models/concerns/atomic_internal_id_spec.rb

@@ -22,7 +22,7 @@ describe AtomicInternalId do
     end
 
     context 'when value is set by ensure_project_iid!' do
-      context 'with iid_always_track true' do
+      context 'with iid_always_track false' do
         before do
           stub_feature_flags(iid_always_track: false)
         end
@@ -33,6 +33,17 @@ describe AtomicInternalId do
           milestone.ensure_project_iid!
           subject
         end
+
+        it 'tracks the iid for the scope that is actually present' 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)
+
+          # group scope is not present here, the milestone does not have a group
+          milestone.track_group_iid!
+          subject
+        end
       end
 
       context 'with iid_always_track enabled' do

spec/support/shared_examples/models/atomic_internal_id_shared_examples.rb

@@ -47,6 +47,10 @@ shared_examples_for 'AtomicInternalId' do |validate_presence: true|
     end
 
     describe 'internal id generation' do
+      before do
+        stub_feature_flags(iid_always_track: false)
+      end
+
       subject { instance.save! }
 
       it 'calls InternalId.generate_next and sets internal id attribute' do