Use approximate counts for big tables

changelogs/unreleased/usage-count.yml

+---
+title: Use approximate count for big tables for usage statistics.
+merge_request:
+author:
+type: fixed

lib/gitlab/database/count.rb

@@ -40,7 +40,7 @@ module Gitlab
           if strategy.enabled?
             models_with_missing_counts = models - counts_by_model.keys
 
-            break if models_with_missing_counts.empty?
+            break counts_by_model if models_with_missing_counts.empty?
 
             counts = strategy.new(models_with_missing_counts).count
 

lib/gitlab/database/count/exact_count_strategy.rb

@@ -20,6 +20,8 @@ module Gitlab
           models.each_with_object({}) do |model, data|
             data[model] = model.count
           end
+        rescue *CONNECTION_ERRORS
+          {}
         end
 
         def self.enabled?

lib/gitlab/usage_data.rb

@@ -2,6 +2,8 @@
 
 module Gitlab
   class UsageData
+    APPROXIMATE_COUNT_MODELS = [Label, MergeRequest, Note, Todo].freeze
+
     class << self
       def data(force_refresh: false)
         Rails.cache.fetch('usage_data', force: force_refresh, expires_in: 2.weeks) { uncached_data }
@@ -73,12 +75,9 @@ module Gitlab
             issues: count(Issue),
             keys: count(Key),
             label_lists: count(List.label),
-            labels: count(Label),
             lfs_objects: count(LfsObject),
-            merge_requests: count(MergeRequest),
             milestone_lists: count(List.milestone),
             milestones: count(Milestone),
-            notes: count(Note),
             pages_domains: count(PagesDomain),
             projects: count(Project),
             projects_imported_from_github: count(Project.where(import_type: 'github')),
@@ -86,10 +85,9 @@ module Gitlab
             releases: count(Release),
             remote_mirrors: count(RemoteMirror),
             snippets: count(Snippet),
-            todos: count(Todo),
             uploads: count(Upload),
             web_hooks: count(WebHook)
-          }.merge(services_usage)
+          }.merge(services_usage).merge(approximate_counts)
         }
       end
       # rubocop: enable CodeReuse/ActiveRecord
@@ -164,6 +162,16 @@ module Gitlab
         fallback
       end
       # rubocop: enable CodeReuse/ActiveRecord
+
+      def approximate_counts
+        approx_counts = Gitlab::Database::Count.approximate_counts(APPROXIMATE_COUNT_MODELS)
+
+        APPROXIMATE_COUNT_MODELS.each_with_object({}) do |model, result|
+          key = model.name.underscore.pluralize.to_sym
+
+          result[key] = approx_counts[model] || -1
+        end
+      end
     end
   end
 end

spec/lib/gitlab/database/count/exact_count_strategy_spec.rb

@@ -16,6 +16,12 @@ describe Gitlab::Database::Count::ExactCountStrategy do
 
       expect(subject).to eq({ Project => 3, Identity => 1 })
     end
+
+    it 'returns default value if count times out' do
+      allow(models.first).to receive(:count).and_raise(ActiveRecord::StatementInvalid.new(''))
+
+      expect(subject).to eq({})
+    end
   end
 
   describe '.enabled?' do

spec/lib/gitlab/usage_data_spec.rb

@@ -213,4 +213,29 @@ describe Gitlab::UsageData do
       expect(described_class.count(relation, fallback: 15)).to eq(15)
     end
   end
+
+  describe '#approximate_counts' do
+    it 'gets approximate counts for selected models' do
+      create(:label)
+
+      expect(Gitlab::Database::Count).to receive(:approximate_counts)
+        .with(described_class::APPROXIMATE_COUNT_MODELS).once.and_call_original
+
+      counts = described_class.approximate_counts.values
+
+      expect(counts.count).to eq(described_class::APPROXIMATE_COUNT_MODELS.count)
+      expect(counts.any? { |count| count < 0 }).to be_falsey
+    end
+
+    it 'returns default values if counts can not be retrieved' do
+      described_class::APPROXIMATE_COUNT_MODELS.map do |model|
+        model.name.underscore.pluralize.to_sym
+      end
+
+      expect(Gitlab::Database::Count).to receive(:approximate_counts)
+        .and_return({})
+
+      expect(described_class.approximate_counts.values.uniq).to eq([-1])
+    end
+  end
 end