Refine bulk insert tags

Gemfile

@@ -194,7 +194,7 @@ end
 # State machine
 gem 'state_machines-activerecord', '~> 0.8.0'
-# Issue tags
+# CI domain tags
 gem 'acts-as-taggable-on', '~> 9.0'
 # Background jobs

app/models/ci/build.rb

@@ -420,6 +420,10 @@ def runnable?
       true
     end
+    def save_tags
+      super unless Thread.current['ci_bulk_insert_tags']
+    end
+
     def archived?
       return true if degenerated?

app/models/commit_status.rb

@@ -221,8 +221,8 @@ def self.locking_enabled?
     false
   end
-  def self.bulk_insert_tags!(statuses, tag_list_by_build)
-    Gitlab::Ci::Tags::BulkInsert.new(statuses, tag_list_by_build).insert!
+  def self.bulk_insert_tags!(statuses)
+    Gitlab::Ci::Tags::BulkInsert.new(statuses).insert!
   end
   def locking_enabled?

lib/gitlab/ci/pipeline/chain/create.rb

@@ -11,11 +11,11 @@ class Create < Chain::Base
           def perform!
             logger.instrument_with_sql(:pipeline_save) do
               BulkInsertableAssociations.with_bulk_insert do
-                tags = extract_tag_list_by_status
-
-                pipeline.transaction do
-                  pipeline.save!
-                  CommitStatus.bulk_insert_tags!(statuses, tags) if bulk_insert_tags?
+                with_bulk_insert_tags do
+                  pipeline.transaction do
+                    pipeline.save!
+                    CommitStatus.bulk_insert_tags!(statuses) if bulk_insert_tags?
+                  end
                 end
               end
             end
@@ -29,32 +29,26 @@ def break?
           private
-          def statuses
-            strong_memoize(:statuses) do
-              pipeline.stages.flat_map(&:statuses)
+          def bulk_insert_tags?
+            strong_memoize(:bulk_insert_tags) do
+              ::Feature.enabled?(:ci_bulk_insert_tags, project, default_enabled: :yaml)
             end
           end
-          # We call `job.tag_list=` to assign tags to the jobs from the
-          # Chain::Seed step which uses the `@tag_list` instance variable to
-          # store them on the record. We remove them here because we want to
-          # bulk insert them, otherwise they would be inserted and assigned one
-          # by one with callbacks. We must use `remove_instance_variable`
-          # because having the instance variable defined would still run the callbacks
-          def extract_tag_list_by_status
-            return {} unless bulk_insert_tags?
-
-            statuses.each.with_object({}) do |job, acc|
-              tag_list = job.clear_memoization(:tag_list)
-              next unless tag_list
-
-              acc[job.name] = tag_list
-            end
+          def with_bulk_insert_tags
+            previous = Thread.current['ci_bulk_insert_tags']
+            Thread.current['ci_bulk_insert_tags'] = bulk_insert_tags?
+            yield
+          ensure
+            Thread.current['ci_bulk_insert_tags'] = previous
           end
-          def bulk_insert_tags?
-            strong_memoize(:bulk_insert_tags) do
-              ::Feature.enabled?(:ci_bulk_insert_tags, project, default_enabled: :yaml)
+          def statuses
+            strong_memoize(:statuses) do
+              pipeline
+                .stages
+                .flat_map(&:statuses)
+                .select { |status| status.respond_to?(:tag_list) }
             end
           end
         end

lib/gitlab/ci/tags/bulk_insert.rb

@@ -4,12 +4,13 @@ module Gitlab
   module Ci
     module Tags
       class BulkInsert
+        include Gitlab::Utils::StrongMemoize
+
         TAGGINGS_BATCH_SIZE = 1000
         TAGS_BATCH_SIZE = 500
-        def initialize(statuses, tag_list_by_status)
+        def initialize(statuses)
           @statuses = statuses
-          @tag_list_by_status = tag_list_by_status
         end
         def insert!
@@ -20,7 +21,18 @@ def insert!
         private
-        attr_reader :statuses, :tag_list_by_status
+        attr_reader :statuses
+
+        def tag_list_by_status
+          strong_memoize(:tag_list_by_status) do
+            statuses.each.with_object({}) do |status, acc|
+              tag_list = status.tag_list
+              next unless tag_list
+
+              acc[status] = tag_list
+            end
+          end
+        end
         def persist_build_tags!
           all_tags = tag_list_by_status.values.flatten.uniq.reject(&:blank?)
@@ -54,7 +66,7 @@ def create_tags(tags)
         def build_taggings_attributes(tag_records_by_name)
           taggings = statuses.flat_map do |status|
-            tag_list = tag_list_by_status[status.name]
+            tag_list = tag_list_by_status[status]
             next unless tag_list
             tags = tag_records_by_name.values_at(*tag_list)

spec/lib/gitlab/ci/pipeline/chain/create_spec.rb

@@ -79,12 +79,11 @@
       it 'extracts an empty tag list' do
         expect(CommitStatus)
           .to receive(:bulk_insert_tags!)
-          .with(stage.statuses, {})
+          .with([job])
           .and_call_original
         step.perform!
-        expect(job.instance_variable_defined?(:@tag_list)).to be_falsey
         expect(job).to be_persisted
         expect(job.tag_list).to eq([])
       end
@@ -98,14 +97,13 @@
       it 'bulk inserts tags' do
         expect(CommitStatus)
           .to receive(:bulk_insert_tags!)
-          .with(stage.statuses, { job.name => %w[tag1 tag2] })
+          .with([job])
           .and_call_original
         step.perform!
-        expect(job.instance_variable_defined?(:@tag_list)).to be_falsey
         expect(job).to be_persisted
-        expect(job.tag_list).to match_array(%w[tag1 tag2])
+        expect(job.reload.tag_list).to match_array(%w[tag1 tag2])
       end
     end
@@ -120,7 +118,6 @@
         step.perform!
-        expect(job.instance_variable_defined?(:@tag_list)).to be_truthy
         expect(job).to be_persisted
         expect(job.reload.tag_list).to match_array(%w[tag1 tag2])
       end

spec/lib/gitlab/ci/tags/bulk_insert_spec.rb

@@ -5,27 +5,37 @@
 RSpec.describe Gitlab::Ci::Tags::BulkInsert do
   let_it_be(:project) { create(:project, :repository) }
   let_it_be(:pipeline) { create(:ci_pipeline, project: project) }
-  let_it_be_with_refind(:job) { create(:ci_build, :unique_name, pipeline: pipeline, project: project) }
-  let_it_be_with_refind(:other_job) { create(:ci_build, :unique_name, pipeline: pipeline, project: project) }
-  let_it_be_with_refind(:bridge) { create(:ci_bridge, pipeline: pipeline, project: project) }
-  let(:statuses) { [job, bridge, other_job] }
-  subject(:service) { described_class.new(statuses, tags_list) }
+  let_it_be_with_refind(:job) { create(:ci_build, :unique_name, pipeline: pipeline) }
+  let_it_be_with_refind(:other_job) { create(:ci_build, :unique_name, pipeline: pipeline) }
+  let(:statuses) { [job, other_job] }
+  subject(:service) { described_class.new(statuses) }
+
+  describe 'gem version' do
+    let(:acceptable_version) { '9.0.0' }
+
+    let(:error_message) do
+      <<~MESSAGE
+      A mechanism depending on internals of 'act-as-taggable-on` has been designed
+      to bulk insert tags for Ci::Build records.
+      Please review the code carefully before updating the gem version
+      https://gitlab.com/gitlab-org/gitlab/-/issues/350053
+      MESSAGE
+    end
+
+    it { expect(ActsAsTaggableOn::VERSION).to eq(acceptable_version), error_message }
+  end
   describe '#insert!' do
     context 'without tags' do
-      let(:tags_list) { {} }
-
       it { expect(service.insert!).to be_falsey }
     end
     context 'with tags' do
-      let(:tags_list) do
-        {
-          job.name => %w[tag1 tag2],
-          other_job.name => %w[tag2 tag3 tag4]
-        }
+      before do
+        job.tag_list = %w[tag1 tag2]
+        other_job.tag_list = %w[tag2 tag3 tag4]
       end
       it 'persists tags' do
@@ -35,5 +45,18 @@
         expect(other_job.reload.tag_list).to match_array(%w[tag2 tag3 tag4])
       end
     end
+
+    context 'with tags for only one job' do
+      before do
+        job.tag_list = %w[tag1 tag2]
+      end
+
+      it 'persists tags' do
+        expect(service.insert!).to be_truthy
+
+        expect(job.reload.tag_list).to match_array(%w[tag1 tag2])
+        expect(other_job.reload.tag_list).to be_empty
+      end
+    end
   end
 end

spec/models/commit_status_spec.rb

@@ -961,18 +961,17 @@ def create_status(**opts)
   describe '.bulk_insert_tags!' do
     let(:statuses) { double('statuses') }
-    let(:tag_list_by_build) { double('tag list') }
     let(:inserter) { double('inserter') }
     it 'delegates to bulk insert class' do
       expect(Gitlab::Ci::Tags::BulkInsert)
         .to receive(:new)
-        .with(statuses, tag_list_by_build)
+        .with(statuses)
         .and_return(inserter)
       expect(inserter).to receive(:insert!)
-      described_class.bulk_insert_tags!(statuses, tag_list_by_build)
+      described_class.bulk_insert_tags!(statuses)
     end
   end