Separate Redis instrumentation per client

This is the groundwork required for a series
of improvements on Redis usage visualization.

It separates the instrumentation storages for each
Redis client being used. These being, Cache, Shared State,
Queues and Action Cable.

That means that these changes will make the following achievable:

- Separate Redis logs per Redis instance
- Categorize the Redis instance in Prometheus metrics
- Separate Redis performance bar data per Redis instance

config/initializers/action_cable.rb

 # frozen_string_literal: true
 
+require 'action_cable/subscription_adapter/redis'
+
 Rails.application.configure do
   # We only mount the ActionCable engine in tests where we run it in-app
   # For other environments, we run it on a standalone Puma server
@@ -7,3 +9,11 @@ Rails.application.configure do
   config.action_cable.url = Gitlab::Utils.append_path(Gitlab.config.gitlab.relative_url_root, '/-/cable')
   config.action_cable.worker_pool_size = Gitlab.config.action_cable.worker_pool_size
 end
+
+# https://github.com/rails/rails/blob/bb5ac1623e8de08c1b7b62b1368758f0d3bb6379/actioncable/lib/action_cable/subscription_adapter/redis.rb#L18
+ActionCable::SubscriptionAdapter::Redis.redis_connector = lambda do |config|
+  args = config.except(:adapter, :channel_prefix)
+    .merge(instrumentation_class: ::Gitlab::Instrumentation::Redis::ActionCable)
+
+  ::Redis.new(args)
+end

lib/gitlab/instrumentation/redis.rb

 # frozen_string_literal: true
 
-require 'redis'
-
 module Gitlab
   module Instrumentation
-    module RedisInterceptor
-      def call(*args, &block)
-        start = Time.now
-        super(*args, &block)
-      ensure
-        duration = (Time.now - start)
-
-        if ::RequestStore.active?
-          ::Gitlab::Instrumentation::Redis.increment_request_count
-          ::Gitlab::Instrumentation::Redis.add_duration(duration)
-          ::Gitlab::Instrumentation::Redis.add_call_details(duration, args)
-        end
-      end
-    end
-
+    # Aggregates Redis measurements from different request storage sources.
     class Redis
-      REDIS_REQUEST_COUNT = :redis_request_count
-      REDIS_CALL_DURATION = :redis_call_duration
-      REDIS_CALL_DETAILS = :redis_call_details
-      REDIS_READ_BYTES = :redis_read_bytes
-      REDIS_WRITE_BYTES = :redis_write_bytes
+      ActionCable = Class.new(RedisBase)
+      Cache = Class.new(RedisBase)
+      Queues = Class.new(RedisBase)
+      SharedState = Class.new(RedisBase)
+
+      STORAGES = [ActionCable, Cache, Queues, SharedState].freeze
 
       # Milliseconds represented in seconds (from 1 to 500 milliseconds).
       QUERY_TIME_BUCKETS = [0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5].freeze
 
-      def self.get_request_count
-        ::RequestStore[REDIS_REQUEST_COUNT] || 0
-      end
-
-      def self.increment_request_count
-        ::RequestStore[REDIS_REQUEST_COUNT] ||= 0
-        ::RequestStore[REDIS_REQUEST_COUNT] += 1
-      end
-
-      def self.increment_read_bytes(num_bytes)
-        ::RequestStore[REDIS_READ_BYTES] ||= 0
-        ::RequestStore[REDIS_READ_BYTES] += num_bytes
-      end
-
-      def self.increment_write_bytes(num_bytes)
-        ::RequestStore[REDIS_WRITE_BYTES] ||= 0
-        ::RequestStore[REDIS_WRITE_BYTES] += num_bytes
-      end
-
-      def self.read_bytes
-        ::RequestStore[REDIS_READ_BYTES] || 0
-      end
-
-      def self.write_bytes
-        ::RequestStore[REDIS_WRITE_BYTES] || 0
-      end
-
-      def self.detail_store
-        ::RequestStore[REDIS_CALL_DETAILS] ||= []
+      class << self
+        def detail_store
+          STORAGES.flat_map(&:detail_store)
         end
 
-      def self.query_time
-        query_time = ::RequestStore[REDIS_CALL_DURATION] || 0
-        query_time.round(::Gitlab::InstrumentationHelper::DURATION_PRECISION)
+        %i[get_request_count query_time read_bytes write_bytes].each do |method|
+          define_method method do
+            STORAGES.sum(&method) # rubocop:disable CodeReuse/ActiveRecord
           end
-
-      def self.add_duration(duration)
-        ::RequestStore[REDIS_CALL_DURATION] ||= 0
-        ::RequestStore[REDIS_CALL_DURATION] += duration
         end
-
-      def self.add_call_details(duration, args)
-        return unless Gitlab::PerformanceBar.enabled_for_request?
-        # redis-rb passes an array (e.g. [:get, key])
-        return unless args.length == 1
-
-        detail_store << {
-          cmd: args.first,
-          duration: duration,
-          backtrace: ::Gitlab::BacktraceCleaner.clean_backtrace(caller)
-        }
       end
     end
   end
 end
-
-class ::Redis::Client
-  prepend ::Gitlab::Instrumentation::RedisInterceptor
-end

lib/gitlab/instrumentation/redis_base.rb

+# frozen_string_literal: true
+
+require 'redis'
+
+module Gitlab
+  module Instrumentation
+    class RedisBase
+      class << self
+        include ::Gitlab::Utils::StrongMemoize
+
+        # TODO: To be used by https://gitlab.com/gitlab-com/gl-infra/scalability/-/issues/395
+        # as a 'label' alias.
+        def storage_key
+          self.name.underscore
+        end
+
+        def add_duration(duration)
+          ::RequestStore[call_duration_key] ||= 0
+          ::RequestStore[call_duration_key] += duration
+        end
+
+        def add_call_details(duration, args)
+          return unless Gitlab::PerformanceBar.enabled_for_request?
+          # redis-rb passes an array (e.g. [[:get, key]])
+          return unless args.length == 1
+
+          # TODO: Add information about current Redis client
+          # being instrumented.
+          # https://gitlab.com/gitlab-com/gl-infra/scalability/-/issues/316.
+          detail_store << {
+            cmd: args.first,
+            duration: duration,
+            backtrace: ::Gitlab::BacktraceCleaner.clean_backtrace(caller)
+          }
+        end
+
+        def increment_request_count
+          ::RequestStore[request_count_key] ||= 0
+          ::RequestStore[request_count_key] += 1
+        end
+
+        def increment_read_bytes(num_bytes)
+          ::RequestStore[read_bytes_key] ||= 0
+          ::RequestStore[read_bytes_key] += num_bytes
+        end
+
+        def increment_write_bytes(num_bytes)
+          ::RequestStore[write_bytes_key] ||= 0
+          ::RequestStore[write_bytes_key] += num_bytes
+        end
+
+        def get_request_count
+          ::RequestStore[request_count_key] || 0
+        end
+
+        def read_bytes
+          ::RequestStore[read_bytes_key] || 0
+        end
+
+        def write_bytes
+          ::RequestStore[write_bytes_key] || 0
+        end
+
+        def detail_store
+          ::RequestStore[call_details_key] ||= []
+        end
+
+        def query_time
+          query_time = ::RequestStore[call_duration_key] || 0
+          query_time.round(::Gitlab::InstrumentationHelper::DURATION_PRECISION)
+        end
+
+        private
+
+        def request_count_key
+          strong_memoize(:request_count_key) { build_key(:redis_request_count) }
+        end
+
+        def read_bytes_key
+          strong_memoize(:read_bytes_key) { build_key(:redis_read_bytes) }
+        end
+
+        def write_bytes_key
+          strong_memoize(:write_bytes_key) { build_key(:redis_write_bytes) }
+        end
+
+        def call_duration_key
+          strong_memoize(:call_duration_key) { build_key(:redis_call_duration) }
+        end
+
+        def call_details_key
+          strong_memoize(:call_details_key) { build_key(:redis_call_details) }
+        end
+
+        def build_key(namespace)
+          "#{storage_key}_#{namespace}"
+        end
+      end
+    end
+  end
+end

lib/gitlab/instrumentation/redis_driver.rb → lib/gitlab/instrumentation/redis_interceptor.rb

@@ -4,7 +4,20 @@ require 'redis'
 
 module Gitlab
   module Instrumentation
-    class RedisDriver < ::Redis::Connection::Ruby
+    module RedisInterceptor
+      def call(*args, &block)
+        start = Time.now
+        super(*args, &block)
+      ensure
+        duration = (Time.now - start)
+
+        if ::RequestStore.active?
+          instrumentation_class.increment_request_count
+          instrumentation_class.add_duration(duration)
+          instrumentation_class.add_call_details(duration, args)
+        end
+      end
+
       def write(command)
         measure_write_size(command) if ::RequestStore.active?
         super
@@ -35,27 +48,40 @@ module Gitlab
           end
         end
 
-        ::Gitlab::Instrumentation::Redis.increment_write_bytes(size)
+        instrumentation_class.increment_write_bytes(size)
       end
 
       def measure_read_size(result)
-        # The superclass can return one of four types of results from read:
+        # The Connection::Ruby#read class can return one of four types of results from read:
         # https://github.com/redis/redis-rb/blob/f597f21a6b954b685cf939febbc638f6c803e3a7/lib/redis/connection/ruby.rb#L406
         #
         # 1. Error (exception, will not reach this line)
         # 2. Status (string)
         # 3. Integer (will be converted to string by to_s.bytesize and thrown away)
         # 4. "Binary" string (i.e. may contain zero byte)
-        # 5. Array of binary string (recurses back into read)
+        # 5. Array of binary string
 
-        # Avoid double-counting array responses: the array elements themselves
-        # are retrieved with 'read'.
-        unless result.is_a? Array
+        size = if result.is_a? Array
                  # This count is an approximation that omits the Redis protocol overhead
                  # of type prefixes, length prefixes and line endings.
-          ::Gitlab::Instrumentation::Redis.increment_read_bytes(result.to_s.bytesize)
+                 result.inject(0) { |sum, y| sum + y.to_s.bytesize }
+               else
+                 result.to_s.bytesize
                end
+
+        instrumentation_class.increment_read_bytes(size)
+      end
+
+      # That's required so it knows which GitLab Redis instance
+      # it's interacting with in order to categorize accordingly.
+      #
+      def instrumentation_class
+        @options[:instrumentation_class] # rubocop:disable Gitlab/ModuleWithInstanceVariables
       end
     end
   end
 end
+
+class ::Redis::Client
+  prepend ::Gitlab::Instrumentation::RedisInterceptor
+end

lib/gitlab/redis/cache.rb

@@ -27,6 +27,10 @@ module Gitlab
           # this will force use of DEFAULT_REDIS_QUEUES_URL when config file is absent
           super
         end
+
+        def instrumentation_class
+          ::Gitlab::Instrumentation::Redis::Cache
+        end
       end
     end
   end

lib/gitlab/redis/queues.rb

@@ -28,6 +28,10 @@ module Gitlab
           # this will force use of DEFAULT_REDIS_QUEUES_URL when config file is absent
           super
         end
+
+        def instrumentation_class
+          ::Gitlab::Instrumentation::Redis::Queues
+        end
       end
     end
   end

lib/gitlab/redis/shared_state.rb

@@ -30,6 +30,10 @@ module Gitlab
           # this will force use of DEFAULT_REDIS_SHARED_STATE_URL when config file is absent
           super
         end
+
+        def instrumentation_class
+          ::Gitlab::Instrumentation::Redis::SharedState
+        end
       end
     end
   end

lib/gitlab/redis/wrapper.rb

@@ -71,6 +71,10 @@ module Gitlab
           # nil will force use of DEFAULT_REDIS_URL when config file is absent
           nil
         end
+
+        def instrumentation_class
+          raise NotImplementedError
+        end
       end
 
       def initialize(rails_env = nil)
@@ -100,7 +104,7 @@ module Gitlab
         redis_url = config.delete(:url)
         redis_uri = URI.parse(redis_url)
 
-        config[:driver] ||= ::Gitlab::Instrumentation::RedisDriver
+        config[:instrumentation_class] ||= self.class.instrumentation_class
 
         if redis_uri.scheme == 'unix'
           # Redis::Store does not handle Unix sockets well, so let's do it for them

lib/peek/views/redis_detailed.rb

@@ -9,6 +9,10 @@ module Peek
         'redis'
       end
 
+      def detail_store
+        ::Gitlab::Instrumentation::Redis.detail_store
+      end
+
       private
 
       def format_call_details(call)

spec/lib/gitlab/instrumentation/redis_base_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Instrumentation::RedisBase, :request_store do
+  let(:instrumentation_class_a) do
+    stub_const('RedisInstanceA', Class.new(described_class))
+  end
+
+  let(:instrumentation_class_b) do
+    stub_const('RedisInstanceB', Class.new(described_class))
+  end
+
+  describe '.storage_key' do
+    it 'returns the class name with underscore' do
+      expect(instrumentation_class_a.storage_key).to eq('redis_instance_a')
+      expect(instrumentation_class_b.storage_key).to eq('redis_instance_b')
+    end
+  end
+
+  describe '.add_duration' do
+    it 'does not lose precision while adding' do
+      precision = 1.0 / (10**::Gitlab::InstrumentationHelper::DURATION_PRECISION)
+      2.times { instrumentation_class_a.add_duration(0.4 * precision) }
+
+      # 2 * 0.4 should be 0.8 and get rounded to 1
+      expect(instrumentation_class_a.query_time).to eq(1 * precision)
+    end
+
+    context 'storage key overlapping' do
+      it 'keys do not overlap across storages' do
+        instrumentation_class_a.add_duration(0.4)
+        instrumentation_class_b.add_duration(0.5)
+
+        expect(instrumentation_class_a.query_time).to eq(0.4)
+        expect(instrumentation_class_b.query_time).to eq(0.5)
+      end
+    end
+  end
+
+  describe '.increment_request_count' do
+    context 'storage key overlapping' do
+      it 'keys do not overlap across storages' do
+        3.times { instrumentation_class_a.increment_request_count }
+        2.times { instrumentation_class_b.increment_request_count }
+
+        expect(instrumentation_class_a.get_request_count).to eq(3)
+        expect(instrumentation_class_b.get_request_count).to eq(2)
+      end
+    end
+  end
+
+  describe '.increment_write_bytes' do
+    context 'storage key overlapping' do
+      it 'keys do not overlap across storages' do
+        2.times do
+          instrumentation_class_a.increment_write_bytes(42)
+          instrumentation_class_b.increment_write_bytes(77)
+        end
+
+        expect(instrumentation_class_a.write_bytes).to eq(42 * 2)
+        expect(instrumentation_class_b.write_bytes).to eq(77 * 2)
+      end
+    end
+  end
+
+  describe '.increment_read_bytes' do
+    context 'storage key overlapping' do
+      it 'keys do not overlap across storages' do
+        2.times do
+          instrumentation_class_a.increment_read_bytes(42)
+          instrumentation_class_b.increment_read_bytes(77)
+        end
+
+        expect(instrumentation_class_a.read_bytes).to eq(42 * 2)
+        expect(instrumentation_class_b.read_bytes).to eq(77 * 2)
+      end
+    end
+  end
+
+  describe '.add_call_details' do
+    before do
+      allow(Gitlab::PerformanceBar).to receive(:enabled_for_request?) { true }
+    end
+
+    context 'storage key overlapping' do
+      it 'keys do not overlap across storages' do
+        2.times do
+          instrumentation_class_a.add_call_details(0.3, [:set])
+          instrumentation_class_b.add_call_details(0.4, [:set])
+        end
+
+        expect(instrumentation_class_a.detail_store).to match(
+          [
+            a_hash_including(cmd: :set, duration: 0.3, backtrace: an_instance_of(Array)),
+            a_hash_including(cmd: :set, duration: 0.3, backtrace: an_instance_of(Array))
+          ]
+        )
+
+        expect(instrumentation_class_b.detail_store).to match(
+          [
+            a_hash_including(cmd: :set, duration: 0.4, backtrace: an_instance_of(Array)),
+            a_hash_including(cmd: :set, duration: 0.4, backtrace: an_instance_of(Array))
+          ]
+        )
+      end
+    end
+  end
+end

spec/lib/gitlab/instrumentation/redis_driver_spec.rb → spec/lib/gitlab/instrumentation/redis_interceptor_spec.rb

@@ -3,7 +3,7 @@
 require 'spec_helper'
 require 'rspec-parameterized'
 
-describe Gitlab::Instrumentation::RedisDriver, :clean_gitlab_redis_shared_state, :request_store do
+describe Gitlab::Instrumentation::RedisInterceptor, :clean_gitlab_redis_shared_state, :request_store do
   using RSpec::Parameterized::TableSyntax
 
   describe 'read and write' do

spec/lib/gitlab/instrumentation/redis_spec.rb

@@ -3,13 +3,26 @@
 require 'spec_helper'
 
 describe Gitlab::Instrumentation::Redis do
-  describe '.add_duration', :request_store do
-    it 'does not lose precision while adding' do
-      precision = 1.0 / (10**::Gitlab::InstrumentationHelper::DURATION_PRECISION)
-      2.times { described_class.add_duration(0.4 * precision) }
+  def stub_storages(method, value)
+    described_class::STORAGES.each do |storage|
+      allow(storage).to receive(method) { value }
+    end
+  end
+
+  shared_examples 'aggregation of redis storage data' do |method|
+    describe "#{method} sum" do
+      it "sums data from all Redis storages" do
+        amount = 0.3
 
-      # 2 * 0.4 should be 0.8 and get rounded to 1
-      expect(described_class.query_time).to eq(1 * precision)
+        stub_storages(method, amount)
+
+        expect(described_class.public_send(method)).to eq(described_class::STORAGES.size * amount)
+      end
     end
   end
+
+  it_behaves_like 'aggregation of redis storage data', :get_request_count
+  it_behaves_like 'aggregation of redis storage data', :query_time
+  it_behaves_like 'aggregation of redis storage data', :read_bytes
+  it_behaves_like 'aggregation of redis storage data', :write_bytes
 end

spec/lib/gitlab/redis/wrapper_spec.rb

@@ -18,7 +18,21 @@ describe Gitlab::Redis::Wrapper do
   let(:config_env_variable_url) {"TEST_GITLAB_REDIS_URL"}
   let(:class_redis_url) { Gitlab::Redis::Wrapper::DEFAULT_REDIS_URL }
 
-  include_examples "redis_shared_examples"
+  include_examples "redis_shared_examples" do
+    before do
+      allow(described_class).to receive(:instrumentation_class) do
+        ::Gitlab::Instrumentation::Redis::Cache
+      end
+    end
+  end
+
+  describe '.instrumentation_class' do
+    it 'raises a NotImplementedError' do
+      expect(described_class).to receive(:instrumentation_class).and_call_original
+
+      expect { described_class.instrumentation_class }.to raise_error(NotImplementedError)
+    end
+  end
 
   describe '.config_file_path' do
     it 'returns the absolute path to the configuration file' do

spec/lib/peek/views/redis_detailed_spec.rb

@@ -17,7 +17,7 @@ describe Peek::Views::RedisDetailed, :request_store do
 
   with_them do
     it 'scrubs Redis commands' do
-      subject.detail_store << { cmd: cmd, duration: 1.second }
+      Gitlab::Instrumentation::Redis::SharedState.detail_store << { cmd: cmd, duration: 1.second }
 
       expect(subject.results[:details].count).to eq(1)
       expect(subject.results[:details].first)
@@ -29,11 +29,12 @@ describe Peek::Views::RedisDetailed, :request_store do
   end
 
   it 'returns aggregated results' do
-    subject.detail_store << { cmd: [:get, 'test'], duration: 0.001 }
-    subject.detail_store << { cmd: [:get, 'test'], duration: 1.second }
+    Gitlab::Instrumentation::Redis::Cache.detail_store << { cmd: [:get, 'test'], duration: 0.001 }
+    Gitlab::Instrumentation::Redis::Cache.detail_store << { cmd: [:get, 'test'], duration: 1.second }
+    Gitlab::Instrumentation::Redis::SharedState.detail_store << { cmd: [:get, 'test'], duration: 1.second }
 
-    expect(subject.results[:calls]).to eq(2)
-    expect(subject.results[:duration]).to eq('1001.00ms')
-    expect(subject.results[:details].count).to eq(2)
+    expect(subject.results[:calls]).to eq(3)
+    expect(subject.results[:duration]).to eq('2001.00ms')
+    expect(subject.results[:details].count).to eq(3)
   end
 end