Merge branch '357094-fix-capacity-phase2' into 'master'

Update Migration::EnqueuerWorker capacity logic

See merge request gitlab-org/gitlab!83901

app/workers/container_registry/migration/enqueuer_worker.rb

@@ -76,7 +76,7 @@ module ContainerRegistry
       end
 
       def below_capacity?
-        current_capacity <= maximum_capacity
+        current_capacity < maximum_capacity
       end
 
       def waiting_time_passed?
@@ -111,11 +111,9 @@ module ContainerRegistry
       end
 
       def current_capacity
-        strong_memoize(:current_capacity) do
-          ContainerRepository.with_migration_states(
-            %w[pre_importing pre_import_done importing]
-          ).count
-        end
+        ContainerRepository.with_migration_states(
+          %w[pre_importing pre_import_done importing]
+        ).count
       end
 
       def maximum_capacity
@@ -145,7 +143,7 @@ module ContainerRegistry
       end
 
       def re_enqueue_if_capacity
-        return unless current_capacity < maximum_capacity
+        return unless below_capacity?
 
         self.class.perform_async
       end

spec/workers/container_registry/migration/enqueuer_worker_spec.rb

@@ -6,6 +6,8 @@ RSpec.describe ContainerRegistry::Migration::EnqueuerWorker, :aggregate_failures
   include ExclusiveLeaseHelpers
 
   let_it_be_with_reload(:container_repository) { create(:container_repository, created_at: 2.days.ago) }
+  let_it_be(:importing_repository) { create(:container_repository, :importing) }
+  let_it_be(:pre_importing_repository) { create(:container_repository, :pre_importing) }
 
   let(:worker) { described_class.new }
 
@@ -26,10 +28,10 @@ RSpec.describe ContainerRegistry::Migration::EnqueuerWorker, :aggregate_failures
       end
     end
 
-    shared_examples 're-enqueuing based on capacity' do
+    shared_examples 're-enqueuing based on capacity' do |capacity_limit: 4|
       context 'below capacity' do
         before do
-          allow(ContainerRegistry::Migration).to receive(:capacity).and_return(9999)
+          allow(ContainerRegistry::Migration).to receive(:capacity).and_return(capacity_limit)
         end
 
         it 're-enqueues the worker' do
@@ -53,14 +55,16 @@ RSpec.describe ContainerRegistry::Migration::EnqueuerWorker, :aggregate_failures
     end
 
     context 'with qualified repository' do
-      it 'starts the pre-import for the next qualified repository' do
+      before do
         method = worker.method(:next_repository)
         allow(worker).to receive(:next_repository) do
           next_qualified_repository = method.call
           allow(next_qualified_repository).to receive(:migration_pre_import).and_return(:ok)
           next_qualified_repository
         end
+      end
 
+      it 'starts the pre-import for the next qualified repository' do
         expect_log_extra_metadata(
           import_type: 'next',
           container_repository_id: container_repository.id,
@@ -73,6 +77,24 @@ RSpec.describe ContainerRegistry::Migration::EnqueuerWorker, :aggregate_failures
         expect(container_repository.reload).to be_pre_importing
       end
 
+      context 'when the new pre-import maxes out the capacity' do
+        before do
+          # set capacity to 10
+          stub_feature_flags(
+            container_registry_migration_phase2_capacity_25: false
+          )
+
+          # Plus 2 created above gives 9 importing repositories
+          create_list(:container_repository, 7, :importing)
+        end
+
+        it 'does not re-enqueue the worker' do
+          expect(described_class).not_to receive(:perform_async)
+
+          subject
+        end
+      end
+
       it_behaves_like 're-enqueuing based on capacity'
     end
 
@@ -124,29 +146,33 @@ RSpec.describe ContainerRegistry::Migration::EnqueuerWorker, :aggregate_failures
     context 'when an aborted import is available' do
       let_it_be(:aborted_repository) { create(:container_repository, :import_aborted) }
 
-      it 'retries the import for the aborted repository' do
-        method = worker.method(:next_aborted_repository)
-        allow(worker).to receive(:next_aborted_repository) do
-          next_aborted_repository = method.call
-          allow(next_aborted_repository).to receive(:migration_import).and_return(:ok)
-          allow(next_aborted_repository.gitlab_api_client).to receive(:import_status).and_return('import_failed')
-          next_aborted_repository
+      context 'with a successful registry request' do
+        before do
+          method = worker.method(:next_aborted_repository)
+          allow(worker).to receive(:next_aborted_repository) do
+            next_aborted_repository = method.call
+            allow(next_aborted_repository).to receive(:migration_import).and_return(:ok)
+            allow(next_aborted_repository.gitlab_api_client).to receive(:import_status).and_return('import_failed')
+            next_aborted_repository
+          end
         end
 
-        expect_log_extra_metadata(
-          import_type: 'retry',
-          container_repository_id: aborted_repository.id,
-          container_repository_path: aborted_repository.path,
-          container_repository_migration_state: 'importing'
-        )
+        it 'retries the import for the aborted repository' do
+          expect_log_extra_metadata(
+            import_type: 'retry',
+            container_repository_id: aborted_repository.id,
+            container_repository_path: aborted_repository.path,
+            container_repository_migration_state: 'importing'
+          )
 
-        subject
+          subject
 
-        expect(aborted_repository.reload).to be_importing
-        expect(container_repository.reload).to be_default
-      end
+          expect(aborted_repository.reload).to be_importing
+          expect(container_repository.reload).to be_default
+        end
 
-      it_behaves_like 're-enqueuing based on capacity'
+        it_behaves_like 're-enqueuing based on capacity'
+      end
 
       context 'when an error occurs' do
         it 'does not abort that migration' do
@@ -204,7 +230,7 @@ RSpec.describe ContainerRegistry::Migration::EnqueuerWorker, :aggregate_failures
         expect(container_repository.migration_skipped_at).not_to be_nil
       end
 
-      it_behaves_like 're-enqueuing based on capacity'
+      it_behaves_like 're-enqueuing based on capacity', capacity_limit: 3
     end
 
     context 'when an error occurs' do