summaryrefslogtreecommitdiff
path: root/lib/gitlab/instrumentation/redis_cluster_validator.rb
blob: 6800e5667f647dbfcc119fcd9b2cfb816779eaf4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
# 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