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/models/container_repository_spec.rb

@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe ContainerRepository do
+RSpec.describe ContainerRepository, :aggregate_failures do
   using RSpec::Parameterized::TableSyntax
 
   let(:group) { create(:group, name: 'group') }
@@ -36,7 +36,282 @@ RSpec.describe ContainerRepository do
   describe 'validations' do
     it { is_expected.to validate_presence_of(:migration_retries_count) }
     it { is_expected.to validate_numericality_of(:migration_retries_count).is_greater_than_or_equal_to(0) }
-    it { is_expected.to validate_presence_of(:migration_state) }
+
+    it { is_expected.to validate_inclusion_of(:migration_aborted_in_state).in_array(ContainerRepository::ACTIVE_MIGRATION_STATES) }
+    it { is_expected.to allow_value(nil).for(:migration_aborted_in_state) }
+
+    context 'migration_state' do
+      it { is_expected.to validate_presence_of(:migration_state) }
+      it { is_expected.to validate_inclusion_of(:migration_state).in_array(ContainerRepository::MIGRATION_STATES) }
+
+      describe 'pre_importing' do
+        it 'validates expected attributes' do
+          expect(build(:container_repository, migration_state: 'pre_importing')).to be_invalid
+          expect(build(:container_repository, :pre_importing)).to be_valid
+        end
+      end
+
+      describe 'pre_import_done' do
+        it 'validates expected attributes' do
+          expect(build(:container_repository, migration_state: 'pre_import_done')).to be_invalid
+          expect(build(:container_repository, :pre_import_done)).to be_valid
+        end
+      end
+
+      describe 'importing' do
+        it 'validates expected attributes' do
+          expect(build(:container_repository, migration_state: 'importing')).to be_invalid
+          expect(build(:container_repository, :importing)).to be_valid
+        end
+      end
+
+      describe 'import_skipped' do
+        it 'validates expected attributes' do
+          expect(build(:container_repository, migration_state: 'import_skipped')).to be_invalid
+          expect(build(:container_repository, :import_skipped)).to be_valid
+        end
+      end
+
+      describe 'import_aborted' do
+        it 'validates expected attributes' do
+          expect(build(:container_repository, migration_state: 'import_aborted')).to be_invalid
+          expect(build(:container_repository, :import_aborted)).to be_valid
+        end
+      end
+    end
+  end
+
+  context ':migration_state state_machine' do
+    shared_examples 'no action when feature flag is disabled' do
+      context 'feature flag disabled' do
+        before do
+          stub_feature_flags(container_registry_migration_phase2_enabled: false)
+        end
+
+        it { is_expected.to eq(false) }
+      end
+    end
+
+    shared_examples 'transitioning to pre_importing', skip_pre_import_success: true do
+      before do
+        repository.update_column(:migration_pre_import_done_at, Time.zone.now)
+      end
+
+      it_behaves_like 'no action when feature flag is disabled'
+
+      context 'successful pre_import request' do
+        it 'sets migration_pre_import_started_at and resets migration_pre_import_done_at' do
+          expect(repository).to receive(:migration_pre_import).and_return(:ok)
+
+          expect { subject }.to change { repository.reload.migration_pre_import_started_at }
+            .and change { repository.migration_pre_import_done_at }.to(nil)
+
+          expect(repository).to be_pre_importing
+        end
+      end
+
+      context 'failed pre_import request' do
+        it 'sets migration_pre_import_started_at and resets migration_pre_import_done_at' do
+          expect(repository).to receive(:migration_pre_import).and_return(:error)
+
+          expect { subject }.to change { repository.reload.migration_pre_import_started_at }
+            .and change { repository.migration_aborted_at }
+            .and change { repository.migration_pre_import_done_at }.to(nil)
+
+          expect(repository.migration_aborted_in_state).to eq('pre_importing')
+          expect(repository).to be_import_aborted
+        end
+      end
+    end
+
+    shared_examples 'transitioning to importing', skip_import_success: true do
+      before do
+        repository.update_columns(migration_import_done_at: Time.zone.now)
+      end
+
+      context 'successful import request' do
+        it 'sets migration_import_started_at and resets migration_import_done_at' do
+          expect(repository).to receive(:migration_import).and_return(:ok)
+
+          expect { subject }.to change { repository.reload.migration_import_started_at }
+            .and change { repository.migration_import_done_at }.to(nil)
+
+          expect(repository).to be_importing
+        end
+      end
+
+      context 'failed import request' do
+        it 'sets migration_import_started_at and resets migration_import_done_at' do
+          expect(repository).to receive(:migration_import).and_return(:error)
+
+          expect { subject }.to change { repository.reload.migration_import_started_at }
+            .and change { repository.migration_aborted_at }
+
+          expect(repository.migration_aborted_in_state).to eq('importing')
+          expect(repository).to be_import_aborted
+        end
+      end
+    end
+
+    shared_examples 'transitioning out of import_aborted' do
+      it 'resets migration_aborted_at and migration_aborted_in_state' do
+        expect { subject }.to change { repository.reload.migration_aborted_in_state }.to(nil)
+          .and change { repository.migration_aborted_at }.to(nil)
+      end
+    end
+
+    shared_examples 'transitioning from allowed states' do |allowed_states|
+      ContainerRepository::MIGRATION_STATES.each do |state|
+        result = allowed_states.include?(state)
+
+        context "when transitioning from #{state}" do
+          let(:repository) { create(:container_repository, state.to_sym) }
+
+          it "returns #{result}" do
+            expect(subject).to eq(result)
+          end
+        end
+      end
+    end
+
+    describe '#start_pre_import' do
+      let_it_be_with_reload(:repository) { create(:container_repository) }
+
+      subject { repository.start_pre_import }
+
+      before do |example|
+        unless example.metadata[:skip_pre_import_success]
+          allow(repository).to receive(:migration_pre_import).and_return(:ok)
+        end
+      end
+
+      it_behaves_like 'transitioning from allowed states', %w[default]
+      it_behaves_like 'transitioning to pre_importing'
+    end
+
+    describe '#retry_pre_import' do
+      let_it_be_with_reload(:repository) { create(:container_repository, :import_aborted) }
+
+      subject { repository.retry_pre_import }
+
+      before do |example|
+        unless example.metadata[:skip_pre_import_success]
+          allow(repository).to receive(:migration_pre_import).and_return(:ok)
+        end
+      end
+
+      it_behaves_like 'transitioning from allowed states', %w[import_aborted]
+      it_behaves_like 'transitioning to pre_importing'
+      it_behaves_like 'transitioning out of import_aborted'
+    end
+
+    describe '#finish_pre_import' do
+      let_it_be_with_reload(:repository) { create(:container_repository, :pre_importing) }
+
+      subject { repository.finish_pre_import }
+
+      it_behaves_like 'transitioning from allowed states', %w[pre_importing]
+
+      it 'sets migration_pre_import_done_at' do
+        expect { subject }.to change { repository.reload.migration_pre_import_done_at }
+
+        expect(repository).to be_pre_import_done
+      end
+    end
+
+    describe '#start_import' do
+      let_it_be_with_reload(:repository) { create(:container_repository, :pre_import_done) }
+
+      subject { repository.start_import }
+
+      before do |example|
+        unless example.metadata[:skip_import_success]
+          allow(repository).to receive(:migration_import).and_return(:ok)
+        end
+      end
+
+      it_behaves_like 'transitioning from allowed states', %w[pre_import_done]
+      it_behaves_like 'transitioning to importing'
+    end
+
+    describe '#retry_import' do
+      let_it_be_with_reload(:repository) { create(:container_repository, :import_aborted) }
+
+      subject { repository.retry_import }
+
+      before do |example|
+        unless example.metadata[:skip_import_success]
+          allow(repository).to receive(:migration_import).and_return(:ok)
+        end
+      end
+
+      it_behaves_like 'transitioning from allowed states', %w[import_aborted]
+      it_behaves_like 'transitioning to importing'
+      it_behaves_like 'no action when feature flag is disabled'
+    end
+
+    describe '#finish_import' do
+      let_it_be_with_reload(:repository) { create(:container_repository, :importing) }
+
+      subject { repository.finish_import }
+
+      it_behaves_like 'transitioning from allowed states', %w[importing]
+
+      it 'sets migration_import_done_at' do
+        expect { subject }.to change { repository.reload.migration_import_done_at }
+
+        expect(repository).to be_import_done
+      end
+    end
+
+    describe '#already_migrated' do
+      let_it_be_with_reload(:repository) { create(:container_repository) }
+
+      subject { repository.already_migrated }
+
+      it_behaves_like 'transitioning from allowed states', %w[default]
+
+      it 'sets migration_import_done_at' do
+        subject
+
+        expect(repository).to be_import_done
+      end
+    end
+
+    describe '#abort_import' do
+      let_it_be_with_reload(:repository) { create(:container_repository, :importing) }
+
+      subject { repository.abort_import }
+
+      it_behaves_like 'transitioning from allowed states', %w[pre_importing importing]
+
+      it 'sets migration_aborted_at and migration_aborted_at and increments the retry count' do
+        expect { subject }.to change { repository.migration_aborted_at }
+          .and change { repository.reload.migration_retries_count }.by(1)
+
+        expect(repository.migration_aborted_in_state).to eq('importing')
+        expect(repository).to be_import_aborted
+      end
+    end
+
+    describe '#skip_import' do
+      let_it_be_with_reload(:repository) { create(:container_repository) }
+
+      subject { repository.skip_import(reason: :too_many_retries) }
+
+      it_behaves_like 'transitioning from allowed states', %w[default pre_importing importing]
+
+      it 'sets migration_skipped_at and migration_skipped_reason' do
+        expect { subject }.to change { repository.reload.migration_skipped_at }
+
+        expect(repository.migration_skipped_reason).to eq('too_many_retries')
+        expect(repository).to be_import_skipped
+      end
+
+      it 'raises and error if a reason is not given' do
+        expect { repository.skip_import }.to raise_error(ArgumentError)
+      end
+    end
   end
 
   describe '#tag' do
@@ -545,20 +820,14 @@ RSpec.describe ContainerRepository do
   end
 
   describe '#migration_importing?' do
-    let_it_be_with_reload(:container_repository) { create(:container_repository, migration_state: 'importing')}
-
     subject { container_repository.migration_importing? }
 
-    it { is_expected.to eq(true) }
+    ContainerRepository::MIGRATION_STATES.each do |state|
+      context "when in #{state} migration_state" do
+        let(:container_repository) { create(:container_repository, state.to_sym)}
 
-    context 'when not importing' do
-      before do
-        # Technical debt: when the state machine is added, we should iterate through all possible states
-        # https://gitlab.com/gitlab-org/gitlab/-/merge_requests/78499
-        container_repository.update_column(:migration_state, 'default')
+        it { is_expected.to eq(state == 'importing') }
       end
-
-      it { is_expected.to eq(false) }
     end
   end
 

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) }