1 files changed, 110 insertions, 0 deletions
diff --git a/lib/gitlab/utils/usage_data.rb b/lib/gitlab/utils/usage_data.rb
new file mode 100644
index 00000000000..afc4e000977
--- /dev/null
+++ b/lib/gitlab/utils/usage_data.rb
@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+# Usage data utilities
+#
+#   * distinct_count(relation, column = nil, batch: true, start: nil, finish: nil)
+#     Does a distinct batch count, smartly reduces batch_size and handles errors
+#
+#     Examples:
+#     issues_using_zoom_quick_actions: distinct_count(ZoomMeeting, :issue_id),
+#
+#   * count(relation, column = nil, batch: true, start: nil, finish: nil)
+#     Does a non-distinct batch count, smartly reduces batch_size and handles errors
+#
+#     Examples:
+#     active_user_count: count(User.active)
+#
+#   * alt_usage_data method
+#     handles StandardError and fallbacks by default into -1 this way not all measures fail if we encounter one exception
+#     there might be cases where we need to set a specific fallback in order to be aligned wih what version app is expecting as a type
+#
+#     Examples:
+#     alt_usage_data { Gitlab::VERSION }
+#     alt_usage_data { Gitlab::CurrentSettings.uuid }
+#     alt_usage_data(fallback: nil) { Gitlab.config.registry.enabled }
+#
+#   * redis_usage_data method
+#     handles ::Redis::CommandError, Gitlab::UsageDataCounters::BaseCounter::UnknownEvent
+#     returns -1 when a block is sent or hash with all values -1 when a counter is sent
+#     different behaviour due to 2 different implementations of redis counter
+#
+#     Examples:
+#     redis_usage_data(Gitlab::UsageDataCounters::WikiPageCounter)
+#     redis_usage_data { ::Gitlab::UsageCounters::PodLogs.usage_totals[:total] }
+
+module Gitlab
+  module Utils
+    module UsageData
+      extend self
+
+      FALLBACK = -1
+
+      def count(relation, column = nil, batch: true, start: nil, finish: nil)
+        if batch
+          Gitlab::Database::BatchCount.batch_count(relation, column, start: start, finish: finish)
+        else
+          relation.count
+        end
+      rescue ActiveRecord::StatementInvalid
+        FALLBACK
+      end
+
+      def distinct_count(relation, column = nil, batch: true, batch_size: nil, start: nil, finish: nil)
+        if batch
+          Gitlab::Database::BatchCount.batch_distinct_count(relation, column, batch_size: batch_size, start: start, finish: finish)
+        else
+          relation.distinct_count_by(column)
+        end
+      rescue ActiveRecord::StatementInvalid
+        FALLBACK
+      end
+
+      def alt_usage_data(value = nil, fallback: FALLBACK, &block)
+        if block_given?
+          yield
+        else
+          value
+        end
+      rescue
+        fallback
+      end
+
+      def redis_usage_data(counter = nil, &block)
+        if block_given?
+          redis_usage_counter(&block)
+        elsif counter.present?
+          redis_usage_data_totals(counter)
+        end
+      end
+
+      def with_prometheus_client
+        if Gitlab::Prometheus::Internal.prometheus_enabled?
+          prometheus_address = Gitlab::Prometheus::Internal.uri
+          yield Gitlab::PrometheusClient.new(prometheus_address, allow_local_requests: true)
+        end
+      end
+
+      def measure_duration
+        result = nil
+        duration = Benchmark.realtime do
+          result = yield
+        end
+        [result, duration]
+      end
+
+      private
+
+      def redis_usage_counter
+        yield
+      rescue ::Redis::CommandError, Gitlab::UsageDataCounters::BaseCounter::UnknownEvent
+        FALLBACK
+      end
+
+      def redis_usage_data_totals(counter)
+        counter.totals
+      rescue ::Redis::CommandError, Gitlab::UsageDataCounters::BaseCounter::UnknownEvent
+        counter.fallback_totals
+      end
+    end
+  end
+end