Move populate_scheduling_type logic

We need this logic where only retrying mechanism occurs.
With this movement, we'll achieve to remove redis calls.
We added "ci_ensure_scheduling_type" FF just in case.

app/models/ci/pipeline.rb

@@ -961,6 +961,14 @@ module Ci
       end
     end
 
+    # Set scheduling type of processables if they were created before scheduling_type
+    # data was deployed (https://gitlab.com/gitlab-org/gitlab/-/merge_requests/22246).
+    def ensure_scheduling_type!
+      return unless ::Gitlab::Ci::Features.ensure_scheduling_type_enabled?
+
+      processables.populate_scheduling_type!
+    end
+
     private
 
     def pipeline_data

app/models/ci/processable.rb

@@ -99,5 +99,13 @@ module Ci
         needs.map { |need| need.attributes.except('id', 'build_id') }
       end
     end
+
+    def ensure_scheduling_type!
+      # If this has a scheduling_type, it means all processables in the pipeline already have.
+      return if scheduling_type
+
+      pipeline.ensure_scheduling_type!
+      reset
+    end
   end
 end

app/services/ci/process_pipeline_service.rb

@@ -10,7 +10,6 @@ module Ci
 
     def execute(trigger_build_ids = nil, initial_process: false)
       update_retried
-      ensure_scheduling_type_for_processables
 
       if Feature.enabled?(:ci_atomic_processing, pipeline.project)
         Ci::PipelineProcessing::AtomicProcessingService
@@ -44,17 +43,5 @@ module Ci
         .update_all(retried: true) if latest_statuses.any?
     end
     # rubocop: enable CodeReuse/ActiveRecord
-
-    # Set scheduling type of processables if they were created before scheduling_type
-    # data was deployed (https://gitlab.com/gitlab-org/gitlab/-/merge_requests/22246).
-    # Given that this service runs multiple times during the pipeline
-    # life cycle we need to ensure we populate the data once.
-    # See more: https://gitlab.com/gitlab-org/gitlab/issues/205426
-    def ensure_scheduling_type_for_processables
-      lease = Gitlab::ExclusiveLease.new("set-scheduling-types:#{pipeline.id}", timeout: 1.hour.to_i)
-      return unless lease.try_obtain
-
-      pipeline.processables.populate_scheduling_type!
-    end
   end
 end

app/services/ci/retry_build_service.rb

@@ -9,6 +9,8 @@ module Ci
                          resource_group scheduling_type].freeze
 
     def execute(build)
+      build.ensure_scheduling_type!
+
       reprocess!(build).tap do |new_build|
         build.pipeline.mark_as_processable_after_stage(build.stage_idx)
 
@@ -31,6 +33,9 @@ module Ci
       end.to_h
 
       attributes[:user] = current_user
+
+      # TODO: we can probably remove this logic
+      # see: https://gitlab.com/gitlab-org/gitlab/-/issues/217930
       attributes[:scheduling_type] ||= build.find_legacy_scheduling_type
 
       Ci::Build.transaction do

app/services/ci/retry_pipeline_service.rb

@@ -11,6 +11,8 @@ module Ci
 
       needs = Set.new
 
+      pipeline.ensure_scheduling_type!
+
       pipeline.retryable_builds.preload_needs.find_each do |build|
         next unless can?(current_user, :update_build, build)
 

lib/gitlab/ci/features.rb

@@ -9,6 +9,10 @@ module Gitlab
       def self.artifacts_exclude_enabled?
         ::Feature.enabled?(:ci_artifacts_exclude, default_enabled: false)
       end
+
+      def self.ensure_scheduling_type_enabled?
+        ::Feature.enabled?(:ci_ensure_scheduling_type, default_enabled: true)
+      end
     end
   end
 end

spec/services/ci/process_pipeline_service_spec.rb

@@ -33,25 +33,6 @@ describe Ci::ProcessPipelineService do
     end
   end
 
-  context 'with a pipeline which has processables with nil scheduling_type', :clean_gitlab_redis_shared_state do
-    let!(:build1) { create_build('build1') }
-    let!(:build2) { create_build('build2') }
-    let!(:build3) { create_build('build3', scheduling_type: :dag) }
-    let!(:build3_on_build2) { create(:ci_build_need, build: build3, name: 'build2') }
-
-    before do
-      pipeline.processables.update_all(scheduling_type: nil)
-    end
-
-    it 'populates scheduling_type before processing' do
-      process_pipeline
-
-      expect(build1.scheduling_type).to eq('stage')
-      expect(build2.scheduling_type).to eq('stage')
-      expect(build3.scheduling_type).to eq('dag')
-    end
-  end
-
   def process_pipeline
     described_class.new(pipeline).execute
   end

spec/services/ci/retry_build_service_spec.rb

@@ -190,6 +190,35 @@ describe Ci::RetryBuildService do
           expect(subsequent_build.reload).to be_created
         end
       end
+
+      context 'when pipeline has other builds' do
+        let!(:stage2) { create(:ci_stage_entity, project: project, pipeline: pipeline, name: 'deploy') }
+        let!(:build2) { create(:ci_build, pipeline: pipeline, stage_id: stage.id ) }
+        let!(:deploy) { create(:ci_build, pipeline: pipeline, stage_id: stage2.id) }
+        let!(:deploy_needs_build2) { create(:ci_build_need, build: deploy, name: build2.name) }
+
+        context 'when build has nil scheduling_type' do
+          before do
+            build.pipeline.processables.update_all(scheduling_type: nil)
+            build.reload
+          end
+
+          it 'populates scheduling_type of processables' do
+            expect(new_build.scheduling_type).to eq('stage')
+            expect(build.reload.scheduling_type).to eq('stage')
+            expect(build2.reload.scheduling_type).to eq('stage')
+            expect(deploy.reload.scheduling_type).to eq('dag')
+          end
+        end
+
+        context 'when build has scheduling_type' do
+          it 'does not call populate_scheduling_type!' do
+            expect_any_instance_of(Ci::Pipeline).not_to receive(:ensure_scheduling_type!)
+
+            expect(new_build.scheduling_type).to eq('stage')
+          end
+        end
+      end
     end
 
     context 'when user does not have ability to execute build' do

spec/services/ci/retry_pipeline_service_spec.rb

@@ -261,6 +261,25 @@ describe Ci::RetryPipelineService, '#execute' do
 
       service.execute(pipeline)
     end
+
+    context 'when pipeline has processables with nil scheduling_type' do
+      let!(:build1) { create_build('build1', :success, 0) }
+      let!(:build2) { create_build('build2', :failed, 0) }
+      let!(:build3) { create_build('build3', :failed, 1) }
+      let!(:build3_needs_build1) { create(:ci_build_need, build: build3, name: build1.name) }
+
+      before do
+        statuses.update_all(scheduling_type: nil)
+      end
+
+      it 'populates scheduling_type of processables' do
+        service.execute(pipeline)
+
+        expect(build1.reload.scheduling_type).to eq('stage')
+        expect(build2.reload.scheduling_type).to eq('stage')
+        expect(build3.reload.scheduling_type).to eq('dag')
+      end
+    end
   end
 
   context 'when user is not allowed to retry pipeline' do