gitlab-org--gitlab-foss/lib/gitlab/instrumentation/redis_cluster_validator.rb

106 lines
3.6 KiB
Ruby

# 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