Add Redis Cluster validator

If we want to use Redis Cluster, we won't be able to use a multi-key
command with keys that hash to different slots. For instance, a simple:

    MGET foo bar

Won't work because 'foo' and 'bar' hash to different slots, and
therefore could be on different shards of the cluster.

This is a client-side validator to ensure that we can annotate existing
cross-slot commands easily, as part of the data gathering for whether or
not we use Redis Cluster.

app/models/active_session.rb

@@ -91,8 +91,11 @@ class ActiveSession
     key_names = session_ids.map { |session_id| key_name(user.id, session_id.public_id) }
 
     redis.srem(lookup_key_name(user.id), session_ids.map(&:public_id))
-    redis.del(key_names)
-    redis.del(rack_session_keys(session_ids))
+
+    Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+      redis.del(key_names)
+      redis.del(rack_session_keys(session_ids))
+    end
   end
 
   def self.cleanup(user)
@@ -136,8 +139,10 @@ class ActiveSession
       session_keys = rack_session_keys(session_ids)
 
       session_keys.each_slice(SESSION_BATCH_SIZE).flat_map do |session_keys_batch|
-        redis.mget(session_keys_batch).compact.map do |raw_session|
-          load_raw_session(raw_session)
+        Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+          redis.mget(session_keys_batch).compact.map do |raw_session|
+            load_raw_session(raw_session)
+          end
         end
       end
     end
@@ -178,7 +183,9 @@ class ActiveSession
 
     entry_keys = session_ids.map { |session_id| key_name(user_id, session_id) }
 
-    redis.mget(entry_keys)
+    Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+      redis.mget(entry_keys)
+    end
   end
 
   def self.active_session_entries(session_ids, user_id, redis)

app/models/ci/build_trace_chunks/redis.rb

@@ -35,7 +35,10 @@ module Ci
         keys = keys.map { |key| key_raw(*key) }
 
         Gitlab::Redis::SharedState.with do |redis|
-          redis.del(keys)
+          # https://gitlab.com/gitlab-org/gitlab/-/issues/224171
+          Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+            redis.del(keys)
+          end
         end
       end
 

doc/development/redis.md

@@ -33,6 +33,8 @@ stop being consulted if the project is renamed. If the contents of the key are
 invalidated by a name change, it is better to include a hook that will expire
 the entry, instead of relying on the key changing.
 
+### Multi-key commands
+
 We don't use [Redis Cluster](https://redis.io/topics/cluster-tutorial) at the
 moment, but may wish to in the future: [#118820](https://gitlab.com/gitlab-org/gitlab/-/issues/118820).
 
@@ -41,3 +43,8 @@ operations that require several keys to be held on the same Redis server - for
 instance, diffing two sets held in Redis - the keys should ensure that by
 enclosing the changeable parts in curly braces, such as, `project:{1}:set_a` and
 `project:{1}:set_b`.
+
+Currently, we validate this in the development and test environments
+with the [`RedisClusterValidator`](https://gitlab.com/gitlab-org/gitlab/-/blob/master/lib/gitlab/instrumentation/redis_cluster_validator.rb),
+which is enabled for the `cache` and `shared_state`
+[Redis instances](https://docs.gitlab.com/omnibus/settings/redis.html#running-with-multiple-redis-instances)..

ee/app/services/elastic/indexing_control_service.rb

@@ -68,7 +68,9 @@ module Elastic
     attr_reader :klass, :queue_name, :redis_set_key, :redis_score_key
 
     def with_redis(&blk)
-      Gitlab::Redis::SharedState.with(&blk) # rubocop:disable CodeReuse/ActiveRecord
+      Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+        Gitlab::Redis::SharedState.with(&blk) # rubocop:disable CodeReuse/ActiveRecord
+      end
     end
 
     def serialize(args, context)

ee/app/services/elastic/process_bookkeeping_service.rb

@@ -38,7 +38,11 @@ module Elastic
       end
 
       def clear_tracking!
-        with_redis { |redis| redis.del(self::REDIS_SET_KEY, self::REDIS_SCORE_KEY) }
+        with_redis do |redis|
+          Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+            redis.del(self::REDIS_SET_KEY, self::REDIS_SCORE_KEY)
+          end
+        end
       end
 
       def logger

lib/gitlab/analytics/unique_visits.rb

@@ -43,7 +43,9 @@ module Gitlab
         keys = TARGET_IDS.map { |target_id| key(target_id, week_of) }
 
         Gitlab::Redis::SharedState.with do |redis|
-          redis.pfcount(*keys)
+          Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+            redis.pfcount(*keys)
+          end
         end
       end
 

lib/gitlab/discussions_diff/highlight_cache.rb

@@ -36,7 +36,9 @@ module Gitlab
 
           content =
             Redis::Cache.with do |redis|
-              redis.mget(keys)
+              Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+                redis.mget(keys)
+              end
             end
 
           content.map! do |lines|
@@ -58,7 +60,11 @@ module Gitlab
 
           keys = raw_keys.map { |id| cache_key_for(id) }
 
-          Redis::Cache.with { |redis| redis.del(keys) }
+          Redis::Cache.with do |redis|
+            Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+              redis.del(keys)
+            end
+          end
         end
 
         def cache_key_for(raw_key)

lib/gitlab/instrumentation/redis.rb

@@ -5,9 +5,9 @@ module Gitlab
     # Aggregates Redis measurements from different request storage sources.
     class Redis
       ActionCable = Class.new(RedisBase)
-      Cache = Class.new(RedisBase)
+      Cache = Class.new(RedisBase).enable_redis_cluster_validation
       Queues = Class.new(RedisBase)
-      SharedState = Class.new(RedisBase)
+      SharedState = Class.new(RedisBase).enable_redis_cluster_validation
 
       STORAGES = [ActionCable, Cache, Queues, SharedState].freeze
 

lib/gitlab/instrumentation/redis_base.rb

@@ -71,6 +71,16 @@ module Gitlab
           query_time.round(::Gitlab::InstrumentationHelper::DURATION_PRECISION)
         end
 
+        def redis_cluster_validate!(command)
+          RedisClusterValidator.validate!(command) if @redis_cluster_validation
+        end
+
+        def enable_redis_cluster_validation
+          @redis_cluster_validation = true
+
+          self
+        end
+
         private
 
         def request_count_key

lib/gitlab/instrumentation/redis_cluster_validator.rb

+# frozen_string_literal: true
+
+require 'rails'
+require 'redis'
+
+module Gitlab
+  module Instrumentation
+    module RedisClusterValidator
+      # Generate with:
+      #
+      # Gitlab::Redis::Cache
+      #   .with { |redis| redis.call('COMMAND') }
+      #   .select { |command| command[3] != command[4] }
+      #   .map { |command| [command[0].upcase, { first: command[3], last: command[4], step: command[5] }] }
+      #   .sort_by(&:first)
+      #   .to_h
+      #
+      MULTI_KEY_COMMANDS = {
+        "BITOP" => { first: 2, last: -1, step: 1 },
+        "BLPOP" => { first: 1, last: -2, step: 1 },
+        "BRPOP" => { first: 1, last: -2, step: 1 },
+        "BRPOPLPUSH" => { first: 1, last: 2, step: 1 },
+        "BZPOPMAX" => { first: 1, last: -2, step: 1 },
+        "BZPOPMIN" => { first: 1, last: -2, step: 1 },
+        "DEL" => { first: 1, last: -1, step: 1 },
+        "EXISTS" => { first: 1, last: -1, step: 1 },
+        "MGET" => { first: 1, last: -1, step: 1 },
+        "MSET" => { first: 1, last: -1, step: 2 },
+        "MSETNX" => { first: 1, last: -1, step: 2 },
+        "PFCOUNT" => { first: 1, last: -1, step: 1 },
+        "PFMERGE" => { first: 1, last: -1, step: 1 },
+        "RENAME" => { first: 1, last: 2, step: 1 },
+        "RENAMENX" => { first: 1, last: 2, step: 1 },
+        "RPOPLPUSH" => { first: 1, last: 2, step: 1 },
+        "SDIFF" => { first: 1, last: -1, step: 1 },
+        "SDIFFSTORE" => { first: 1, last: -1, step: 1 },
+        "SINTER" => { first: 1, last: -1, step: 1 },
+        "SINTERSTORE" => { first: 1, last: -1, step: 1 },
+        "SMOVE" => { first: 1, last: 2, step: 1 },
+        "SUNION" => { first: 1, last: -1, step: 1 },
+        "SUNIONSTORE" => { first: 1, last: -1, step: 1 },
+        "UNLINK" => { first: 1, last: -1, step: 1 },
+        "WATCH" => { first: 1, last: -1, step: 1 }
+      }.freeze
+
+      CrossSlotError = Class.new(StandardError)
+
+      class << self
+        def validate!(command)
+          return unless Rails.env.development? || Rails.env.test?
+          return if allow_cross_slot_commands?
+
+          command_name = command.first.to_s.upcase
+          argument_positions = MULTI_KEY_COMMANDS[command_name]
+
+          return unless argument_positions
+
+          arguments = command.flatten[argument_positions[:first]..argument_positions[:last]]
+
+          key_slots = arguments.each_slice(argument_positions[:step]).map do |args|
+            key_slot(args.first)
+          end
+
+          unless key_slots.uniq.length == 1
+            raise CrossSlotError.new("Redis command #{command_name} arguments hash to different slots. See https://docs.gitlab.com/ee/development/redis.html#multi-key-commands")
+          end
+        end
+
+        # Keep track of the call stack to allow nested calls to work.
+        def allow_cross_slot_commands
+          Thread.current[:allow_cross_slot_commands] ||= 0
+          Thread.current[:allow_cross_slot_commands] += 1
+
+          yield
+        ensure
+          Thread.current[:allow_cross_slot_commands] -= 1
+        end
+
+        private
+
+        def allow_cross_slot_commands?
+          Thread.current[:allow_cross_slot_commands].to_i > 0
+        end
+
+        def key_slot(key)
+          ::Redis::Cluster::KeySlotConverter.convert(extract_hash_tag(key))
+        end
+
+        # This is almost identical to Redis::Cluster::Command#extract_hash_tag,
+        # except that it returns the original string if no hash tag is found.
+        #
+        def extract_hash_tag(key)
+          s = key.index('{')
+
+          return key unless s
+
+          e = key.index('}', s + 1)
+
+          return key unless e
+
+          key[s + 1..e - 1]
+        end
+      end
+    end
+  end
+end

lib/gitlab/instrumentation/redis_interceptor.rb

@@ -7,6 +7,9 @@ module Gitlab
     module RedisInterceptor
       def call(*args, &block)
         start = Time.now
+
+        instrumentation_class.redis_cluster_validate!(args.first)
+
         super(*args, &block)
       ensure
         duration = (Time.now - start)

lib/gitlab/set_cache.rb

@@ -20,7 +20,10 @@ module Gitlab
 
       with do |redis|
         keys = keys.map { |key| cache_key(key) }
-        unlink_or_delete(redis, keys)
+
+        Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+          unlink_or_delete(redis, keys)
+        end
       end
     end
 

lib/tasks/cache.rake

@@ -18,7 +18,9 @@ namespace :cache do
               count: REDIS_CLEAR_BATCH_SIZE
             )
 
-            redis.del(*keys) if keys.any?
+            Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+              redis.del(*keys) if keys.any?
+            end
 
             break if cursor == REDIS_SCAN_START_STOP
           end

spec/lib/gitlab/instrumentation/redis_cluster_validator_spec.rb

+# frozen_string_literal: true
+
+require 'fast_spec_helper'
+require 'support/helpers/rails_helpers'
+require 'rspec-parameterized'
+
+describe Gitlab::Instrumentation::RedisClusterValidator do
+  include RailsHelpers
+
+  describe '.validate!' do
+    using RSpec::Parameterized::TableSyntax
+
+    context 'Rails environments' do
+      where(:env, :should_raise) do
+        'production' | false
+        'staging' | false
+        'development' | true
+        'test' | true
+      end
+
+      with_them do
+        it do
+          stub_rails_env(env)
+
+          args = [:mget, 'foo', 'bar']
+
+          if should_raise
+            expect { described_class.validate!(args) }
+              .to raise_error(described_class::CrossSlotError)
+          else
+            expect { described_class.validate!(args) }.not_to raise_error
+          end
+        end
+      end
+    end
+
+    where(:command, :arguments, :should_raise) do
+      :rename | %w(foo bar) | true
+      :RENAME | %w(foo bar) | true
+      'rename' | %w(foo bar) | true
+      'RENAME' | %w(foo bar) | true
+      :rename | %w(iaa ahy) | false # 'iaa' and 'ahy' hash to the same slot
+      :rename | %w({foo}:1 {foo}:2) | false
+      :rename | %w(foo foo bar) | false # This is not a valid command but should not raise here
+      :mget | %w(foo bar) | true
+      :mget | %w(foo foo bar) | true
+      :mget | %w(foo foo) | false
+      :blpop | %w(foo bar 1) | true
+      :blpop | %w(foo foo 1) | false
+      :mset | %w(foo a bar a) | true
+      :mset | %w(foo a foo a) | false
+      :del | %w(foo bar) | true
+      :del | [%w(foo bar)] | true # Arguments can be a nested array
+      :del | %w(foo foo) | false
+      :hset | %w(foo bar) | false # Not a multi-key command
+    end
+
+    with_them do
+      it do
+        args = [command] + arguments
+
+        if should_raise
+          expect { described_class.validate!(args) }
+            .to raise_error(described_class::CrossSlotError)
+        else
+          expect { described_class.validate!(args) }.not_to raise_error
+        end
+      end
+    end
+  end
+
+  describe '.allow_cross_slot_commands' do
+    it 'does not raise for invalid arguments' do
+      expect do
+        described_class.allow_cross_slot_commands do
+          described_class.validate!([:mget, 'foo', 'bar'])
+        end
+      end.not_to raise_error
+    end
+
+    it 'allows nested invocation' do
+      expect do
+        described_class.allow_cross_slot_commands do
+          described_class.allow_cross_slot_commands do
+            described_class.validate!([:mget, 'foo', 'bar'])
+          end
+
+          described_class.validate!([:mget, 'foo', 'bar'])
+        end
+      end.not_to raise_error
+    end
+  end
+end