More flexible way of internal id generation.

app/models/concerns/atomic_internal_id.rb

@@ -2,15 +2,19 @@ module AtomicInternalId
   extend ActiveSupport::Concern
 
   included do
-    before_validation(on: :create) do
-      set_iid
-    end
-
-    validates :iid, presence: true, numericality: true
-  end
+    class << self
+      def has_internal_id(on, scope:, usage: nil, init: nil)
+        before_validation(on: :create) do
+          if self.public_send(on).blank? # rubocop:disable GitlabSecurity/PublicSend
+            usage = (usage || self.class.name.tableize).to_sym
+            new_iid = InternalId.generate_next(self, scope, usage, init)
+            self.public_send("#{on}=", new_iid) # rubocop:disable GitlabSecurity/PublicSend
+          end
+        end
 
-  def set_iid
-    self.iid = InternalId.generate_next(self.project, :issues) if iid.blank?
+        validates on, presence: true, numericality: true
+      end
+    end
   end
 
   def to_param

app/models/internal_id.rb

@@ -23,17 +23,6 @@ class InternalId < ActiveRecord::Base
     last_value
   end
 
-  before_create :calculate_last_value!
-
-  # Calculate #last_value by counting the number of
-  # existing records for this usage.
-  def calculate_last_value!
-    return if last_value
-
-    parent = project # ??|| group
-    self.last_value = parent.send(usage.to_sym).maximum(:iid) || 0 # rubocop:disable GitlabSecurity/PublicSend
-  end
-
   class << self
     # Generate next internal id for a given project and usage.
     #
@@ -45,12 +34,21 @@ class InternalId < ActiveRecord::Base
     # 3) The generated sequence is gapless.
     # 4) In the absence of a record in the internal_ids table, one will be created
     #    and last_value will be calculated on the fly.
-    def generate_next(project, usage)
-      raise 'project not set - this is required' unless project
+    def generate_next(subject, scope, usage, init)
+      scope = [scope].flatten.compact
+      raise 'scope is not well-defined, need at least one column for scope (given: 0)' if scope.empty?
+      raise "usage #{usage} is unknown. Supported values are InternalId.usages = #{InternalId.usages.keys.to_s}" unless InternalId.usages.include?(usage.to_sym)
+
+      init ||= ->(s) { 0 }
+
+      scope_attrs = scope.inject({}) do |h, e|
+        h[e] = subject.public_send(e)
+        h
+      end
 
-      project.transaction do
+      transaction do
         # Create a record in internal_ids if one does not yet exist
-        id = (lookup(project, usage) || create_record(project, usage))
+        id = (lookup(scope_attrs, usage) || create_record(scope_attrs, usage, init, subject))
 
         # This will lock the InternalId record with ROW SHARE
         # and increment #last_value
@@ -61,8 +59,8 @@ class InternalId < ActiveRecord::Base
     private
 
     # Retrieve InternalId record for (project, usage) combination, if it exists
-    def lookup(project, usage)
-      project.internal_ids.find_by(usage: usages[usage.to_s])
+    def lookup(scope_attrs, usage)
+      InternalId.find_by(usage: usages[usage.to_s], **scope_attrs)
     end
 
     # Create InternalId record for (project, usage) combination, if it doesn't exist
@@ -71,13 +69,13 @@ class InternalId < ActiveRecord::Base
     # was faster in doing this, we'll realize once we hit the unique key constraint
     # violation. We can safely roll-back the nested transaction and perform
     # a lookup instead to retrieve the record.
-    def create_record(project, usage)
+    def create_record(scope_attrs, usage, init, subject)
       begin
-        project.transaction(requires_new: true) do
-          create!(project: project, usage: usages[usage.to_s])
+        transaction(requires_new: true) do
+          create!(usage: usages[usage.to_s], **scope_attrs, last_value: init.call(subject) || 0)
         end
       rescue ActiveRecord::RecordNotUnique
-        lookup(project, usage)
+        lookup(scope_attrs, usage)
       end
     end
   end

app/models/issue.rb

@@ -24,6 +24,8 @@ class Issue < ActiveRecord::Base
   belongs_to :project
   belongs_to :moved_to, class_name: 'Issue'
 
+  has_internal_id :iid, scope: :project, init: ->(s) { s.project.issues.maximum(:iid) }
+
   has_many :events, as: :target, dependent: :destroy # rubocop:disable Cop/ActiveRecordDependent
 
   has_many :merge_requests_closing_issues,

spec/factories/internal_ids.rb

 FactoryBot.define do
   factory :internal_id do
     project
-    usage { InternalId.usages.keys.first }
+    usage :issues
+    last_value { project.issues.maximum(:iid) || 0 }
   end
 end

spec/models/internal_id_spec.rb

@@ -3,6 +3,9 @@ require 'spec_helper'
 describe InternalId do
   let(:project) { create(:project) }
   let(:usage) { :issues }
+  let(:issue) { build(:issue, project: project) }
+  let(:scope) { :project }
+  let(:init) { ->(s) { project.issues.size } }
 
   context 'validations' do
     it { is_expected.to validate_presence_of(:usage) }
@@ -11,7 +14,7 @@ describe InternalId do
 
   describe '.generate_next' do
     context 'in the absence of a record' do
-      subject { described_class.generate_next(project, usage) }
+      subject { described_class.generate_next(issue, scope, usage, init) }
 
       it 'creates a record if not yet present' do
         expect { subject }.to change { described_class.count }.from(0).to(1)
@@ -22,13 +25,14 @@ describe InternalId do
 
         described_class.first.tap do |record|
           expect(record.project).to eq(project)
-          expect(record.usage).to eq(usage.to_s) # TODO
+          expect(record.usage).to eq(usage.to_s)
         end
       end
 
       context 'with existing issues' do
         before do
           rand(10).times { create(:issue, project: project) }
+          InternalId.delete_all
         end
 
         it 'calculates last_value values automatically' do
@@ -39,7 +43,7 @@ describe InternalId do
 
     it 'generates a strictly monotone, gapless sequence' do
       seq = (0..rand(1000)).map do
-        described_class.generate_next(project, usage)
+        described_class.generate_next(issue, scope, usage, init)
       end
       normalized = seq.map { |i| i - seq.min }
       expect(normalized).to eq((0..seq.size - 1).to_a)
@@ -68,20 +72,4 @@ describe InternalId do
       end
     end
   end
-
-  describe '#calculate_last_value! (for issues)' do
-    subject do
-      build(:internal_id, project: project, usage: :issues)
-    end
-
-    context 'with existing issues' do
-      before do
-        rand(10).times { create(:issue, project: project) }
-      end
-
-      it 'counts related issues and saves' do
-        expect { subject.calculate_last_value! }.to change { subject.last_value }.from(nil).to(project.issues.size)
-      end
-    end
-  end
 end