1 files changed, 49 insertions, 7 deletions
diff --git a/app/workers/merge_request_cleanup_refs_worker.rb b/app/workers/merge_request_cleanup_refs_worker.rb
index 162c6dc2a88..408d070d56f 100644
--- a/app/workers/merge_request_cleanup_refs_worker.rb
+++ b/app/workers/merge_request_cleanup_refs_worker.rb
@@ -2,6 +2,8 @@
 
 class MergeRequestCleanupRefsWorker
   include ApplicationWorker
+  include LimitedCapacity::Worker
+  include Gitlab::Utils::StrongMemoize
 
   sidekiq_options retry: 3
 
@@ -9,20 +11,60 @@ class MergeRequestCleanupRefsWorker
   tags :exclude_from_kubernetes
   idempotent!
 
-  def perform(merge_request_id)
-    return unless Feature.enabled?(:merge_request_refs_cleanup, default_enabled: false)
+  # Hard-coded to 4 for now. Will be configurable later on via application settings.
+  # This means, there can only be 4 jobs running at the same time at maximum.
+  MAX_RUNNING_JOBS = 4
+  FAILURE_THRESHOLD = 3
 
-    merge_request = MergeRequest.find_by_id(merge_request_id)
+  def perform_work
+    return unless Feature.enabled?(:merge_request_refs_cleanup, default_enabled: false)
 
     unless merge_request
-      logger.error("Failed to find merge request with ID: #{merge_request_id}")
+      logger.error('No existing merge request to be cleaned up.')
       return
     end
 
-    result = ::MergeRequests::CleanupRefsService.new(merge_request).execute
+    log_extra_metadata_on_done(:merge_request_id, merge_request.id)
+
+    result = MergeRequests::CleanupRefsService.new(merge_request).execute
+
+    if result[:status] == :success
+      merge_request_cleanup_schedule.complete!
+    else
+      if merge_request_cleanup_schedule.failed_count < FAILURE_THRESHOLD
+        merge_request_cleanup_schedule.retry!
+      else
+        merge_request_cleanup_schedule.mark_as_failed!
+      end
+
+      log_extra_metadata_on_done(:message, result[:message])
+    end
+
+    log_extra_metadata_on_done(:status, merge_request_cleanup_schedule.status)
+  end
+
+  def remaining_work_count
+    MergeRequest::CleanupSchedule
+      .scheduled_and_unstarted
+      .limit(max_running_jobs)
+      .count
+  end
+
+  def max_running_jobs
+    MAX_RUNNING_JOBS
+  end
+
+  private
 
-    return if result[:status] == :success
+  def merge_request
+    strong_memoize(:merge_request) do
+      merge_request_cleanup_schedule&.merge_request
+    end
+  end
 
-    logger.error("Failed cleanup refs of merge request (#{merge_request_id}): #{result[:message]}")
+  def merge_request_cleanup_schedule
+    strong_memoize(:merge_request_cleanup_schedule) do
+      MergeRequest::CleanupSchedule.start_next
+    end
   end
 end