1 files changed, 30 insertions, 51 deletions

diff --git a/lib/gitlab/instrumentation/redis.rb b/lib/gitlab/instrumentation/redis.rb
index cc99e828251..82b4701872f 100644
--- a/lib/gitlab/instrumentation/redis.rb
+++ b/lib/gitlab/instrumentation/redis.rb
@@ -1,67 +1,46 @@
 # 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
+      ActionCable = Class.new(RedisBase)
+      Cache = Class.new(RedisBase)
+      Queues = Class.new(RedisBase)
+      SharedState = Class.new(RedisBase)
 
-      def self.get_request_count
-        ::RequestStore[REDIS_REQUEST_COUNT] || 0
-      end
+      STORAGES = [ActionCable, Cache, Queues, SharedState].freeze
 
-      def self.increment_request_count
-        ::RequestStore[REDIS_REQUEST_COUNT] ||= 0
-        ::RequestStore[REDIS_REQUEST_COUNT] += 1
-      end
+      # Milliseconds represented in seconds (from 1 millisecond to 2 seconds).
+      QUERY_TIME_BUCKETS = [0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2].freeze
 
-      def self.detail_store
-        ::RequestStore[REDIS_CALL_DETAILS] ||= []
-      end
+      class << self
+        include ::Gitlab::Instrumentation::RedisPayload
 
-      def self.query_time
-        query_time = ::RequestStore[REDIS_CALL_DURATION] || 0
-        query_time.round(::Gitlab::InstrumentationHelper::DURATION_PRECISION)
-      end
+        def storage_key
+          nil
+        end
 
-      def self.add_duration(duration)
-        total_time = query_time + duration
-        ::RequestStore[REDIS_CALL_DURATION] = total_time
-      end
+        def known_payload_keys
+          super + STORAGES.flat_map(&:known_payload_keys)
+        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
+        def payload
+          super.merge(*STORAGES.flat_map(&:payload))
+        end
 
-        detail_store << {
-          cmd: args.first,
-          duration: duration,
-          backtrace: ::Gitlab::BacktraceCleaner.clean_backtrace(caller)
-        }
+        def detail_store
+          STORAGES.flat_map do |storage|
+            storage.detail_store.map { |details| details.merge(storage: storage.name.demodulize) }
+          end
+        end
+
+        %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
+        end
       end
     end
   end
 end
-
-class ::Redis::Client
-  prepend ::Gitlab::Instrumentation::RedisInterceptor
-end