1 files changed, 70 insertions, 0 deletions
diff --git a/lib/gitlab/metrics/subscribers/load_balancing.rb b/lib/gitlab/metrics/subscribers/load_balancing.rb
new file mode 100644
index 00000000000..333fc63ebef
--- /dev/null
+++ b/lib/gitlab/metrics/subscribers/load_balancing.rb
@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Metrics
+    module Subscribers
+      class LoadBalancing < ActiveSupport::Subscriber
+        attach_to :load_balancing
+
+        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 Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
+
+          result = event.payload[:result]
+          counter_name = counter(result)
+
+          increment(counter_name)
+        end
+
+        # we want to update Prometheus counter after the controller/action are set
+        def web_transaction_completed(_event)
+          return unless Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
+
+          LOG_COUNTERS.keys.each { |result| increment_prometheus_for_result_label(result) }
+        end
+
+        def self.load_balancing_payload
+          return {} unless Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
+
+          {}.tap do |payload|
+            LOG_COUNTERS.values.each do |counter|
+              value = Gitlab::SafeRequestStore[counter]
+
+              payload[counter] = value.to_i if value
+            end
+          end
+        end
+
+        private
+
+        def increment(counter)
+          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 = Gitlab::SafeRequestStore[counter_name])
+
+          increment_prometheus(labels: { result: label_value }, value: counter_value.to_i)
+        end
+
+        def increment_prometheus(labels:, value:)
+          current_transaction&.increment(PROMETHEUS_COUNTER, value, labels) do
+            docstring 'Caught up replica pick result'
+            label_keys labels.keys
+          end
+        end
+
+        def counter(result)
+          LOG_COUNTERS[result]
+        end
+
+        def current_transaction
+          ::Gitlab::Metrics::WebTransaction.current
+        end
+      end
+    end
+  end
+end