1 files changed, 10 insertions, 6 deletions

diff --git a/lib/gitlab/shard_health_cache.rb b/lib/gitlab/shard_health_cache.rb
index eeb0cc75ef9..a34e4e9c8d1 100644
--- a/lib/gitlab/shard_health_cache.rb
+++ b/lib/gitlab/shard_health_cache.rb
@@ -7,17 +7,17 @@ module Gitlab
 
     # Clears the Redis set storing the list of healthy shards
     def self.clear
-      Gitlab::Redis::Cache.with { |redis| redis.del(HEALTHY_SHARDS_KEY) }
+      with_redis { |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|
+      with_redis do |redis|
         redis.multi do |m|
           m.del(HEALTHY_SHARDS_KEY)
-          shards.each { |shard_name| m.sadd(HEALTHY_SHARDS_KEY, shard_name) }
+          m.sadd(HEALTHY_SHARDS_KEY, shards) unless shards.blank?
           m.expire(HEALTHY_SHARDS_KEY, HEALTHY_SHARDS_TIMEOUT)
         end
       end
@@ -25,19 +25,23 @@ module Gitlab
 
     # Returns an array of strings of healthy shards
     def self.cached_healthy_shards
-      Gitlab::Redis::Cache.with { |redis| redis.smembers(HEALTHY_SHARDS_KEY) }
+      with_redis { |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) }
+      with_redis { |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) }
+      with_redis { |redis| redis.scard(HEALTHY_SHARDS_KEY) }
+    end
+
+    def self.with_redis(&block)
+      Gitlab::Redis::Cache.with(&block) # rubocop:disable CodeReuse/ActiveRecord
     end
   end
 end