Add RepositoryCheck::DispatchWorker to start worker per shard

The RepositoryCheck::DispatchWorker will start a
RepositoryCheck::BatchWorker for each healthy shard.

Closes gitlab-org/gitlab-ce#48042

3 files changed, 30 insertions, 5 deletions

diff --git a/app/workers/all_queues.yml b/app/workers/all_queues.yml
index 026f756582d..d06f51b1828 100644
--- a/app/workers/all_queues.yml
+++ b/app/workers/all_queues.yml
@@ -11,7 +11,7 @@
 - cronjob:remove_old_web_hook_logs
 - cronjob:remove_unreferenced_lfs_objects
 - cronjob:repository_archive_cache
-- cronjob:repository_check_batch
+- cronjob:repository_check_dispatch
 - cronjob:requests_profiles
 - cronjob:schedule_update_user_activity
 - cronjob:stuck_ci_jobs
@@ -71,6 +71,7 @@
 - pipeline_processing:update_head_pipeline_for_merge_request
 
 - repository_check:repository_check_clear
+- repository_check:repository_check_batch
 - repository_check:repository_check_single_repository
 
 - default
diff --git a/app/workers/repository_check/batch_worker.rb b/app/workers/repository_check/batch_worker.rb
index 898bca976ec..051382a08a9 100644
--- a/app/workers/repository_check/batch_worker.rb
+++ b/app/workers/repository_check/batch_worker.rb
@@ -3,13 +3,18 @@
 module RepositoryCheck
   class BatchWorker
     include ApplicationWorker
-    include CronjobQueue
+    include RepositoryCheckQueue
 
     RUN_TIME = 3600
     BATCH_SIZE = 10_000
 
-    def perform
+    attr_reader :shard_name
+
+    def perform(shard_name)
+      @shard_name = shard_name
+
       return unless Gitlab::CurrentSettings.repository_checks_enabled
+      return unless Gitlab::ShardHealthCache.healthy_shard?(shard_name)
 
       start = Time.now
 
@@ -39,18 +44,22 @@ module RepositoryCheck
     end
 
     def never_checked_project_ids(batch_size)
-      Project.where(last_repository_check_at: nil)
+      projects_on_shard.where(last_repository_check_at: nil)
         .where('created_at < ?', 24.hours.ago)
         .limit(batch_size).pluck(:id)
     end
 
     def old_checked_project_ids(batch_size)
-      Project.where.not(last_repository_check_at: nil)
+      projects_on_shard.where.not(last_repository_check_at: nil)
         .where('last_repository_check_at < ?', 1.month.ago)
         .reorder(last_repository_check_at: :asc)
         .limit(batch_size).pluck(:id)
     end
 
+    def projects_on_shard
+      Project.where(repository_storage: shard_name)
+    end
+
     def try_obtain_lease(id)
       # Use a 24-hour timeout because on servers/projects where 'git fsck' is
       # super slow we definitely do not want to run it twice in parallel.
diff --git a/app/workers/repository_check/dispatch_worker.rb b/app/workers/repository_check/dispatch_worker.rb
new file mode 100644
index 00000000000..403479a8497
--- /dev/null
+++ b/app/workers/repository_check/dispatch_worker.rb
@@ -0,0 +1,15 @@
+module RepositoryCheck
+  class DispatchWorker
+    include ApplicationWorker
+    include CronjobQueue
+    include ::EachShardWorker
+
+    def perform
+      return unless Gitlab::CurrentSettings.repository_checks_enabled
+
+      each_shard do |shard_name|
+        RepositoryCheck::BatchWorker.perform_async(shard_name)
+      end
+    end
+  end
+end