Move cleanup policies metrics to the parent worker

app/workers/container_expiration_policies/cleanup_container_repository_worker.rb

@@ -49,15 +49,11 @@ module ContainerExpirationPolicies
     end
 
     def remaining_work_count
-      total_count = cleanup_scheduled_count + cleanup_unfinished_count
+      count = cleanup_scheduled_count
 
-      log_info(
-        cleanup_scheduled_count: cleanup_scheduled_count,
-        cleanup_unfinished_count: cleanup_unfinished_count,
-        cleanup_total_count: total_count
-      )
+      return count if count > max_running_jobs
 
-      total_count
+      count + cleanup_unfinished_count
     end
 
     private

app/workers/container_expiration_policy_worker.rb

@@ -17,6 +17,7 @@ class ContainerExpirationPolicyWorker # rubocop:disable Scalability/IdempotentWo
     process_stale_ongoing_cleanups
     disable_policies_without_container_repositories
     throttling_enabled? ? perform_throttled : perform_unthrottled
+    log_counts
   end
 
   private
@@ -28,6 +29,26 @@ class ContainerExpirationPolicyWorker # rubocop:disable Scalability/IdempotentWo
     end
   end
 
+  def log_counts
+    use_replica_if_available do
+      required_count = ContainerRepository.requiring_cleanup.count
+      unfinished_count = ContainerRepository.with_unfinished_cleanup.count
+
+      log_extra_metadata_on_done(:cleanup_required_count, required_count)
+      log_extra_metadata_on_done(:cleanup_unfinished_count, unfinished_count)
+      log_extra_metadata_on_done(:cleanup_total_count, required_count + unfinished_count)
+    end
+  end
+
+  # data_consistency :delayed not used as this is a cron job and those jobs are
+  # not perfomed with a delay
+  # https://gitlab.com/gitlab-org/gitlab/-/merge_requests/63635#note_603771207
+  def use_replica_if_available(&blk)
+    return yield unless ::Gitlab::Database::LoadBalancing.enable?
+
+    ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&blk)
+  end
+
   def process_stale_ongoing_cleanups
     threshold = delete_tags_service_timeout.seconds + 30.minutes
     ContainerRepository.with_stale_ongoing_cleanup(threshold.ago)

db/migrate/20210705130919_create_container_repos_on_exp_cleanup_status_project_id_start_date_index.rb

+# frozen_string_literal: true
+
+# See https://docs.gitlab.com/ee/development/migration_style_guide.html
+# for more information on how to write migrations for GitLab.
+
+class CreateContainerReposOnExpCleanupStatusProjectIdStartDateIndex < ActiveRecord::Migration[6.1]
+  include Gitlab::Database::MigrationHelpers
+
+  DOWNTIME = false
+  OLD_INDEX_NAME = 'idx_container_repositories_on_exp_cleanup_status_and_start_date'
+  NEW_INDEX_NAME = 'idx_container_repos_on_exp_cleanup_status_project_id_start_date'
+
+  disable_ddl_transaction!
+
+  def up
+    add_concurrent_index(:container_repositories, [:expiration_policy_cleanup_status, :project_id, :expiration_policy_started_at], name: NEW_INDEX_NAME)
+    remove_concurrent_index(:container_repositories, [:expiration_policy_cleanup_status, :expiration_policy_started_at], name: OLD_INDEX_NAME)
+  end
+
+  def down
+    add_concurrent_index(:container_repositories, [:expiration_policy_cleanup_status, :expiration_policy_started_at], name: OLD_INDEX_NAME)
+    remove_concurrent_index(:container_repositories, [:expiration_policy_cleanup_status, :project_id, :expiration_policy_started_at], name: NEW_INDEX_NAME)
+  end
+end

db/schema_migrations/20210705130919

+c33dd2c63d5a8c6e3c2f49e640b1780734b4bfca88378fac67ea5f5bd24fb2b4
\ No newline at end of file

db/structure.sql

@@ -22481,7 +22481,7 @@ CREATE INDEX idx_container_exp_policies_on_project_id_next_run_at ON container_e
 
 CREATE INDEX idx_container_exp_policies_on_project_id_next_run_at_enabled ON container_expiration_policies USING btree (project_id, next_run_at, enabled);
 
-CREATE INDEX idx_container_repositories_on_exp_cleanup_status_and_start_date ON container_repositories USING btree (expiration_policy_cleanup_status, expiration_policy_started_at);
+CREATE INDEX idx_container_repos_on_exp_cleanup_status_project_id_start_date ON container_repositories USING btree (expiration_policy_cleanup_status, project_id, expiration_policy_started_at);
 
 CREATE INDEX idx_deployment_clusters_on_cluster_id_and_kubernetes_namespace ON deployment_clusters USING btree (cluster_id, kubernetes_namespace);
 

spec/workers/container_expiration_policies/cleanup_container_repository_worker_spec.rb

@@ -413,20 +413,30 @@ RSpec.describe ContainerExpirationPolicies::CleanupContainerRepositoryWorker do
       disabled_repository.project.container_expiration_policy.update_column(:enabled, false)
     end
 
+    context 'counts and capacity' do
+      where(:scheduled_count, :unfinished_count, :capacity, :expected_count) do
+        2 | 2 | 10 | 4
+        2 | 0 | 10 | 2
+        0 | 2 | 10 | 2
+        4 | 2 | 2  | 4
+        4 | 0 | 2  | 4
+        0 | 4 | 2  | 4
+      end
+
+      with_them do
+        before do
+          allow(worker).to receive(:cleanup_scheduled_count).and_return(scheduled_count)
+          allow(worker).to receive(:cleanup_unfinished_count).and_return(unfinished_count)
+        end
+
+        it { is_expected.to eq(expected_count) }
+      end
+    end
+
     context 'with container repositories waiting for cleanup' do
       let_it_be(:unfinished_repositories) { create_list(:container_repository, 2, :cleanup_unfinished) }
 
       it { is_expected.to eq(3) }
-
-      it 'logs the work count' do
-        expect_log_info(
-          cleanup_scheduled_count: 1,
-          cleanup_unfinished_count: 2,
-          cleanup_total_count: 3
-        )
-
-        subject
-      end
     end
 
     context 'with no container repositories waiting for cleanup' do
@@ -436,16 +446,6 @@ RSpec.describe ContainerExpirationPolicies::CleanupContainerRepositoryWorker do
       end
 
       it { is_expected.to eq(0) }
-
-      it 'logs 0 work count' do
-        expect_log_info(
-          cleanup_scheduled_count: 0,
-          cleanup_unfinished_count: 0,
-          cleanup_total_count: 0
-        )
-
-        subject
-      end
     end
   end
 
@@ -468,9 +468,4 @@ RSpec.describe ContainerExpirationPolicies::CleanupContainerRepositoryWorker do
       it { is_expected.to eq(0) }
     end
   end
-
-  def expect_log_info(structure)
-    expect(worker.logger)
-      .to receive(:info).with(worker.structured_payload(structure))
-  end
 end

spec/workers/container_expiration_policy_worker_spec.rb

@@ -113,8 +113,8 @@ RSpec.describe ContainerExpirationPolicyWorker do
 
     context 'process stale ongoing cleanups' do
       let_it_be(:stuck_cleanup) { create(:container_repository, :cleanup_ongoing, expiration_policy_started_at: 1.day.ago) }
-      let_it_be(:container_repository) { create(:container_repository, :cleanup_scheduled) }
-      let_it_be(:container_repository) { create(:container_repository, :cleanup_unfinished) }
+      let_it_be(:container_repository1) { create(:container_repository, :cleanup_scheduled) }
+      let_it_be(:container_repository2) { create(:container_repository, :cleanup_unfinished) }
 
       it 'set them as unfinished' do
         expect { subject }
@@ -137,5 +137,36 @@ RSpec.describe ContainerExpirationPolicyWorker do
         expect(container_expiration_policy3.reload.enabled).to be false
       end
     end
+
+    context 'counts logging' do
+      let_it_be(:container_repository1) { create(:container_repository, :cleanup_scheduled) }
+      let_it_be(:container_repository2) { create(:container_repository, :cleanup_unfinished) }
+      let_it_be(:container_repository3) { create(:container_repository, :cleanup_unfinished) }
+
+      before do
+        ContainerExpirationPolicy.update_all(enabled: true)
+        container_repository1.project.container_expiration_policy.update_column(:next_run_at, 5.minutes.ago)
+      end
+
+      it 'logs all the counts' do
+        expect(worker).to receive(:log_extra_metadata_on_done).with(:cleanup_required_count, 1)
+        expect(worker).to receive(:log_extra_metadata_on_done).with(:cleanup_unfinished_count, 2)
+        expect(worker).to receive(:log_extra_metadata_on_done).with(:cleanup_total_count, 3)
+
+        subject
+      end
+
+      context 'with load balancing enabled' do
+        before do
+          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+        end
+
+        it 'reads the counts from 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