Container repository migration_state machine

Add a state_machine for managing registry
imports.

app/models/container_repository.rb

@@ -8,12 +8,16 @@ class ContainerRepository < ApplicationRecord
 
   WAITING_CLEANUP_STATUSES = %i[cleanup_scheduled cleanup_unfinished].freeze
   REQUIRING_CLEANUP_STATUSES = %i[cleanup_unscheduled cleanup_scheduled].freeze
+  IDLE_MIGRATION_STATES = %w[default pre_import_done import_done import_aborted import_skipped].freeze
+  ACTIVE_MIGRATION_STATES = %w[pre_importing importing].freeze
+  MIGRATION_STATES = (IDLE_MIGRATION_STATES + ACTIVE_MIGRATION_STATES).freeze
 
   belongs_to :project
 
   validates :name, length: { minimum: 0, allow_nil: false }
   validates :name, uniqueness: { scope: :project_id }
-  validates :migration_state, presence: true
+  validates :migration_state, presence: true, inclusion: { in: MIGRATION_STATES }
+  validates :migration_aborted_in_state, inclusion: { in: ACTIVE_MIGRATION_STATES }, allow_nil: true
 
   validates :migration_retries_count, presence: true,
                                       numericality: { greater_than_or_equal_to: 0 },
@@ -41,6 +45,119 @@ class ContainerRepository < ApplicationRecord
   scope :expiration_policy_started_at_nil_or_before, ->(timestamp) { where('expiration_policy_started_at < ? OR expiration_policy_started_at IS NULL', timestamp) }
   scope :with_stale_ongoing_cleanup, ->(threshold) { cleanup_ongoing.where('expiration_policy_started_at < ?', threshold) }
 
+  state_machine :migration_state, initial: :default do
+    state :pre_importing do
+      validates :migration_pre_import_started_at, presence: true
+      validates :migration_pre_import_done_at, presence: false
+    end
+
+    state :pre_import_done do
+      validates :migration_pre_import_started_at,
+                :migration_pre_import_done_at,
+                presence: true
+    end
+
+    state :importing do
+      validates :migration_import_started_at, presence: true
+      validates :migration_import_done_at, presence: false
+    end
+
+    state :import_done
+
+    state :import_skipped do
+      validates :migration_skipped_reason,
+                :migration_skipped_at,
+                presence: true
+    end
+
+    state :import_aborted do
+      validates :migration_aborted_at, presence: true
+      validates :migration_retries_count, presence: true, numericality: { greater_than_or_equal_to: 1 }
+    end
+
+    event :start_pre_import do
+      transition default: :pre_importing
+    end
+
+    event :finish_pre_import do
+      transition pre_importing: :pre_import_done
+    end
+
+    event :start_import do
+      transition pre_import_done: :importing
+    end
+
+    event :finish_import do
+      transition importing: :import_done
+    end
+
+    event :already_migrated do
+      transition default: :import_done
+    end
+
+    event :abort_import do
+      transition %i[pre_importing importing] => :import_aborted
+    end
+
+    event :skip_import do
+      transition %i[default pre_importing importing] => :import_skipped
+    end
+
+    event :retry_pre_import do
+      transition import_aborted: :pre_importing
+    end
+
+    event :retry_import do
+      transition import_aborted: :importing
+    end
+
+    before_transition any => :pre_importing do |container_repository|
+      container_repository.migration_pre_import_started_at = Time.zone.now
+      container_repository.migration_pre_import_done_at = nil
+    end
+
+    after_transition any => :pre_importing do |container_repository|
+      container_repository.abort_import unless container_repository.migration_pre_import == :ok
+    end
+
+    before_transition pre_importing: :pre_import_done do |container_repository|
+      container_repository.migration_pre_import_done_at = Time.zone.now
+    end
+
+    before_transition any => :importing do |container_repository|
+      container_repository.migration_import_started_at = Time.zone.now
+      container_repository.migration_import_done_at = nil
+    end
+
+    after_transition any => :importing do |container_repository|
+      container_repository.abort_import unless container_repository.migration_import == :ok
+    end
+
+    before_transition importing: :import_done do |container_repository|
+      container_repository.migration_import_done_at = Time.zone.now
+    end
+
+    before_transition any => :import_aborted do |container_repository|
+      container_repository.migration_aborted_in_state = container_repository.migration_state
+      container_repository.migration_aborted_at = Time.zone.now
+      container_repository.migration_retries_count += 1
+    end
+
+    before_transition import_aborted: any do |container_repository|
+      container_repository.migration_aborted_at = nil
+      container_repository.migration_aborted_in_state = nil
+    end
+
+    before_transition any => :import_skipped do |container_repository|
+      container_repository.migration_skipped_at = Time.zone.now
+    end
+
+    before_transition any => %i[import_done import_aborted] do
+      # EnqueuerJob.enqueue perform_async or perform_in depending on the speed FF
+      # To be implemented in https://gitlab.com/gitlab-org/gitlab/-/issues/349744
+    end
+  end
+
   def self.exists_by_path?(path)
     where(
       project: path.repository_project,
@@ -64,6 +181,30 @@ class ContainerRepository < ApplicationRecord
     with_enabled_policy.cleanup_unfinished
   end
 
+  def skip_import(reason:)
+    self.migration_skipped_reason = reason
+
+    super
+  end
+
+  def start_pre_import
+    return false unless ContainerRegistry::Migration.enabled?
+
+    super
+  end
+
+  def retry_pre_import
+    return false unless ContainerRegistry::Migration.enabled?
+
+    super
+  end
+
+  def retry_import
+    return false unless ContainerRegistry::Migration.enabled?
+
+    super
+  end
+
   # rubocop: disable CodeReuse/ServiceClass
   def registry
     @registry ||= begin

spec/factories/container_repositories.rb

@@ -33,6 +33,52 @@ FactoryBot.define do
       expiration_policy_cleanup_status { :cleanup_ongoing }
     end
 
+    trait :default do
+      migration_state { 'default' }
+    end
+
+    trait :pre_importing do
+      migration_state { 'pre_importing' }
+      migration_pre_import_started_at { Time.zone.now }
+    end
+
+    trait :pre_import_done do
+      migration_state { 'pre_import_done' }
+      migration_pre_import_started_at { Time.zone.now }
+      migration_pre_import_done_at { Time.zone.now }
+    end
+
+    trait :importing do
+      migration_state { 'importing' }
+      migration_pre_import_started_at { Time.zone.now }
+      migration_pre_import_done_at { Time.zone.now }
+      migration_import_started_at { Time.zone.now }
+    end
+
+    trait :import_done do
+      migration_state { 'import_done' }
+      migration_pre_import_started_at { Time.zone.now }
+      migration_pre_import_done_at { Time.zone.now }
+      migration_import_started_at { Time.zone.now }
+      migration_import_done_at { Time.zone.now }
+    end
+
+    trait :import_aborted do
+      migration_state { 'import_aborted' }
+      migration_pre_import_started_at { Time.zone.now }
+      migration_pre_import_done_at { Time.zone.now }
+      migration_import_started_at { Time.zone.now }
+      migration_aborted_at { Time.zone.now }
+      migration_aborted_in_state { 'importing' }
+      migration_retries_count { 1 }
+    end
+
+    trait :import_skipped do
+      migration_state { 'import_skipped' }
+      migration_skipped_at { Time.zone.now }
+      migration_skipped_reason { :too_many_tags }
+    end
+
     after(:build) do |repository, evaluator|
       next if evaluator.tags.to_a.none?
 

spec/support/shared_examples/services/container_registry_auth_service_shared_examples.rb

@@ -1142,7 +1142,7 @@ RSpec.shared_examples 'a container registry auth service' do
   end
 
   context 'when importing' do
-    let_it_be(:container_repository) { create(:container_repository, :root, migration_state: 'importing') }
+    let_it_be(:container_repository) { create(:container_repository, :root, :importing) }
     let_it_be(:current_project) { container_repository.project }
     let_it_be(:current_user) { create(:user) }