1 files changed, 51 insertions, 4 deletions
diff --git a/app/models/merge_request/cleanup_schedule.rb b/app/models/merge_request/cleanup_schedule.rb
index 79817269be2..35194b2b318 100644
--- a/app/models/merge_request/cleanup_schedule.rb
+++ b/app/models/merge_request/cleanup_schedule.rb
@@ -1,14 +1,61 @@
 # frozen_string_literal: true
 
 class MergeRequest::CleanupSchedule < ApplicationRecord
+  STATUSES = {
+    unstarted: 0,
+    running: 1,
+    completed: 2,
+    failed: 3
+  }.freeze
+
   belongs_to :merge_request, inverse_of: :cleanup_schedule
 
   validates :scheduled_at, presence: true
 
-  def self.scheduled_merge_request_ids(limit)
-    where('completed_at IS NULL AND scheduled_at <= NOW()')
+  state_machine :status, initial: :unstarted do
+    state :unstarted, value: STATUSES[:unstarted]
+    state :running, value: STATUSES[:running]
+    state :completed, value: STATUSES[:completed]
+    state :failed, value: STATUSES[:failed]
+
+    event :run do
+      transition unstarted: :running
+    end
+
+    event :retry do
+      transition running: :unstarted
+    end
+
+    event :complete do
+      transition running: :completed
+    end
+
+    event :mark_as_failed do
+      transition running: :failed
+    end
+
+    before_transition to: [:completed] do |cleanup_schedule, _transition|
+      cleanup_schedule.completed_at = Time.current
+    end
+
+    before_transition from: :running, to: [:unstarted, :failed] do |cleanup_schedule, _transition|
+      cleanup_schedule.failed_count += 1
+    end
+  end
+
+  scope :scheduled_and_unstarted, -> {
+    where('completed_at IS NULL AND scheduled_at <= NOW() AND status = ?', STATUSES[:unstarted])
       .order('scheduled_at DESC')
-      .limit(limit)
-      .pluck(:merge_request_id)
+  }
+
+  def self.start_next
+    MergeRequest::CleanupSchedule.transaction do
+      cleanup_schedule = scheduled_and_unstarted.lock('FOR UPDATE SKIP LOCKED').first
+
+      next if cleanup_schedule.blank?
+
+      cleanup_schedule.run!
+      cleanup_schedule
+    end
   end
 end