Registry import observer worker

Worker to log counts of container repositories
in different migration states.

Changelog: added

app/workers/all_queues.yml

@@ -282,6 +282,15 @@
   :weight: 1
   :idempotent: true
   :tags: []
+- :name: cronjob:container_registry_migration_observer
+  :worker_name: ContainerRegistry::Migration::ObserverWorker
+  :feature_category: :container_registry
+  :has_external_dependencies:
+  :urgency: :low
+  :resource_boundary: :unknown
+  :weight: 1
+  :idempotent: true
+  :tags: []
 - :name: cronjob:database_batched_background_migration
   :worker_name: Database::BatchedBackgroundMigrationWorker
   :feature_category: :database

app/workers/container_registry/migration/observer_worker.rb

+# frozen_string_literal: true
+
+module ContainerRegistry
+  module Migration
+    class ObserverWorker
+      include ApplicationWorker
+      # This worker does not perform work scoped to a context
+      include CronjobQueue # rubocop:disable Scalability/CronWorkerContext
+
+      COUNT_BATCH_SIZE = 50000
+
+      data_consistency :sticky
+      feature_category :container_registry
+      urgency :low
+      deduplicate :until_executed, including_scheduled: true
+      idempotent!
+
+      def perform
+        return unless ::ContainerRegistry::Migration.enabled?
+
+        use_replica_if_available do
+          ContainerRepository::MIGRATION_STATES.each do |state|
+            relation = ContainerRepository.with_migration_state(state)
+            count = ::Gitlab::Database::BatchCount.batch_count(
+              relation, batch_size: COUNT_BATCH_SIZE
+            )
+            name = "#{state}_count".to_sym
+            log_extra_metadata_on_done(name, count)
+          end
+        end
+      end
+
+      private
+
+      def use_replica_if_available(&block)
+        ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
+      end
+    end
+  end
+end

config/initializers/1_settings.rb

@@ -542,6 +542,9 @@ Settings.cron_jobs['container_expiration_policy_worker']['job_class'] = 'Contain
 Settings.cron_jobs['container_registry_migration_guard_worker'] ||= Settingslogic.new({})
 Settings.cron_jobs['container_registry_migration_guard_worker']['cron'] ||= '*/10 * * * *'
 Settings.cron_jobs['container_registry_migration_guard_worker']['job_class'] = 'ContainerRegistry::Migration::GuardWorker'
+Settings.cron_jobs['container_registry_migration_observer_worker'] ||= Settingslogic.new({})
+Settings.cron_jobs['container_registry_migration_observer_worker']['cron'] ||= '*/30 * * * *'
+Settings.cron_jobs['container_registry_migration_observer_worker']['job_class'] = 'ContainerRegistry::Migration::ObserverWorker'
 Settings.cron_jobs['image_ttl_group_policy_worker'] ||= Settingslogic.new({})
 Settings.cron_jobs['image_ttl_group_policy_worker']['cron'] ||= '40 0 * * *'
 Settings.cron_jobs['image_ttl_group_policy_worker']['job_class'] = 'DependencyProxy::ImageTtlGroupPolicyWorker'

db/migrate/20220202034409_add_tmp_index_on_id_and_migration_state_to_container_repositories.rb

+# frozen_string_literal: true
+
+class AddTmpIndexOnIdAndMigrationStateToContainerRepositories < Gitlab::Database::Migration[1.0]
+  INDEX_NAME = 'tmp_index_container_repositories_on_id_migration_state'
+
+  disable_ddl_transaction!
+
+  # Temporary index to be removed https://gitlab.com/gitlab-org/gitlab/-/issues/351783
+  def up
+    add_concurrent_index :container_repositories, [:id, :migration_state], name: INDEX_NAME
+  end
+
+  def down
+    remove_concurrent_index_by_name :container_repositories, INDEX_NAME
+  end
+end

db/schema_migrations/20220202034409

+0efe482aa626cf80912feaa1176837253b094fc434f273bee35b5fe3e8ce4243
\ No newline at end of file

db/structure.sql

@@ -28317,6 +28317,8 @@ CREATE INDEX tmp_idx_deduplicate_vulnerability_occurrences ON vulnerability_occu
 
 CREATE INDEX tmp_idx_vulnerability_occurrences_on_id_where_report_type_7_99 ON vulnerability_occurrences USING btree (id) WHERE (report_type = ANY (ARRAY[7, 99]));
 
+CREATE INDEX tmp_index_container_repositories_on_id_migration_state ON container_repositories USING btree (id, migration_state);
+
 CREATE INDEX tmp_index_for_namespace_id_migration_on_routes ON routes USING btree (id) WHERE ((namespace_id IS NULL) AND ((source_type)::text = 'Namespace'::text));
 
 CREATE INDEX tmp_index_members_on_state ON members USING btree (state) WHERE (state = 2);

spec/workers/container_registry/migration/observer_worker_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe ContainerRegistry::Migration::ObserverWorker, :aggregate_failures do
+  let(:worker) { described_class.new }
+
+  describe '#perform' do
+    subject { worker.perform }
+
+    context 'when the migration feature flag is disabled' do
+      before do
+        stub_feature_flags(container_registry_migration_phase2_enabled: false)
+      end
+
+      it 'does nothing' do
+        expect(worker).not_to receive(:log_extra_metadata_on_done)
+
+        subject
+      end
+    end
+
+    context 'when the migration is enabled' do
+      before do
+        create_list(:container_repository, 3)
+        create(:container_repository, :pre_importing)
+        create(:container_repository, :pre_import_done)
+        create_list(:container_repository, 2, :importing)
+        create(:container_repository, :import_aborted)
+        # batch_count is not allowed within a transaction but
+        # all rspec tests run inside of a transaction.
+        # This mocks the false positive.
+        allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false) # rubocop:disable Database/MultipleDatabases
+      end
+
+      it 'logs all the counts' do
+        expect(worker).to receive(:log_extra_metadata_on_done).with(:default_count, 3)
+        expect(worker).to receive(:log_extra_metadata_on_done).with(:pre_importing_count, 1)
+        expect(worker).to receive(:log_extra_metadata_on_done).with(:pre_import_done_count, 1)
+        expect(worker).to receive(:log_extra_metadata_on_done).with(:importing_count, 2)
+        expect(worker).to receive(:log_extra_metadata_on_done).with(:import_done_count, 0)
+        expect(worker).to receive(:log_extra_metadata_on_done).with(:import_aborted_count, 1)
+        expect(worker).to receive(:log_extra_metadata_on_done).with(:import_skipped_count, 0)
+
+        subject
+      end
+
+      context 'with load balancing enabled', :db_load_balancing do
+        it 'uses the replica' do
+          expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
+
+          subject
+        end
+      end
+    end
+  end
+end