Fix usage ping timeouts with batch counters

Usage pings time out on large instances with large tables. This change
adds batch counters to fix the issue

changelogs/unreleased/193075-usage-ping-timing-out-for-larger-instances-batch.yml

+---
+title: Fix usage ping timeouts with batch counters
+merge_request: 22705
+author:
+type: performance

ee/lib/ee/gitlab/usage_data.rb

@@ -19,7 +19,8 @@ module EE
           # to time out on GitLab.com. Since we are mostly interested in gathering these statistics for
           # self hosted instances, prevent them from running on GitLab.com and allow instance maintainers
           # to disable them via a feature flag.
-          return super if ::Gitlab.com? || ::Feature.disabled?(:usage_activity_by_stage, default_enabled: true)
+          return super if (::Feature.disabled?(:usage_ping_batch_counter) && ::Gitlab.com?) ||
+                          ::Feature.disabled?(:usage_activity_by_stage, default_enabled: true)
 
           super.merge(usage_activity_by_stage)
         end
@@ -94,7 +95,7 @@ module EE
             sast: :sast_jobs
           }
 
-          results = count(::Ci::Build.where(name: types.keys).group(:name), fallback: Hash.new(-1))
+          results = count(::Ci::Build.where(name: types.keys).group(:name), fallback: Hash.new(-1), batch: false)
 
           license_scan_count = results.delete("license_scanning")
           if license_scan_count && results["license_management"]
@@ -119,7 +120,7 @@ module EE
 
         def operations_dashboard_usage
           users_with_ops_dashboard_as_default = count(::User.active.with_dashboard('operations'))
-          users_with_projects_added = count(UsersOpsDashboardProject.distinct_users(::User.active))
+          users_with_projects_added = count(UsersOpsDashboardProject.distinct_users(::User.active), batch: false)
 
           {
             operations_dashboard_default_dashboard: users_with_ops_dashboard_as_default,
@@ -137,15 +138,15 @@ module EE
                                          geo_nodes: count(::GeoNode),
                                          ldap_group_links: count(::LdapGroupLink),
                                          ldap_keys: count(::LDAPKey),
-                                         ldap_users: count(::User.ldap),
+                                         ldap_users: count(::User.ldap, 'users.id'),
                                          pod_logs_usages_total: ::Gitlab::UsageCounters::PodLogs.usage_totals[:total],
-                                         projects_enforcing_code_owner_approval: count(::Project.without_deleted.non_archived.requiring_code_owner_approval),
-                                         projects_mirrored_with_pipelines_enabled: count(::Project.mirrored_with_enabled_pipelines),
-                                         projects_reporting_ci_cd_back_to_github: count(::GithubService.without_defaults.active),
-                                         projects_with_packages: count(::Packages::Package.select('distinct project_id')),
-                                         projects_with_prometheus_alerts: count(PrometheusAlert.distinct_projects),
-                                         projects_with_tracing_enabled: count(ProjectTracingSetting),
-                                         template_repositories:  count(::Project.with_repos_templates) + count(::Project.with_groups_level_repos_templates)
+                                         projects_enforcing_code_owner_approval: count(::Project.without_deleted.non_archived.requiring_code_owner_approval, batch: false),
+                                         projects_mirrored_with_pipelines_enabled: count(::Project.mirrored_with_enabled_pipelines, batch: false),
+                                         projects_reporting_ci_cd_back_to_github: count(::GithubService.without_defaults.active, batch: false),
+                                         projects_with_packages: count(::Packages::Package.select('distinct project_id'), batch: false),
+                                         projects_with_prometheus_alerts: count(PrometheusAlert.distinct_projects, batch: false),
+                                         projects_with_tracing_enabled: count(ProjectTracingSetting, batch: false),
+                                         template_repositories:  count(::Project.with_repos_templates, batch: false) + count(::Project.with_groups_level_repos_templates, batch: false)
                                        },
                                        service_desk_counts,
                                        security_products_usage,
@@ -191,62 +192,62 @@ module EE
             clusters_applications_helm: ::Clusters::Applications::Helm.distinct_by_user,
             clusters_applications_ingress: ::Clusters::Applications::Ingress.distinct_by_user,
             clusters_applications_knative: ::Clusters::Applications::Knative.distinct_by_user,
-            clusters_disabled: ::Clusters::Cluster.disabled.distinct_count_by(:user_id),
-            clusters_enabled: ::Clusters::Cluster.enabled.distinct_count_by(:user_id),
-            clusters_platforms_gke: ::Clusters::Cluster.gcp_installed.enabled.distinct_count_by(:user_id),
-            clusters_platforms_eks: ::Clusters::Cluster.aws_installed.enabled.distinct_count_by(:user_id),
-            clusters_platforms_user: ::Clusters::Cluster.user_provided.enabled.distinct_count_by(:user_id),
-            group_clusters_disabled: ::Clusters::Cluster.disabled.group_type.distinct_count_by(:user_id),
-            group_clusters_enabled: ::Clusters::Cluster.enabled.group_type.distinct_count_by(:user_id),
-            project_clusters_disabled: ::Clusters::Cluster.disabled.project_type.distinct_count_by(:user_id),
-            project_clusters_enabled: ::Clusters::Cluster.enabled.project_type.distinct_count_by(:user_id),
-            projects_slack_notifications_active: ::Project.with_slack_service.distinct_count_by(:creator_id),
-            projects_slack_slash_active: ::Project.with_slack_slash_commands_service.distinct_count_by(:creator_id),
-            projects_with_prometheus_alerts: ::Project.with_prometheus_service.distinct_count_by(:creator_id)
+            clusters_disabled: distinct_count(::Clusters::Cluster.disabled, :user_id),
+            clusters_enabled: distinct_count(::Clusters::Cluster.enabled, :user_id),
+            clusters_platforms_gke: distinct_count(::Clusters::Cluster.gcp_installed.enabled, :user_id),
+            clusters_platforms_eks: distinct_count(::Clusters::Cluster.aws_installed.enabled, :user_id),
+            clusters_platforms_user: distinct_count(::Clusters::Cluster.user_provided.enabled, :user_id),
+            group_clusters_disabled: distinct_count(::Clusters::Cluster.disabled.group_type, :user_id),
+            group_clusters_enabled: distinct_count(::Clusters::Cluster.enabled.group_type, :user_id),
+            project_clusters_disabled: distinct_count(::Clusters::Cluster.disabled.project_type, :user_id),
+            project_clusters_enabled: distinct_count(::Clusters::Cluster.enabled.project_type, :user_id),
+            projects_slack_notifications_active: distinct_count(::Project.with_slack_service, :creator_id),
+            projects_slack_slash_active: distinct_count(::Project.with_slack_slash_commands_service, :creator_id),
+            projects_with_prometheus_alerts: distinct_count(::Project.with_prometheus_service, :creator_id)
           }
         end
 
         # Omitted because no user, creator or author associated: `lfs_objects`, `pool_repositories`, `web_hooks`
         def usage_activity_by_stage_create
           {
-            deploy_keys: ::DeployKey.distinct_count_by(:user_id),
-            keys: ::Key.regular_keys.distinct_count_by(:user_id),
-            merge_requests: ::MergeRequest.distinct_count_by(:author_id),
-            projects_enforcing_code_owner_approval: ::Project.requiring_code_owner_approval.distinct_count_by(:creator_id),
-            projects_imported_from_github: ::Project.github_imported.distinct_count_by(:creator_id),
-            projects_with_repositories_enabled: ::Project.with_repositories_enabled.distinct_count_by(:creator_id),
-            protected_branches: ::Project.with_protected_branches.distinct_count_by(:creator_id),
-            remote_mirrors: ::Project.with_remote_mirrors.distinct_count_by(:creator_id),
-            snippets: ::Snippet.distinct_count_by(:author_id),
-            suggestions: ::Note.with_suggestions.distinct_count_by(:author_id)
+            deploy_keys: distinct_count(::DeployKey, :user_id),
+            keys: distinct_count(::Key.regular_keys, :user_id),
+            merge_requests: distinct_count(::MergeRequest, :author_id),
+            projects_enforcing_code_owner_approval: distinct_count(::Project.requiring_code_owner_approval, :creator_id),
+            projects_imported_from_github: distinct_count(::Project.github_imported, :creator_id),
+            projects_with_repositories_enabled: distinct_count(::Project.with_repositories_enabled, :creator_id),
+            protected_branches: distinct_count(::Project.with_protected_branches, :creator_id),
+            remote_mirrors: distinct_count(::Project.with_remote_mirrors, :creator_id),
+            snippets: distinct_count(::Snippet, :author_id),
+            suggestions: distinct_count(::Note.with_suggestions, :author_id)
           }
         end
 
         # Omitted because no user, creator or author associated: `campaigns_imported_from_github`, `ldap_group_links`
         def usage_activity_by_stage_manage
           {
-            events: ::Event.distinct_count_by(:author_id),
-            groups: ::GroupMember.distinct_count_by(:user_id),
-            ldap_keys: ::LDAPKey.distinct_count_by(:user_id),
-            ldap_users: ::GroupMember.of_ldap_type.distinct_count_by(:user_id)
+            events: distinct_count(::Event, :author_id),
+            groups: distinct_count(::GroupMember, :user_id),
+            ldap_keys: distinct_count(::LDAPKey, :user_id),
+            ldap_users: distinct_count(::GroupMember.of_ldap_type, :user_id)
           }
         end
 
         def usage_activity_by_stage_monitor
           {
-            clusters: ::Clusters::Cluster.distinct_count_by(:user_id),
+            clusters: distinct_count(::Clusters::Cluster, :user_id),
             clusters_applications_prometheus: ::Clusters::Applications::Prometheus.distinct_by_user,
             operations_dashboard_default_dashboard: count(::User.active.with_dashboard('operations')),
-            operations_dashboard_users_with_projects_added: count(UsersOpsDashboardProject.distinct_users(::User.active)),
-            projects_prometheus_active: ::Project.with_active_prometheus_service.distinct_count_by(:creator_id),
-            projects_with_error_tracking_enabled: ::Project.with_enabled_error_tracking.distinct_count_by(:creator_id),
-            projects_with_tracing_enabled: ::Project.with_tracing_enabled.distinct_count_by(:creator_id)
+            operations_dashboard_users_with_projects_added: count(UsersOpsDashboardProject.distinct_users(::User.active), batch: false),
+            projects_prometheus_active: distinct_count(::Project.with_active_prometheus_service, :creator_id),
+            projects_with_error_tracking_enabled: distinct_count(::Project.with_enabled_error_tracking, :creator_id),
+            projects_with_tracing_enabled: distinct_count(::Project.with_tracing_enabled, :creator_id)
           }
         end
 
         def usage_activity_by_stage_package
           {
-            projects_with_packages: ::Project.with_packages.distinct_count_by(:creator_id)
+            projects_with_packages: distinct_count(::Project.with_packages, :creator_id)
           }
         end
 
@@ -255,46 +256,46 @@ module EE
         # Omitted because of encrypted properties: `projects_jira_cloud_active`, `projects_jira_server_active`
         def usage_activity_by_stage_plan
           {
-            assignee_lists: ::List.assignee.distinct_count_by(:user_id),
-            epics: ::Epic.distinct_count_by(:author_id),
-            issues: ::Issue.distinct_count_by(:author_id),
-            label_lists: ::List.label.distinct_count_by(:user_id),
-            milestone_lists: ::List.milestone.distinct_count_by(:user_id),
-            notes: ::Note.distinct_count_by(:author_id),
-            projects: ::Project.distinct_count_by(:creator_id),
-            projects_jira_active: ::Project.with_active_jira_services.distinct_count_by(:creator_id),
-            projects_jira_dvcs_cloud_active: ::Project.with_active_jira_services.with_jira_dvcs_cloud.distinct_count_by(:creator_id),
-            projects_jira_dvcs_server_active: ::Project.with_active_jira_services.with_jira_dvcs_server.distinct_count_by(:creator_id),
-            service_desk_enabled_projects: ::Project.with_active_services.service_desk_enabled.distinct_count_by(:creator_id),
-            service_desk_issues: ::Issue.service_desk.distinct_count_by,
-            todos: ::Todo.distinct_count_by(:author_id)
+            assignee_lists: distinct_count(::List.assignee, :user_id),
+            epics: distinct_count(::Epic, :author_id),
+            issues: distinct_count(::Issue, :author_id),
+            label_lists: distinct_count(::List.label, :user_id),
+            milestone_lists: distinct_count(::List.milestone, :user_id),
+            notes: distinct_count(::Note, :author_id),
+            projects: distinct_count(::Project, :creator_id),
+            projects_jira_active: distinct_count(::Project.with_active_jira_services, :creator_id),
+            projects_jira_dvcs_cloud_active: distinct_count(::Project.with_active_jira_services.with_jira_dvcs_cloud, :creator_id),
+            projects_jira_dvcs_server_active: distinct_count(::Project.with_active_jira_services.with_jira_dvcs_server, :creator_id),
+            service_desk_enabled_projects: distinct_count(::Project.with_active_services.service_desk_enabled, :creator_id),
+            service_desk_issues: distinct_count(::Issue.service_desk),
+            todos: distinct_count(::Todo, :author_id)
           }
         end
 
         # Omitted because no user, creator or author associated: `environments`, `feature_flags`, `in_review_folder`, `pages_domains`
         def usage_activity_by_stage_release
           {
-            deployments: ::Deployment.distinct_count_by(:user_id),
-            failed_deployments: ::Deployment.failed.distinct_count_by(:user_id),
-            projects_mirrored_with_pipelines_enabled: ::Project.mirrored_with_enabled_pipelines.distinct_count_by(:creator_id),
-            releases: ::Release.distinct_count_by(:author_id),
-            successful_deployments: ::Deployment.success.distinct_count_by(:user_id)
+            deployments: distinct_count(::Deployment, :user_id),
+            failed_deployments: distinct_count(::Deployment.failed, :user_id),
+            projects_mirrored_with_pipelines_enabled: distinct_count(::Project.mirrored_with_enabled_pipelines, :creator_id),
+            releases: distinct_count(::Release, :author_id),
+            successful_deployments: distinct_count(::Deployment.success, :user_id)
           }
         end
 
         # Omitted because no user, creator or author associated: `ci_runners`
         def usage_activity_by_stage_verify
           {
-            ci_builds: ::Ci::Build.distinct_count_by(:user_id),
-            ci_external_pipelines: ::Ci::Pipeline.external.distinct_count_by(:user_id),
-            ci_internal_pipelines: ::Ci::Pipeline.internal.distinct_count_by(:user_id),
-            ci_pipeline_config_auto_devops: ::Ci::Pipeline.auto_devops_source.distinct_count_by(:user_id),
-            ci_pipeline_config_repository: ::Ci::Pipeline.repository_source.distinct_count_by(:user_id),
-            ci_pipeline_schedules: ::Ci::PipelineSchedule.distinct_count_by(:owner_id),
-            ci_pipelines: ::Ci::Pipeline.distinct_count_by(:user_id),
-            ci_triggers: ::Ci::Trigger.distinct_count_by(:owner_id),
+            ci_builds: distinct_count(::Ci::Build, :user_id),
+            ci_external_pipelines: distinct_count(::Ci::Pipeline.external, :user_id),
+            ci_internal_pipelines: distinct_count(::Ci::Pipeline.internal, :user_id),
+            ci_pipeline_config_auto_devops: distinct_count(::Ci::Pipeline.auto_devops_source, :user_id),
+            ci_pipeline_config_repository: distinct_count(::Ci::Pipeline.repository_source, :user_id),
+            ci_pipeline_schedules: distinct_count(::Ci::PipelineSchedule, :owner_id),
+            ci_pipelines: distinct_count(::Ci::Pipeline, :user_id),
+            ci_triggers: distinct_count(::Ci::Trigger, :owner_id),
             clusters_applications_runner: ::Clusters::Applications::Runner.distinct_by_user,
-            projects_reporting_ci_cd_back_to_github: ::Project.with_github_service_pipeline_events.distinct_count_by(:creator_id)
+            projects_reporting_ci_cd_back_to_github: distinct_count(::Project.with_github_service_pipeline_events, :creator_id)
           }
         end
 

ee/spec/controllers/admin/instance_review_controller_spec.rb

@@ -17,6 +17,7 @@ describe Admin::InstanceReviewController do
 
     context 'with usage ping enabled' do
       before do
+        allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
         stub_application_setting(usage_ping_enabled: true)
         ::Gitlab::UsageData.data(force_refresh: true)
         subject

ee/spec/lib/ee/gitlab/usage_data_spec.rb

@@ -3,14 +3,24 @@
 require 'spec_helper'
 
 describe Gitlab::UsageData do
+  before do
+    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+  end
+
+  [true, false].each do |usage_ping_batch_counter_on|
+    describe "when the feature flag usage_ping_batch_counter is set to #{usage_ping_batch_counter_on}" do
+      before do
+        stub_feature_flags(usage_ping_batch_counter: usage_ping_batch_counter_on)
+      end
+
       describe '.uncached_data' do
         context 'when on Gitlab.com' do
           before do
             allow(Gitlab).to receive(:com?).and_return(true)
           end
 
-      it 'does not include usage_activity_by_stage data' do
-        expect(described_class.uncached_data).not_to include(:usage_activity_by_stage)
+          it "does #{'not' unless usage_ping_batch_counter_on} include usage_activity_by_stage data" do
+            expect(described_class.uncached_data.include?(:usage_activity_by_stage)).to be(usage_ping_batch_counter_on)
           end
 
           context 'when feature is enabled' do
@@ -18,8 +28,8 @@ describe Gitlab::UsageData do
               stub_feature_flags(usage_activity_by_stage: true)
             end
 
-        it 'does not include usage_activity_by_stage data' do
-          expect(described_class.uncached_data).not_to include(:usage_activity_by_stage)
+            it "does #{'not' unless usage_ping_batch_counter_on} include usage_activity_by_stage data" do
+              expect(described_class.uncached_data.include?(:usage_activity_by_stage)).to be(usage_ping_batch_counter_on)
             end
           end
         end
@@ -29,7 +39,7 @@ describe Gitlab::UsageData do
             stub_feature_flags(usage_activity_by_stage: false)
           end
 
-      it 'does not include usage_activity_by_stage data' do
+          it "does not include usage_activity_by_stage data" do
             expect(described_class.uncached_data).not_to include(:usage_activity_by_stage)
           end
         end
@@ -251,4 +261,6 @@ describe Gitlab::UsageData do
           end
         end
       end
+    end
+  end
 end

ee/spec/lib/gitlab/usage_data_spec.rb

@@ -3,6 +3,10 @@
 require 'spec_helper'
 
 describe Gitlab::UsageData do
+  before do
+    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+  end
+
   describe '#data' do
     # using Array.new to create a different creator User for each of the projects
     let_it_be(:projects) { Array.new(3) { create(:project, :repository, creator: create(:user, group_view: :security_dashboard)) } }

lib/gitlab/database/batch_count.rb

+# frozen_string_literal: true
+
+# For large tables, PostgreSQL can take a long time to count rows due to MVCC.
+# Implements a distinct and ordinary batch counter
+# Needs indexes on the column below to calculate max, min and range queries
+# For larger tables just set use higher batch_size with index optimization
+# See https://gitlab.com/gitlab-org/gitlab/-/merge_requests/22705
+# Examples:
+#  extend ::Gitlab::Database::BatchCount
+#  batch_count(User.active)
+#  batch_count(::Clusters::Cluster.aws_installed.enabled, :cluster_id)
+#  batch_distinct_count(::Project, :creator_id)
+module Gitlab
+  module Database
+    module BatchCount
+      def batch_count(relation, column = nil, batch_size: nil)
+        BatchCounter.new(relation, column: column).count(batch_size: batch_size)
+      end
+
+      def batch_distinct_count(relation, column = nil, batch_size: nil)
+        BatchCounter.new(relation, column: column).count(mode: :distinct, batch_size: batch_size)
+      end
+
+      class << self
+        include BatchCount
+      end
+    end
+
+    class BatchCounter
+      FALLBACK = -1
+      MIN_REQUIRED_BATCH_SIZE = 2_000
+      MAX_ALLOWED_LOOPS = 10_000
+      SLEEP_TIME_IN_SECONDS = 0.01 # 10 msec sleep
+      # Each query should take <<500ms https://gitlab.com/gitlab-org/gitlab/-/merge_requests/22705
+      DEFAULT_DISTINCT_BATCH_SIZE = 100_000
+      DEFAULT_BATCH_SIZE = 10_000
+
+      def initialize(relation, column: nil)
+        @relation = relation
+        @column = column || relation.primary_key
+      end
+
+      def unwanted_configuration?(finish, batch_size, start)
+        batch_size <= MIN_REQUIRED_BATCH_SIZE ||
+          (finish - start) / batch_size >= MAX_ALLOWED_LOOPS ||
+          start > finish
+      end
+
+      def count(batch_size: nil, mode: :itself)
+        raise 'BatchCount can not be run inside a transaction' if ActiveRecord::Base.connection.transaction_open?
+        raise "The mode #{mode.inspect} is not supported" unless [:itself, :distinct].include?(mode)
+
+        # non-distinct have better performance
+        batch_size ||= mode == :distinct ? DEFAULT_BATCH_SIZE : DEFAULT_DISTINCT_BATCH_SIZE
+
+        start = @relation.minimum(@column) || 0
+        finish = @relation.maximum(@column) || 0
+
+        raise "Batch counting expects positive values only for #{@column}" if start < 0 || finish < 0
+        return FALLBACK if unwanted_configuration?(finish, batch_size, start)
+
+        counter = 0
+        batch_start = start
+
+        while batch_start <= finish
+          begin
+            counter += batch_fetch(batch_start, batch_start + batch_size, mode)
+            batch_start += batch_size
+          rescue ActiveRecord::QueryCanceled
+            # retry with a safe batch size & warmer cache
+            if batch_size >= 2 * MIN_REQUIRED_BATCH_SIZE
+              batch_size /= 2
+            else
+              return FALLBACK
+            end
+          end
+          sleep(SLEEP_TIME_IN_SECONDS)
+        end
+
+        counter
+      end
+
+      def batch_fetch(start, finish, mode)
+        # rubocop:disable GitlabSecurity/PublicSend
+        @relation.select(@column).public_send(mode).where(@column => start..(finish - 1)).count
+      end
+    end
+  end
+end

lib/gitlab/usage_data.rb

@@ -67,8 +67,8 @@ module Gitlab
             clusters_disabled: count(::Clusters::Cluster.disabled),
             project_clusters_disabled: count(::Clusters::Cluster.disabled.project_type),
             group_clusters_disabled: count(::Clusters::Cluster.disabled.group_type),
-            clusters_platforms_eks: count(::Clusters::Cluster.aws_installed.enabled),
-            clusters_platforms_gke: count(::Clusters::Cluster.gcp_installed.enabled),
+            clusters_platforms_eks: count(::Clusters::Cluster.aws_installed.enabled, batch: false),
+            clusters_platforms_gke: count(::Clusters::Cluster.gcp_installed.enabled, batch: false),
             clusters_platforms_user: count(::Clusters::Cluster.user_provided.enabled),
             clusters_applications_helm: count(::Clusters::Applications::Helm.available),
             clusters_applications_ingress: count(::Clusters::Applications::Ingress.available),
@@ -85,7 +85,7 @@ module Gitlab
             issues: count(Issue),
             issues_created_from_gitlab_error_tracking_ui: count(SentryIssue),
             issues_with_associated_zoom_link: count(ZoomMeeting.added_to_issue),
-            issues_using_zoom_quick_actions: count(ZoomMeeting.select(:issue_id).distinct),
+            issues_using_zoom_quick_actions: count(ZoomMeeting.select(:issue_id).distinct, batch: false),
             issues_with_embedded_grafana_charts_approx: ::Gitlab::GrafanaEmbedUsageData.issue_count,
             incident_issues: count(::Issue.authored(::User.alert_bot)),
             keys: count(Key),
@@ -99,7 +99,7 @@ module Gitlab
             projects_imported_from_github: count(Project.where(import_type: 'github')),
             projects_with_repositories_enabled: count(ProjectFeature.where('repository_access_level > ?', ProjectFeature::DISABLED)),
             projects_with_error_tracking_enabled: count(::ErrorTracking::ProjectErrorTrackingSetting.where(enabled: true)),
-            projects_with_alerts_service_enabled: count(AlertsService.active),
+            projects_with_alerts_service_enabled: count(AlertsService.active, batch: false),
             protected_branches: count(ProtectedBranch),
             releases: count(Release),
             remote_mirrors: count(RemoteMirror),
@@ -181,7 +181,7 @@ module Gitlab
 
       # rubocop: disable CodeReuse/ActiveRecord
       def services_usage
-        service_counts = count(Service.active.where(template: false).where.not(type: 'JiraService').group(:type), fallback: Hash.new(-1))
+        service_counts = count(Service.active.where(template: false).where.not(type: 'JiraService').group(:type), fallback: Hash.new(-1), batch: false)
 
         results = Service.available_services_names.each_with_object({}) do |service_name, response|
           response["projects_#{service_name}_active".to_sym] = service_counts["#{service_name}_service".camelize] || 0
@@ -217,9 +217,9 @@ module Gitlab
           results[:projects_jira_server_active] += counts[:server].count if counts[:server]
           results[:projects_jira_cloud_active] += counts[:cloud].count if counts[:cloud]
           if results[:projects_jira_active] == -1
-            results[:projects_jira_active] = count(services)
+            results[:projects_jira_active] = count(services, batch: false)
           else
-            results[:projects_jira_active] += count(services)
+            results[:projects_jira_active] += count(services, batch: false)
           end
         end
 
@@ -231,8 +231,22 @@ module Gitlab
         {} # augmented in EE
       end
 
-      def count(relation, fallback: -1)
+      def count(relation, column = nil, fallback: -1, batch: true)
+        if batch && Feature.enabled?(:usage_ping_batch_counter)
+          Gitlab::Database::BatchCount.batch_count(relation, column)
+        else
           relation.count
+        end
+      rescue ActiveRecord::StatementInvalid
+        fallback
+      end
+
+      def distinct_count(relation, column = nil, fallback: -1, batch: true)
+        if batch && Feature.enabled?(:usage_ping_batch_counter)
+          Gitlab::Database::BatchCount.batch_distinct_count(relation, column)
+        else
+          relation.distinct_count_by(column)
+        end
       rescue ActiveRecord::StatementInvalid
         fallback
       end

spec/controllers/admin/application_settings_controller_spec.rb

@@ -16,6 +16,7 @@ describe Admin::ApplicationSettingsController do
 
   describe 'GET #usage_data with no access' do
     before do
+      allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
       sign_in(user)
     end
 
@@ -28,6 +29,7 @@ describe Admin::ApplicationSettingsController do
 
   describe 'GET #usage_data' do
     before do
+      allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
       sign_in(admin)
     end
 

spec/features/admin/admin_settings_spec.rb

@@ -326,6 +326,8 @@ describe 'Admin updates settings', :clean_gitlab_redis_shared_state, :do_not_moc
       end
 
       it 'loads usage ping payload on click', :js do
+        allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+
         expect(page).to have_button 'Preview payload'
 
         find('.js-usage-ping-payload-trigger').click

spec/lib/gitlab/database/batch_count_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Database::BatchCount do
+  let(:model) { Issue }
+  let(:column) { :author_id }
+
+  let(:in_transaction) { false }
+  let(:user) { create(:user) }
+  let(:another_user) { create(:user) }
+
+  before do
+    create_list(:issue, 3, author: user )
+    create_list(:issue, 2, author: another_user )
+    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(in_transaction)
+  end
+
+  describe '#batch_count' do
+    it 'counts table' do
+      expect(described_class.batch_count(model)).to eq(5)
+    end
+
+    it 'counts with :id field' do
+      expect(described_class.batch_count(model, :id)).to eq(5)
+    end
+
+    it 'counts with "id" field' do
+      expect(described_class.batch_count(model, 'id')).to eq(5)
+    end
+
+    it 'counts with table.id field' do
+      expect(described_class.batch_count(model, "#{model.table_name}.id")).to eq(5)
+    end
+
+    it 'counts table with batch_size 50K' do
+      expect(described_class.batch_count(model, batch_size: 50_000)).to eq(5)
+    end
+
+    it 'will not count table with batch_size 1K' do
+      fallback = ::Gitlab::Database::BatchCounter::FALLBACK
+      expect(described_class.batch_count(model, batch_size: fallback / 2)).to eq(fallback)
+    end
+
+    it 'counts with a small edge case batch_sizes than result' do
+      stub_const('Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE', 0)
+
+      [1, 2, 4, 5, 6].each { |i| expect(described_class.batch_count(model, batch_size: i)).to eq(5) }
+    end
+
+    context 'in a transaction' do
+      let(:in_transaction) { true }
+
+      it 'cannot count' do
+        expect do
+          described_class.batch_count(model)
+        end.to raise_error 'BatchCount can not be run inside a transaction'
+      end
+    end
+  end
+
+  describe '#batch_distinct_count' do
+    it 'counts with :id field' do
+      expect(described_class.batch_distinct_count(model, :id)).to eq(5)
+    end
+
+    it 'counts with column field' do
+      expect(described_class.batch_distinct_count(model, column)).to eq(2)
+    end
+
+    it 'counts with "id" field' do
+      expect(described_class.batch_distinct_count(model, "#{column}")).to eq(2)
+    end
+
+    it 'counts with table.column field' do
+      expect(described_class.batch_distinct_count(model, "#{model.table_name}.#{column}")).to eq(2)
+    end
+
+    it 'counts with :column field with batch_size of 50K' do
+      expect(described_class.batch_distinct_count(model, column, batch_size: 50_000)).to eq(2)
+    end
+
+    it 'will not count table with batch_size 1K' do
+      fallback = ::Gitlab::Database::BatchCounter::FALLBACK
+      expect(described_class.batch_distinct_count(model, column, batch_size: fallback / 2)).to eq(fallback)
+    end
+
+    it 'counts with a small edge case batch_sizes than result' do
+      stub_const('Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE', 0)
+
+      [1, 2, 4, 5, 6].each { |i| expect(described_class.batch_distinct_count(model, column, batch_size: i)).to eq(2) }
+    end
+  end
+end

spec/lib/gitlab/usage_data_spec.rb

@@ -6,6 +6,15 @@ describe Gitlab::UsageData do
   let(:projects) { create_list(:project, 4) }
   let!(:board) { create(:board, project: projects[0]) }
 
+  before do
+    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+  end
+  [true, false].each do |usage_ping_batch_counter_on|
+    describe "when the feature flag usage_ping_batch_counter is set to #{usage_ping_batch_counter_on}" do
+      before do
+        stub_feature_flags(usage_ping_batch_counter: usage_ping_batch_counter_on)
+      end
+
       describe '#data' do
         before do
           create(:jira_service, project: projects[0])
@@ -351,13 +360,13 @@ describe Gitlab::UsageData do
         it 'returns the count when counting succeeds' do
           allow(relation).to receive(:count).and_return(1)
 
-      expect(described_class.count(relation)).to eq(1)
+          expect(described_class.count(relation, batch: false)).to eq(1)
         end
 
         it 'returns the fallback value when counting fails' do
           allow(relation).to receive(:count).and_raise(ActiveRecord::StatementInvalid.new(''))
 
-      expect(described_class.count(relation, fallback: 15)).to eq(15)
+          expect(described_class.count(relation, fallback: 15, batch: false)).to eq(15)
         end
       end
 
@@ -383,4 +392,6 @@ describe Gitlab::UsageData do
           expect(described_class.approximate_counts.values.uniq).to eq([-1])
         end
       end
+    end
+  end
 end

spec/services/submit_usage_ping_service_spec.rb

@@ -76,6 +76,7 @@ describe SubmitUsagePingService do
 
   context 'when usage ping is enabled' do
     before do
+      allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
       stub_application_setting(usage_ping_enabled: true)
     end