Extract EachShardWorker into a concern

1 files changed, 36 insertions, 0 deletions

diff --git a/app/workers/concerns/each_shard_worker.rb b/app/workers/concerns/each_shard_worker.rb
new file mode 100644
index 00000000000..f063846427e
--- /dev/null
+++ b/app/workers/concerns/each_shard_worker.rb
@@ -0,0 +1,36 @@
+module EachShardWorker
+  extend ActiveSupport::Concern
+  include ::Gitlab::Utils::StrongMemoize
+
+  HEALTHY_SHARD_CHECKS = [
+    Gitlab::HealthChecks::GitalyCheck
+  ].freeze
+
+  def each_shard
+    eligible_shard_names.each do |shard_name|
+      yield shard_name
+    end
+  end
+
+  # override when you want to filter out some shards
+  def eligible_shard_names
+    healthy_shard_names
+  end
+
+  def healthy_shard_names
+    strong_memoize(:healthy_shard_names) do
+      # For now, we need to perform both Gitaly and direct filesystem checks to ensure
+      # the shard is healthy. We take the intersection of the successful checks
+      # as the healthy shards.
+      healthy_ready_shards.map { |result| result.labels[:shard] }.compact.uniq
+    end
+  end
+
+  def healthy_ready_shards
+    ready_shards.map { |result| result.select(&:success) }.inject(:&)
+  end
+
+  def ready_shards
+    HEALTHY_SHARD_CHECKS.map(&:readiness)
+  end
+end