Use RequestStore to minimise feature flag lookup

This commit addresses other review nits too.

config/feature_flags/ops/validate_allowed_cross_slot_commands.yml → config/feature_flags/development/validate_allowed_cross_slot_commands.yml

@@ -3,6 +3,6 @@ name: validate_allowed_cross_slot_commands
 introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/105302
 rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/384909
 milestone: '15.7'
-type: ops
+type: development
 group: group::scalability
 default_enabled: false

lib/gitlab/instrumentation/redis_base.rb

@@ -75,10 +75,17 @@ def query_time
           query_time.round(::Gitlab::InstrumentationHelper::DURATION_PRECISION)
         end
 
+        def validate_allowed_commands?
+          flag = ::RequestStore[validate_allowed_commands_key]
+          return flag unless flag.nil?
+
+          ::RequestStore[validate_allowed_commands_key] = Feature.enabled?(:validate_allowed_cross_slot_commands, type: :development)
+        end
+
         def redis_cluster_validate!(commands)
           return true unless @redis_cluster_validation
 
-          result = ::Gitlab::Instrumentation::RedisClusterValidator.validate(commands)
+          result = ::Gitlab::Instrumentation::RedisClusterValidator.validate(commands, validate_allowed_commands?)
           return true if result.nil?
 
           if !result[:valid] && !result[:allowed] && (Rails.env.development? || Rails.env.test?)
@@ -144,6 +151,10 @@ def cross_slots_key
           strong_memoize(:cross_slots_key) { build_key(:redis_cross_slot_request_count) }
         end
 
+        def validate_allowed_commands_key
+          strong_memoize(:validate_allowed_commands_key) { "validate_allowed_commands_flag" }
+        end
+
         def build_key(namespace)
           "#{storage_key}_#{namespace}"
         end

lib/gitlab/instrumentation/redis_cluster_validator.rb

@@ -183,17 +183,14 @@ module RedisClusterValidator
       CrossSlotError = Class.new(StandardError)
 
       class << self
-        def validate(commands)
+        def validate(commands, validate_allowed_cmd)
+          return if allow_cross_slot_commands? && !validate_allowed_cmd
           return if commands.empty?
 
           # early exit for single-command (non-pipelined) if it is a single-key-command
           command_name = commands.size > 1 ? "PIPELINE/MULTI" : commands.first.first.to_s.upcase
           return if commands.size == 1 && REDIS_COMMANDS.dig(command_name, :single_key)
 
-          # Checking Feature.enabled? after single-command check
-          # as Feature.enabled? runs a `get` command which is single key. This prevents recursive calls.
-          return if allow_cross_slot_commands? && !::Feature.enabled?(:validate_allowed_cross_slot_commands, type: :ops)
-
           key_slots = commands.map { |command| key_slots(command) }.flatten
 
           {

spec/lib/gitlab/instrumentation/redis_base_spec.rb

@@ -153,6 +153,12 @@
   end
 
   describe '.redis_cluster_validate!' do
+    let(:args) { [[:mget, 'foo', 'bar']] }
+
+    before do
+      instrumentation_class_a.enable_redis_cluster_validation
+    end
+
     context 'Rails environments' do
       where(:env, :allowed, :should_raise) do
         'production' | false | false
@@ -165,16 +171,10 @@
         'test' | false | true
       end
 
-      before do
-        instrumentation_class_a.enable_redis_cluster_validation
-      end
-
       with_them do
         it do
           stub_rails_env(env)
 
-          args = [[:mget, 'foo', 'bar']]
-
           validation = -> { instrumentation_class_a.redis_cluster_validate!(args) }
           under_test = if allowed
                          -> { Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands(&validation) }
@@ -190,5 +190,53 @@
         end
       end
     end
+
+    context 'validate_allowed_cross_slot_commands feature flag' do
+      context 'when disabled' do
+        before do
+          stub_feature_flags(validate_allowed_cross_slot_commands: false)
+        end
+
+        it 'skips check' do
+          expect(
+            Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+              instrumentation_class_a.redis_cluster_validate!(args)
+            end
+          ).to eq(true)
+        end
+      end
+
+      context 'when enabled' do
+        before do
+          stub_feature_flags(validate_allowed_cross_slot_commands: true)
+        end
+
+        it 'performs check' do
+          expect(
+            Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+              instrumentation_class_a.redis_cluster_validate!(args)
+            end
+          ).to eq(false)
+        end
+      end
+
+      it 'looks up feature-flag once per request' do
+        stub_feature_flags(validate_allowed_cross_slot_commands: true)
+        expect(
+          Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+            instrumentation_class_a.redis_cluster_validate!(args)
+          end
+        ).to eq(false)
+
+        # even with validate set to false, redis_cluster_validate! will use cached
+        # feature flag value and perform validation
+        stub_feature_flags(validate_allowed_cross_slot_commands: false)
+        expect(
+          Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+            instrumentation_class_a.redis_cluster_validate!(args)
+          end
+        ).to eq(false)
+      end
+    end
   end
 end

spec/lib/gitlab/instrumentation/redis_cluster_validator_spec.rb

 # frozen_string_literal: true
 
-require 'spec_helper'
+require 'fast_spec_helper'
 require 'support/helpers/rails_helpers'
 require 'rspec-parameterized'
 
@@ -10,10 +10,6 @@
   describe '.validate' do
     using RSpec::Parameterized::TableSyntax
 
-    before do
-      stub_feature_flags(validate_allowed_cross_slot_commands: true)
-    end
-
     where(:command, :arguments, :keys, :is_valid) do
       :rename | %w(foo bar) | 2 | false
       :RENAME | %w(foo bar) | 2 | false
@@ -41,53 +37,59 @@
       it do
         args = [[command] + arguments]
         if is_valid.nil?
-          expect(described_class.validate(args)).to eq(nil)
+          expect(described_class.validate(args, true)).to eq(nil)
         else
-          expect(described_class.validate(args)[:valid]).to eq(is_valid)
-          expect(described_class.validate(args)[:allowed]).to eq(false)
-          expect(described_class.validate(args)[:command_name]).to eq(command.to_s.upcase)
-          expect(described_class.validate(args)[:key_count]).to eq(keys)
+          expect(described_class.validate(args, true)[:valid]).to eq(is_valid)
+          expect(described_class.validate(args, true)[:allowed]).to eq(false)
+          expect(described_class.validate(args, true)[:command_name]).to eq(command.to_s.upcase)
+          expect(described_class.validate(args, true)[:key_count]).to eq(keys)
         end
       end
     end
 
     where(:arguments, :should_raise, :output) do
-      [[:get, "foo"],
-       [:get,
-        "bar"]] | true | { valid: false, key_count: 2, command_name: 'PIPELINE/MULTI', allowed: false,
-                           command: "get foo" }
-      [[:get, "foo"],
-       [:mget, "foo",
-        "bar"]] | true | { valid: false, key_count: 3, command_name: 'PIPELINE/MULTI', allowed: false,
-                           command: "get foo" }
-      [[:get, "{foo}:name"],
-       [:get,
-        "{foo}:profile"]] | false | { valid: true, key_count: 2, command_name: 'PIPELINE/MULTI', allowed: false,
-                                      command: "get {foo}:name" }
-      [[:del, "foo"],
-       [:del,
-        "bar"]] | true | { valid: false, key_count: 2, command_name: 'PIPELINE/MULTI', allowed: false,
-                           command: "del foo" }
-      [] | false | nil # pipeline or transaction opened and closed without ops
+      [
+        [
+          [[:get, "foo"], [:get, "bar"]],
+          true,
+          { valid: false, key_count: 2, command_name: 'PIPELINE/MULTI', allowed: false, command: "get foo" }
+        ],
+        [
+          [[:get, "foo"], [:mget, "foo", "bar"]],
+          true,
+          { valid: false, key_count: 3, command_name: 'PIPELINE/MULTI', allowed: false, command: "get foo" }
+        ],
+        [
+          [[:get, "{foo}:name"], [:get, "{foo}:profile"]],
+          false,
+          { valid: true, key_count: 2, command_name: 'PIPELINE/MULTI', allowed: false, command: "get {foo}:name" }
+        ],
+        [
+          [[:del, "foo"], [:del, "bar"]],
+          true,
+          { valid: false, key_count: 2, command_name: 'PIPELINE/MULTI', allowed: false, command: "del foo" }
+        ],
+        [
+          [],
+          false,
+          nil # pipeline or transaction opened and closed without ops
+        ]
+      ]
     end
 
     with_them do
       it do
-        expect(described_class.validate(arguments)).to eq(output)
+        expect(described_class.validate(arguments, true)).to eq(output)
       end
     end
   end
 
   describe '.allow_cross_slot_commands' do
     context 'with feature flags enabled' do
-      before do
-        stub_feature_flags(validate_allowed_cross_slot_commands: true)
-      end
-
       it 'does not raise for invalid arguments' do
         expect(
           described_class.allow_cross_slot_commands do
-            described_class.validate([[:mget, 'foo', 'bar']])
+            described_class.validate([[:mget, 'foo', 'bar']], true)
           end
         ).to eq({ valid: false, key_count: 2, command_name: 'MGET', allowed: true,
                   command: "mget foo bar" })
@@ -97,10 +99,10 @@
         expect(
           described_class.allow_cross_slot_commands do
             described_class.allow_cross_slot_commands do
-              described_class.validate([[:mget, 'foo', 'bar']])
+              described_class.validate([[:mget, 'foo', 'bar']], true)
             end
 
-            described_class.validate([[:mget, 'foo', 'bar']])
+            described_class.validate([[:mget, 'foo', 'bar']], true)
           end
         ).to eq({ valid: false, key_count: 2, command_name: 'MGET', allowed: true,
                   command: "mget foo bar" })
@@ -108,14 +110,10 @@
     end
 
     context 'with feature flag disabled' do
-      before do
-        stub_feature_flags(validate_allowed_cross_slot_commands: false)
-      end
-
       it 'does not run for allowed commands' do
         expect(
           described_class.allow_cross_slot_commands do
-            described_class.validate([[:mget, 'foo', 'bar']])
+            described_class.validate([[:mget, 'foo', 'bar']], false)
           end
         ).to eq(nil)
       end