Add `stack_trace_record` into transaction to observe `COMMIT`

This fixes usage of `after_commit` - a set of operations
that are executed after the object is committed. Before that
change the operations would be treated as being executed
in context of a transaction, even though the transaction
was already finished.

This introduces trace record object that receives callback
once the after commit/rollback hooks are to be executed.

app/models/concerns/cross_database_modification.rb

@@ -3,6 +3,38 @@
 module CrossDatabaseModification
   extend ActiveSupport::Concern
 
+  class TransactionStackTrackRecord
+    def initialize(subject, gitlab_schema)
+      @subject = subject
+      @gitlab_schema = gitlab_schema
+      @subject.gitlab_transactions_stack.push(gitlab_schema)
+    end
+
+    def done!
+      unless @done
+        @done = true
+        @subject.gitlab_transactions_stack.pop
+      end
+
+      true
+    end
+
+    def trigger_transactional_callbacks?
+      false
+    end
+
+    def before_committed!
+    end
+
+    def rolledback!(force_restore_state: false, should_run_callbacks: true)
+      done!
+    end
+
+    def committed!(should_run_callbacks: true)
+      done!
+    end
+  end
+
   included do
     private_class_method :gitlab_schema
   end
@@ -14,12 +46,19 @@ module CrossDatabaseModification
 
     def transaction(**options, &block)
       if track_gitlab_schema_in_current_transaction?
-        gitlab_transactions_stack.push(gitlab_schema)
+        super(**options) do
+          # Hook into current transaction to ensure that once
+          # the `COMMIT` is executed the `gitlab_transactions_stack`
+          # will be allowing to execute `after_commit_queue`
+          record = TransactionStackTrackRecord.new(self, gitlab_schema)
+
+          begin
+            connection.current_transaction.add_record(record)
 
-        begin
-          super(**options, &block)
-        ensure
-          gitlab_transactions_stack.pop
+            yield
+          ensure
+            record.done!
+          end
         end
       else
         super(**options, &block)

config/feature_flags/development/track_gitlab_schema_in_current_transaction.yml

@@ -2,7 +2,7 @@
 name: track_gitlab_schema_in_current_transaction
 introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/76717
 rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/349944
-milestone: '14.7'
+milestone: '14.8'
 type: development
 group: group::sharding
 default_enabled: false

ee/spec/models/ci/pipeline_spec.rb

@@ -361,11 +361,9 @@ RSpec.describe Ci::Pipeline do
     end
 
     context 'when pipeline project has downstream subscriptions' do
-      let(:pipeline) { create(:ci_empty_pipeline, project: create(:project, :public)) }
-
-      before do
-        pipeline.project.downstream_projects << create(:project)
-      end
+      let(:downstream_project) { create(:project) }
+      let(:project) { create(:project, :public, downstream_projects: [downstream_project]) }
+      let(:pipeline) { create(:ci_empty_pipeline, project: project) }
 
       context 'when pipeline runs on a tag' do
         before do

ee/spec/models/project_spec.rb

@@ -3231,9 +3231,8 @@ RSpec.describe Project do
 
   describe '#upstream_projects' do
     it 'returns the upstream projects' do
-      primary_project = create(:project, :public)
       upstream_project = create(:project, :public)
-      primary_project.upstream_projects << upstream_project
+      primary_project = create(:project, :public, upstream_projects: [upstream_project])
 
       with_cross_joins_prevented do
         expect(primary_project.upstream_projects).to eq([upstream_project])
@@ -3243,9 +3242,8 @@ RSpec.describe Project do
 
   describe '#upstream_projects_count' do
     it 'returns the upstream projects count' do
-      primary_project = create(:project, :public)
       upstream_projects = create_list(:project, 2, :public)
-      primary_project.upstream_projects = upstream_projects
+      primary_project = create(:project, :public, upstream_projects: upstream_projects)
 
       with_cross_joins_prevented do
         expect(primary_project.upstream_projects_count).to eq(2)
@@ -3255,9 +3253,8 @@ RSpec.describe Project do
 
   describe '#downstream_projects' do
     it 'returns the downstream projects' do
-      primary_project = create(:project, :public)
       downstream_project = create(:project, :public)
-      primary_project.downstream_projects << downstream_project
+      primary_project = create(:project, :public, downstream_projects: [downstream_project])
 
       with_cross_joins_prevented do
         expect(primary_project.downstream_projects).to eq([downstream_project])
@@ -3267,9 +3264,8 @@ RSpec.describe Project do
 
   describe '#downstream_projects_count' do
     it 'returns the downstream projects count' do
-      primary_project = create(:project, :public)
       downstream_projects = create_list(:project, 2, :public)
-      primary_project.downstream_projects = downstream_projects
+      primary_project = create(:project, :public, downstream_projects: downstream_projects)
 
       with_cross_joins_prevented do
         expect(primary_project.downstream_projects_count).to eq(2)

spec/models/concerns/cross_database_modification_spec.rb

@@ -45,12 +45,10 @@ RSpec.describe CrossDatabaseModification do
 
       expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
 
-      recorder = ActiveRecord::QueryRecorder.new do
-        Ci::ApplicationRecord.transaction do
-          expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_ci)
+      Ci::ApplicationRecord.transaction do
+        expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_ci)
 
-          Project.first
-        end
+        Project.first
       end
 
       expect(ApplicationRecord.gitlab_transactions_stack).to be_empty