Keep track of all storage keys in a setbvl-circuitbreaker-keys-set

That way we don't need the to scan the entire keyspace to get all
known keys

6 files changed, 42 insertions, 23 deletions

diff --git a/changelogs/unreleased/bvl-circuitbreaker-keys-set.yml b/changelogs/unreleased/bvl-circuitbreaker-keys-set.yml
new file mode 100644
index 00000000000..a56456240df
--- /dev/null
+++ b/changelogs/unreleased/bvl-circuitbreaker-keys-set.yml
@@ -0,0 +1,5 @@
+---
+title: Keep track of all circuitbreaker keys in a set
+merge_request: 15613
+author:
+type: performance
diff --git a/lib/gitlab/git/storage.rb b/lib/gitlab/git/storage.rb
index 99518c9b1e4..5933312b0b5 100644
--- a/lib/gitlab/git/storage.rb
+++ b/lib/gitlab/git/storage.rb
@@ -15,6 +15,7 @@ module Gitlab
       Failing = Class.new(Inaccessible)
 
       REDIS_KEY_PREFIX = 'storage_accessible:'.freeze
+      REDIS_KNOWN_KEYS = "#{REDIS_KEY_PREFIX}known_keys_set".freeze
 
       def self.redis
         Gitlab::Redis::SharedState
diff --git a/lib/gitlab/git/storage/circuit_breaker.rb b/lib/gitlab/git/storage/circuit_breaker.rb
index be7598ef011..4328c0ea29b 100644
--- a/lib/gitlab/git/storage/circuit_breaker.rb
+++ b/lib/gitlab/git/storage/circuit_breaker.rb
@@ -13,10 +13,8 @@ module Gitlab
         delegate :last_failure, :failure_count, to: :failure_info
 
         def self.reset_all!
-          pattern = "#{Gitlab::Git::Storage::REDIS_KEY_PREFIX}*"
-
           Gitlab::Git::Storage.redis.with do |redis|
-            all_storage_keys = redis.scan_each(match: pattern).to_a
+            all_storage_keys = redis.zrange(Gitlab::Git::Storage::REDIS_KNOWN_KEYS, 0, -1)
             redis.del(*all_storage_keys) unless all_storage_keys.empty?
           end
 
@@ -135,23 +133,29 @@ module Gitlab
               redis.hset(cache_key, :last_failure, last_failure.to_i)
               redis.hincrby(cache_key, :failure_count, 1)
               redis.expire(cache_key, failure_reset_time)
+              maintain_known_keys(redis)
             end
           end
         end
 
         def track_storage_accessible
-          return if no_failures?
-
           @failure_info = FailureInfo.new(nil, 0)
 
           Gitlab::Git::Storage.redis.with do |redis|
             redis.pipelined do
               redis.hset(cache_key, :last_failure, nil)
               redis.hset(cache_key, :failure_count, 0)
+              maintain_known_keys(redis)
             end
           end
         end
 
+        def maintain_known_keys(redis)
+          expire_time = Time.now.to_i + failure_reset_time
+          redis.zadd(Gitlab::Git::Storage::REDIS_KNOWN_KEYS, expire_time, cache_key)
+          redis.zremrangebyscore(Gitlab::Git::Storage::REDIS_KNOWN_KEYS, '-inf', Time.now.to_i)
+        end
+
         def get_failure_info
           last_failure, failure_count = Gitlab::Git::Storage.redis.with do |redis|
             redis.hmget(cache_key, :last_failure, :failure_count)
diff --git a/lib/gitlab/git/storage/health.rb b/lib/gitlab/git/storage/health.rb
index 7049772fe3b..90bbe85fd37 100644
--- a/lib/gitlab/git/storage/health.rb
+++ b/lib/gitlab/git/storage/health.rb
@@ -4,8 +4,8 @@ module Gitlab
       class Health
         attr_reader :storage_name, :info
 
-        def self.pattern_for_storage(storage_name)
-          "#{Gitlab::Git::Storage::REDIS_KEY_PREFIX}#{storage_name}:*"
+        def self.prefix_for_storage(storage_name)
+          "#{Gitlab::Git::Storage::REDIS_KEY_PREFIX}#{storage_name}:"
         end
 
         def self.for_all_storages
@@ -25,26 +25,15 @@ module Gitlab
 
         private_class_method def self.all_keys_for_storages(storage_names, redis)
           keys_per_storage = {}
+          all_keys = redis.zrange(Gitlab::Git::Storage::REDIS_KNOWN_KEYS, 0, -1)
 
-          redis.pipelined do
-            storage_names.each do |storage_name|
-              pattern = pattern_for_storage(storage_name)
-              matched_keys = redis.scan_each(match: pattern)
+          storage_names.each do |storage_name|
+            prefix = prefix_for_storage(storage_name)
 
-              keys_per_storage[storage_name] = matched_keys
-            end
+            keys_per_storage[storage_name] = all_keys.select { |key| key.starts_with?(prefix) }
           end
 
-          # We need to make sure each lazy-loaded `Enumerator` for matched keys
-          # is loaded into an array.
-          #
-          # Otherwise it would be loaded in the second `Redis#pipelined` block
-          # within `.load_for_keys`. In this pipelined call, the active
-          # Redis-client changes again, so the values would not be available
-          # until the end of that pipelined-block.
-          keys_per_storage.each do |storage_name, key_future|
-            keys_per_storage[storage_name] = key_future.to_a
-          end
+          keys_per_storage
         end
 
         private_class_method def self.load_for_keys(keys_per_storage, redis)
diff --git a/spec/lib/gitlab/git/storage/circuit_breaker_spec.rb b/spec/lib/gitlab/git/storage/circuit_breaker_spec.rb
index 72dabca793a..f34c9f09057 100644
--- a/spec/lib/gitlab/git/storage/circuit_breaker_spec.rb
+++ b/spec/lib/gitlab/git/storage/circuit_breaker_spec.rb
@@ -27,6 +27,7 @@ describe Gitlab::Git::Storage::CircuitBreaker, clean_gitlab_redis_shared_state:
 
   def set_in_redis(name, value)
     Gitlab::Git::Storage.redis.with do |redis|
+      redis.zadd(Gitlab::Git::Storage::REDIS_KNOWN_KEYS, 0, cache_key)
       redis.hmset(cache_key, name, value)
     end.first
   end
@@ -181,6 +182,24 @@ describe Gitlab::Git::Storage::CircuitBreaker, clean_gitlab_redis_shared_state:
       expect(circuit_breaker.last_failure).to be_nil
     end
 
+    it 'maintains known storage keys' do
+      Timecop.freeze do
+        # Insert an old key to expire
+        old_entry = Time.now.to_i - 3.days.to_i
+        Gitlab::Git::Storage.redis.with do |redis|
+          redis.zadd(Gitlab::Git::Storage::REDIS_KNOWN_KEYS, old_entry, 'to_be_removed')
+        end
+
+        circuit_breaker.perform { '' }
+
+        known_keys = Gitlab::Git::Storage.redis.with do |redis|
+          redis.zrange(Gitlab::Git::Storage::REDIS_KNOWN_KEYS, 0, -1)
+        end
+
+        expect(known_keys).to contain_exactly(cache_key)
+      end
+    end
+
     it 'only performs the accessibility check once' do
       expect(Gitlab::Git::Storage::ForkedStorageCheck)
         .to receive(:storage_available?).once.and_call_original
diff --git a/spec/lib/gitlab/git/storage/health_spec.rb b/spec/lib/gitlab/git/storage/health_spec.rb
index 4a14a5201d1..d7a52a04fbb 100644
--- a/spec/lib/gitlab/git/storage/health_spec.rb
+++ b/spec/lib/gitlab/git/storage/health_spec.rb
@@ -6,6 +6,7 @@ describe Gitlab::Git::Storage::Health, clean_gitlab_redis_shared_state: true, br
 
   def set_in_redis(cache_key, value)
     Gitlab::Git::Storage.redis.with do |redis|
+      redis.zadd(Gitlab::Git::Storage::REDIS_KNOWN_KEYS, 0, cache_key)
       redis.hmset(cache_key, :failure_count, value)
     end.first
   end