Add latest changes from gitlab-org/gitlab@master

13 files changed, 167 insertions, 20 deletions

diff --git a/app/models/merge_request.rb b/app/models/merge_request.rb
index d110b55b3e1..2b4d0949a92 100644
--- a/app/models/merge_request.rb
+++ b/app/models/merge_request.rb
@@ -1528,6 +1528,14 @@ class MergeRequest < ApplicationRecord
     "refs/#{Repository::REF_MERGE_REQUEST}/#{iid}/train"
   end
 
+  def schedule_cleanup_refs(only: :all)
+    if Feature.enabled?(:merge_request_cleanup_ref_worker_async, target_project)
+      MergeRequests::CleanupRefWorker.perform_async(id, only.to_s)
+    else
+      cleanup_refs(only: only)
+    end
+  end
+
   def cleanup_refs(only: :all)
     target_refs = []
     target_refs << ref_path       if %i[all head].include?(only)
diff --git a/app/services/concerns/projects/remove_refs.rb b/app/services/concerns/projects/remove_refs.rb
new file mode 100644
index 00000000000..d133aa0ced6
--- /dev/null
+++ b/app/services/concerns/projects/remove_refs.rb
@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Projects
+  module RemoveRefs
+    extend ActiveSupport::Concern
+    include Gitlab::ExclusiveLeaseHelpers
+
+    LOCK_RETRY = 3
+    LOCK_TTL = 5.minutes
+    LOCK_SLEEP = 0.5.seconds
+
+    def serialized_remove_refs(project_id, &blk)
+      in_lock("projects/#{project_id}/serialized_remove_refs", **lock_params, &blk)
+    end
+
+    def lock_params
+      {
+        ttl: LOCK_TTL,
+        retries: LOCK_RETRY,
+        sleep_sec: LOCK_SLEEP
+      }
+    end
+  end
+end
diff --git a/app/services/merge_requests/cleanup_refs_service.rb b/app/services/merge_requests/cleanup_refs_service.rb
index 2094ea00160..5081655601b 100644
--- a/app/services/merge_requests/cleanup_refs_service.rb
+++ b/app/services/merge_requests/cleanup_refs_service.rb
@@ -16,7 +16,6 @@ module MergeRequests
       @merge_request = merge_request
       @repository = merge_request.project.repository
       @ref_path = merge_request.ref_path
-      @merge_ref_path = merge_request.merge_ref_path
       @ref_head_sha = @repository.commit(merge_request.ref_path)&.id
       @merge_ref_sha = merge_request.merge_ref_head&.id
     end
@@ -42,7 +41,7 @@ module MergeRequests
 
     private
 
-    attr_reader :repository, :ref_path, :merge_ref_path, :ref_head_sha, :merge_ref_sha
+    attr_reader :repository, :ref_path, :ref_head_sha, :merge_ref_sha
 
     def scheduled?
       merge_request.cleanup_schedule.present? && merge_request.cleanup_schedule.scheduled_at <= Time.current
@@ -79,7 +78,7 @@ module MergeRequests
     end
 
     def delete_refs
-      repository.delete_refs(ref_path, merge_ref_path)
+      merge_request.schedule_cleanup_refs
     end
 
     def update_schedule
diff --git a/app/workers/all_queues.yml b/app/workers/all_queues.yml
index d599e7414c8..316b11ad791 100644
--- a/app/workers/all_queues.yml
+++ b/app/workers/all_queues.yml
@@ -2964,6 +2964,15 @@
   :weight: 1
   :idempotent: true
   :tags: []
+- :name: merge_requests_cleanup_ref
+  :worker_name: MergeRequests::CleanupRefWorker
+  :feature_category: :code_review_workflow
+  :has_external_dependencies: false
+  :urgency: :low
+  :resource_boundary: :unknown
+  :weight: 1
+  :idempotent: true
+  :tags: []
 - :name: merge_requests_close_issue
   :worker_name: MergeRequests::CloseIssueWorker
   :feature_category: :code_review_workflow
diff --git a/app/workers/ci/pipeline_cleanup_ref_worker.rb b/app/workers/ci/pipeline_cleanup_ref_worker.rb
index a367c0183aa..291e1090c18 100644
--- a/app/workers/ci/pipeline_cleanup_ref_worker.rb
+++ b/app/workers/ci/pipeline_cleanup_ref_worker.rb
@@ -3,15 +3,11 @@
 module Ci
   class PipelineCleanupRefWorker
     include ApplicationWorker
-    include Gitlab::ExclusiveLeaseHelpers
+    include Projects::RemoveRefs
 
     sidekiq_options retry: 3
     include PipelineQueue
 
-    LOCK_RETRY = 3
-    LOCK_TTL = 5.minutes
-    LOCK_SLEEP = 0.5.seconds
-
     idempotent!
     deduplicate :until_executed, if_deduplicated: :reschedule_once, ttl: 1.minute
     data_consistency :always # rubocop:disable SidekiqLoadBalancing/WorkerDataConsistency
@@ -31,19 +27,9 @@ module Ci
       return unless pipeline
       return unless pipeline.persistent_ref.should_delete?
 
-      in_lock("#{self.class.name.underscore}/#{pipeline.project_id}", **lock_params) do
+      serialized_remove_refs(pipeline.project_id) do
         pipeline.reset.persistent_ref.delete
       end
     end
-
-    private
-
-    def lock_params
-      {
-        ttl: LOCK_TTL,
-        retries: LOCK_RETRY,
-        sleep_sec: LOCK_SLEEP
-      }
-    end
   end
 end
diff --git a/app/workers/merge_requests/cleanup_ref_worker.rb b/app/workers/merge_requests/cleanup_ref_worker.rb
new file mode 100644
index 00000000000..c714b976a2b
--- /dev/null
+++ b/app/workers/merge_requests/cleanup_ref_worker.rb
@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module MergeRequests
+  class CleanupRefWorker
+    include ApplicationWorker
+    include Projects::RemoveRefs
+
+    sidekiq_options retry: 3
+    loggable_arguments 2
+    feature_category :code_review_workflow
+
+    idempotent!
+    deduplicate :until_executed, if_deduplicated: :reschedule_once, ttl: 1.minute
+    data_consistency :delayed
+
+    urgency :low
+
+    # Even though this worker is de-duplicated we need to acquire lock
+    # on a project to avoid running many concurrent refs removals
+    #
+    # TODO: Once underlying fix is done we can remove `in_lock`
+    #
+    # Related to:
+    # - https://gitlab.com/gitlab-org/gitaly/-/issues/5368
+    # - https://gitlab.com/gitlab-org/gitaly/-/issues/5369
+    def perform(merge_request_id, only)
+      merge_request = MergeRequest.find_by_id(merge_request_id)
+      return unless merge_request
+
+      serialized_remove_refs(merge_request.target_project_id) do
+        merge_request.cleanup_refs(only: only.to_sym)
+      end
+    end
+  end
+end
diff --git a/config/feature_flags/development/merge_request_cleanup_ref_worker_async.yml b/config/feature_flags/development/merge_request_cleanup_ref_worker_async.yml
new file mode 100644
index 00000000000..75d7541fdef
--- /dev/null
+++ b/config/feature_flags/development/merge_request_cleanup_ref_worker_async.yml
@@ -0,0 +1,8 @@
+---
+name: merge_request_cleanup_ref_worker_async
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/125109
+rollout_issue_url: https://gitlab.com/gitlab-org/gitaly/-/issues/5369
+milestone: '16.2'
+type: development
+group: group::tenant scale
+default_enabled: false
diff --git a/config/sidekiq_queues.yml b/config/sidekiq_queues.yml
index 2e643f8e977..35cb020c2e3 100644
--- a/config/sidekiq_queues.yml
+++ b/config/sidekiq_queues.yml
@@ -343,6 +343,8 @@
   - 1
 - - merge_requests_capture_suggested_reviewers_accepted
   - 1
+- - merge_requests_cleanup_ref
+  - 1
 - - merge_requests_close_issue
   - 1
 - - merge_requests_create_approval_event
diff --git a/spec/models/merge_request_spec.rb b/spec/models/merge_request_spec.rb
index cd5460fb4c8..6dd962326d3 100644
--- a/spec/models/merge_request_spec.rb
+++ b/spec/models/merge_request_spec.rb
@@ -5105,6 +5105,32 @@ RSpec.describe MergeRequest, factory_default: :keep, feature_category: :code_rev
     end
   end
 
+  describe '#schedule_cleanup_refs' do
+    subject { merge_request.schedule_cleanup_refs(only: :train) }
+
+    let(:merge_request) { build(:merge_request, source_project: create(:project, :repository)) }
+
+    it 'does schedule MergeRequests::CleanupRefWorker' do
+      expect(MergeRequests::CleanupRefWorker).to receive(:perform_async).with(merge_request.id, 'train')
+
+      subject
+    end
+
+    context 'when merge_request_cleanup_ref_worker_async is disabled' do
+      before do
+        stub_feature_flags(merge_request_cleanup_ref_worker_async: false)
+      end
+
+      it 'deletes all refs from the target project' do
+        expect(merge_request.target_project.repository)
+          .to receive(:delete_refs)
+          .with(merge_request.train_ref_path)
+
+        subject
+      end
+    end
+  end
+
   describe '#cleanup_refs' do
     subject { merge_request.cleanup_refs(only: only) }
 
diff --git a/spec/services/merge_requests/cleanup_refs_service_spec.rb b/spec/services/merge_requests/cleanup_refs_service_spec.rb
index 960b8101c36..efb6265e3d8 100644
--- a/spec/services/merge_requests/cleanup_refs_service_spec.rb
+++ b/spec/services/merge_requests/cleanup_refs_service_spec.rb
@@ -18,6 +18,8 @@ RSpec.describe MergeRequests::CleanupRefsService, feature_category: :code_review
 
   describe '#execute' do
     before do
+      stub_feature_flags(merge_request_cleanup_ref_worker_async: false)
+
       # Need to re-enable this as it's being stubbed in spec_helper for
       # performance reasons but is needed to run for this test.
       allow(Gitlab::Git::KeepAround).to receive(:execute).and_call_original
diff --git a/spec/workers/ci/merge_requests/cleanup_ref_worker_spec.rb b/spec/workers/ci/merge_requests/cleanup_ref_worker_spec.rb
new file mode 100644
index 00000000000..41ee1047adf
--- /dev/null
+++ b/spec/workers/ci/merge_requests/cleanup_ref_worker_spec.rb
@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe MergeRequests::CleanupRefWorker, :sidekiq_inline, feature_category: :code_review_workflow do
+  include ExclusiveLeaseHelpers
+
+  let_it_be(:source_project) { create(:project, :repository) }
+  let_it_be(:merge_request) { create(:merge_request, source_project: source_project) }
+
+  let(:worker) { described_class.new }
+  let(:only) { :all }
+
+  subject { worker.perform(merge_request.id, only) }
+
+  it 'does remove all merge request refs' do
+    expect(MergeRequest).to receive(:find_by_id).with(merge_request.id).and_return(merge_request)
+    expect(merge_request.target_project.repository)
+      .to receive(:delete_refs)
+      .with(merge_request.ref_path, merge_request.merge_ref_path, merge_request.train_ref_path)
+
+    subject
+  end
+
+  context 'when only is :train' do
+    let(:only) { :train }
+
+    it 'does remove only car merge request train ref' do
+      expect(MergeRequest).to receive(:find_by_id).with(merge_request.id).and_return(merge_request)
+      expect(merge_request.target_project.repository)
+        .to receive(:delete_refs)
+        .with(merge_request.train_ref_path)
+
+      subject
+    end
+  end
+
+  context 'when max retry attempts reach' do
+    let(:lease_key) { "projects/#{merge_request.target_project_id}/serialized_remove_refs" }
+    let!(:lease) { stub_exclusive_lease_taken(lease_key) }
+
+    it 'does not raise error' do
+      expect(lease).to receive(:try_obtain).exactly(described_class::LOCK_RETRY + 1).times
+      expect { subject }.to raise_error(Gitlab::ExclusiveLeaseHelpers::FailedToObtainLockError)
+    end
+  end
+end
diff --git a/spec/workers/ci/pipeline_cleanup_ref_worker_spec.rb b/spec/workers/ci/pipeline_cleanup_ref_worker_spec.rb
index 2913412abcd..c2564822bd9 100644
--- a/spec/workers/ci/pipeline_cleanup_ref_worker_spec.rb
+++ b/spec/workers/ci/pipeline_cleanup_ref_worker_spec.rb
@@ -50,7 +50,7 @@ RSpec.describe Ci::PipelineCleanupRefWorker, :sidekiq_inline, feature_category:
   end
 
   context 'when max retry attempts reach' do
-    let(:lease_key) { "#{described_class.name.underscore}/#{pipeline.project_id}" }
+    let(:lease_key) { "projects/#{pipeline.project_id}/serialized_remove_refs" }
     let!(:lease) { stub_exclusive_lease_taken(lease_key) }
 
     it 'does not raise error' do
diff --git a/spec/workers/every_sidekiq_worker_spec.rb b/spec/workers/every_sidekiq_worker_spec.rb
index 2298f2e086f..4df3b4fee8e 100644
--- a/spec/workers/every_sidekiq_worker_spec.rb
+++ b/spec/workers/every_sidekiq_worker_spec.rb
@@ -349,6 +349,7 @@ RSpec.describe 'Every Sidekiq worker', feature_category: :shared do
         'MergeRequestMergeabilityCheckWorker' => 3,
         'MergeRequestResetApprovalsWorker' => 3,
         'MergeRequests::CaptureSuggestedReviewersAcceptedWorker' => 3,
+        'MergeRequests::CleanupRefWorker' => 3,
         'MergeRequests::CreatePipelineWorker' => 3,
         'MergeRequests::DeleteSourceBranchWorker' => 3,
         'MergeRequests::FetchSuggestedReviewersWorker' => 3,