Bring Gitlab::ShardHealthCache to CE

It already existed in EE in the Geo namespace. This change brings it
to CE.

3 files changed, 95 insertions, 0 deletions

diff --git a/app/workers/concerns/each_shard_worker.rb b/app/workers/concerns/each_shard_worker.rb
index f063846427e..0bc2d63adcb 100644
--- a/app/workers/concerns/each_shard_worker.rb
+++ b/app/workers/concerns/each_shard_worker.rb
@@ -7,6 +7,8 @@ module EachShardWorker
   ].freeze
 
   def each_shard
+    Gitlab::ShardHealthCache.update(eligible_shard_names)
+
     eligible_shard_names.each do |shard_name|
       yield shard_name
     end
diff --git a/lib/gitlab/shard_health_cache.rb b/lib/gitlab/shard_health_cache.rb
new file mode 100644
index 00000000000..3f03f46d4b1
--- /dev/null
+++ b/lib/gitlab/shard_health_cache.rb
@@ -0,0 +1,41 @@
+module Gitlab
+  class ShardHealthCache
+    HEALTHY_SHARDS_KEY = 'gitlab-healthy-shards'.freeze
+    HEALTHY_SHARDS_TIMEOUT = 300
+
+    # Clears the Redis set storing the list of healthy shards
+    def self.clear
+      Gitlab::Redis::Cache.with { |redis| redis.del(HEALTHY_SHARDS_KEY) }
+    end
+
+    # Updates the list of healthy shards using a Redis set
+    #
+    # shards - An array of shard names to store
+    def self.update(shards)
+      Gitlab::Redis::Cache.with do |redis|
+        redis.multi do |m|
+          m.del(HEALTHY_SHARDS_KEY)
+          shards.each { |shard_name| m.sadd(HEALTHY_SHARDS_KEY, shard_name) }
+          m.expire(HEALTHY_SHARDS_KEY, HEALTHY_SHARDS_TIMEOUT)
+        end
+      end
+    end
+
+    # Returns an array of strings of healthy shards
+    def self.cached_healthy_shards
+      Gitlab::Redis::Cache.with { |redis| redis.smembers(HEALTHY_SHARDS_KEY) }
+    end
+
+    # Checks whether the given shard name is in the list of healthy shards.
+    #
+    # shard_name - The string to check
+    def self.healthy_shard?(shard_name)
+      Gitlab::Redis::Cache.with { |redis| redis.sismember(HEALTHY_SHARDS_KEY, shard_name) }
+    end
+
+    # Returns the number of healthy shards in the Redis set
+    def self.healthy_shard_count
+      Gitlab::Redis::Cache.with { |redis| redis.scard(HEALTHY_SHARDS_KEY) }
+    end
+  end
+end
diff --git a/spec/lib/gitlab/shard_health_cache_spec.rb b/spec/lib/gitlab/shard_health_cache_spec.rb
new file mode 100644
index 00000000000..e1a69261939
--- /dev/null
+++ b/spec/lib/gitlab/shard_health_cache_spec.rb
@@ -0,0 +1,52 @@
+require 'spec_helper'
+
+describe Gitlab::ShardHealthCache, :clean_gitlab_redis_cache do
+  let(:shards) { %w(foo bar) }
+
+  before do
+    described_class.update(shards)
+  end
+
+  describe '.clear' do
+    it 'leaves no shards around' do
+      described_class.clear
+
+      expect(described_class.healthy_shard_count).to eq(0)
+    end
+  end
+
+  describe '.update' do
+    it 'returns the healthy shards' do
+      expect(described_class.cached_healthy_shards).to match_array(shards)
+    end
+
+    it 'replaces the existing set' do
+      new_set = %w(test me more)
+      described_class.update(new_set)
+
+      expect(described_class.cached_healthy_shards).to match_array(new_set)
+    end
+  end
+
+  describe '.healthy_shard_count' do
+    it 'returns the healthy shard count' do
+      expect(described_class.healthy_shard_count).to eq(2)
+    end
+
+    it 'returns 0 if no shards are available' do
+      described_class.update([])
+
+      expect(described_class.healthy_shard_count).to eq(0)
+    end
+  end
+
+  describe '.healthy_shard?' do
+    it 'returns true for a healthy shard' do
+      expect(described_class.healthy_shard?('foo')).to be_truthy
+    end
+
+    it 'returns false for an unknown shard' do
+      expect(described_class.healthy_shard?('unknown')).to be_falsey
+    end
+  end
+end