1 files changed, 6 insertions, 8 deletions
diff --git a/lib/gitlab/metrics/subscribers/load_balancing.rb b/lib/gitlab/metrics/subscribers/load_balancing.rb
index d7fe33dbe89..bd77e8c3c3f 100644
--- a/lib/gitlab/metrics/subscribers/load_balancing.rb
+++ b/lib/gitlab/metrics/subscribers/load_balancing.rb
@@ -6,13 +6,11 @@ module Gitlab
       class LoadBalancing < ActiveSupport::Subscriber
         attach_to :load_balancing
 
-        InstrumentationStorage = ::Gitlab::Instrumentation::Storage
-
         PROMETHEUS_COUNTER = :gitlab_transaction_caught_up_replica_pick_count_total
         LOG_COUNTERS = { true => :caught_up_replica_pick_ok, false => :caught_up_replica_pick_fail }.freeze
 
         def caught_up_replica_pick(event)
-          return unless InstrumentationStorage.active?
+          return unless Gitlab::SafeRequestStore.active?
 
           result = event.payload[:result]
           counter_name = counter(result)
@@ -22,17 +20,17 @@ module Gitlab
 
         # we want to update Prometheus counter after the controller/action are set
         def web_transaction_completed(_event)
-          return unless InstrumentationStorage.active?
+          return unless Gitlab::SafeRequestStore.active?
 
           LOG_COUNTERS.keys.each { |result| increment_prometheus_for_result_label(result) }
         end
 
         def self.load_balancing_payload
-          return {} unless InstrumentationStorage.active?
+          return {} unless Gitlab::SafeRequestStore.active?
 
           {}.tap do |payload|
             LOG_COUNTERS.values.each do |counter|
-              value = InstrumentationStorage[counter]
+              value = Gitlab::SafeRequestStore[counter]
 
               payload[counter] = value.to_i if value
             end
@@ -42,12 +40,12 @@ module Gitlab
         private
 
         def increment(counter)
-          InstrumentationStorage[counter] = InstrumentationStorage[counter].to_i + 1
+          Gitlab::SafeRequestStore[counter] = Gitlab::SafeRequestStore[counter].to_i + 1
         end
 
         def increment_prometheus_for_result_label(label_value)
           counter_name = counter(label_value)
-          return unless (counter_value = InstrumentationStorage[counter_name])
+          return unless (counter_value = Gitlab::SafeRequestStore[counter_name])
 
           increment_prometheus(labels: { result: label_value }, value: counter_value.to_i)
         end