Track if IIDs are assigned within transaction

app/models/internal_id.rb

@@ -47,18 +47,10 @@ class InternalId < ApplicationRecord
   def update_and_save(&block)
     lock!
     yield
-    update_and_save_counter.increment(usage: usage, changed: last_value_changed?)
     save!
     last_value
   end
 
-  # Instrumentation to track for-update locks
-  def update_and_save_counter
-    strong_memoize(:update_and_save_counter) do
-      Gitlab::Metrics.counter(:gitlab_internal_id_for_update_lock, 'Number of ROW SHARE (FOR UPDATE) locks on individual records from internal_ids')
-    end
-  end
-
   class << self
     def track_greatest(subject, scope, usage, new_value, init)
       InternalIdGenerator.new(subject, scope, usage, init)
@@ -88,6 +80,8 @@ class InternalId < ApplicationRecord
   end
 
   class InternalIdGenerator
+    extend Gitlab::Utils::StrongMemoize
+
     # Generate next internal id for a given scope and usage.
     #
     # For currently supported usages, see #usage enum.
@@ -123,6 +117,8 @@ class InternalId < ApplicationRecord
     # init: Block that gets called to initialize InternalId record if not present
     #       Make sure to not throw exceptions in the absence of records (if this is expected).
     def generate
+      self.class.internal_id_transactions_increment(operation: :generate, usage: usage)
+
       subject.transaction do
         # Create a record in internal_ids if one does not yet exist
         # and increment its last value
@@ -138,6 +134,8 @@ class InternalId < ApplicationRecord
     def reset(value)
       return false unless value
 
+      self.class.internal_id_transactions_increment(operation: :reset, usage: usage)
+
       updated =
         InternalId
           .where(**scope, usage: usage_value)
@@ -152,6 +150,8 @@ class InternalId < ApplicationRecord
     #
     # Note this will acquire a ROW SHARE lock on the InternalId record
     def track_greatest(new_value)
+      self.class.internal_id_transactions_increment(operation: :track_greatest, usage: usage)
+
       subject.transaction do
         record.track_greatest_and_save!(new_value)
       end
@@ -162,6 +162,8 @@ class InternalId < ApplicationRecord
     end
 
     def with_lock(&block)
+      self.class.internal_id_transactions_increment(operation: :with_lock, usage: usage)
+
       record.with_lock(&block)
     end
 
@@ -197,5 +199,22 @@ class InternalId < ApplicationRecord
     rescue ActiveRecord::RecordNotUnique
       lookup
     end
+
+    def self.internal_id_transactions_increment(operation:, usage:)
+      self.internal_id_transactions_total.increment(
+        operation: operation,
+        usage: usage.to_s,
+        in_transaction: ActiveRecord::Base.connection.transaction_open?.to_s
+      )
+    end
+
+    def self.internal_id_transactions_total
+      strong_memoize(:internal_id_transactions_total) do
+        name = :gitlab_internal_id_transactions_total
+        comment = 'Counts all the internal ids happening within transaction'
+
+        Gitlab::Metrics.counter(name, comment)
+      end
+    end
   end
 end

spec/models/internal_id_spec.rb

@@ -97,6 +97,25 @@ RSpec.describe InternalId do
         expect(subject).to eq(1)
       end
     end
+
+    context 'when executed outside of transaction' do
+      it 'increments counter with in_transaction: "false"' do
+        expect(ActiveRecord::Base.connection).to receive(:transaction_open?) { false }
+        expect(InternalId::InternalIdGenerator.internal_id_transactions_total).to receive(:increment)
+          .with(operation: :generate, usage: 'issues', in_transaction: 'false').and_call_original
+
+        subject
+      end
+    end
+
+    context 'when executed within transaction' do
+      it 'increments counter with in_transaction: "true"' do
+        expect(InternalId::InternalIdGenerator.internal_id_transactions_total).to receive(:increment)
+          .with(operation: :generate, usage: 'issues', in_transaction: 'true').and_call_original
+
+        InternalId.transaction { subject }
+      end
+    end
   end
 
   describe '.reset' do
@@ -134,6 +153,29 @@ RSpec.describe InternalId do
         described_class.generate_next(issue, scope, usage, init)
       end
     end
+
+    context 'when executed outside of transaction' do
+      let(:value) { 2 }
+
+      it 'increments counter with in_transaction: "false"' do
+        expect(ActiveRecord::Base.connection).to receive(:transaction_open?) { false }
+        expect(InternalId::InternalIdGenerator.internal_id_transactions_total).to receive(:increment)
+          .with(operation: :reset, usage: 'issues', in_transaction: 'false').and_call_original
+
+        subject
+      end
+    end
+
+    context 'when executed within transaction' do
+      let(:value) { 2 }
+
+      it 'increments counter with in_transaction: "true"' do
+        expect(InternalId::InternalIdGenerator.internal_id_transactions_total).to receive(:increment)
+          .with(operation: :reset, usage: 'issues', in_transaction: 'true').and_call_original
+
+        InternalId.transaction { subject }
+      end
+    end
   end
 
   describe '.track_greatest' do
@@ -183,6 +225,25 @@ RSpec.describe InternalId do
         expect(subject).to eq(value)
       end
     end
+
+    context 'when executed outside of transaction' do
+      it 'increments counter with in_transaction: "false"' do
+        expect(ActiveRecord::Base.connection).to receive(:transaction_open?) { false }
+        expect(InternalId::InternalIdGenerator.internal_id_transactions_total).to receive(:increment)
+          .with(operation: :track_greatest, usage: 'issues', in_transaction: 'false').and_call_original
+
+        subject
+      end
+    end
+
+    context 'when executed within transaction' do
+      it 'increments counter with in_transaction: "true"' do
+        expect(InternalId::InternalIdGenerator.internal_id_transactions_total).to receive(:increment)
+          .with(operation: :track_greatest, usage: 'issues', in_transaction: 'true').and_call_original
+
+        InternalId.transaction { subject }
+      end
+    end
   end
 
   describe '#increment_and_save!' do