Fix fast admin counters not working when PostgreSQL has secondariessh-fix-admin-page-counts-take-2

This commit does a number of things:

1. Reduces the number of queries needed by perform a single query to get all
the tuples for the relevant rows.

2. Uses a transaction to query the tuple counts to ensure that the data
is retrieved from the primary.

Closes #46742

7 files changed, 130 insertions, 65 deletions

diff --git a/app/controllers/admin/dashboard_controller.rb b/app/controllers/admin/dashboard_controller.rb
index d6a6bc7d4a1..737942f3eb2 100644
--- a/app/controllers/admin/dashboard_controller.rb
+++ b/app/controllers/admin/dashboard_controller.rb
@@ -1,7 +1,11 @@
 class Admin::DashboardController < Admin::ApplicationController
   include CountHelper
 
+  COUNTED_ITEMS = [Project, User, Group, ForkedProjectLink, Issue, MergeRequest,
+                   Note, Snippet, Key, Milestone].freeze
+
   def index
+    @counts = Gitlab::Database::Count.approximate_counts(COUNTED_ITEMS)
     @projects = Project.order_id_desc.without_deleted.with_route.limit(10)
     @users = User.order_id_desc.limit(10)
     @groups = Group.order_id_desc.with_route.limit(10)
diff --git a/app/helpers/count_helper.rb b/app/helpers/count_helper.rb
index 24ee62e68ba..5cd98f40f78 100644
--- a/app/helpers/count_helper.rb
+++ b/app/helpers/count_helper.rb
@@ -1,5 +1,9 @@
 module CountHelper
-  def approximate_count_with_delimiters(model)
-    number_with_delimiter(Gitlab::Database::Count.approximate_count(model))
+  def approximate_count_with_delimiters(count_data, model)
+    count = count_data[model]
+
+    raise "Missing model #{model} from count data" unless count
+
+    number_with_delimiter(count)
   end
 end
diff --git a/app/views/admin/dashboard/index.html.haml b/app/views/admin/dashboard/index.html.haml
index 3df4ce93fa8..3cdeb103bb8 100644
--- a/app/views/admin/dashboard/index.html.haml
+++ b/app/views/admin/dashboard/index.html.haml
@@ -12,7 +12,7 @@
             = link_to admin_projects_path do
               %h3.text-center
                 Projects:
-                = approximate_count_with_delimiters(Project)
+                = approximate_count_with_delimiters(@counts, Project)
             %hr
             = link_to('New project', new_project_path, class: "btn btn-new")
       .col-sm-4
@@ -21,7 +21,7 @@
             = link_to admin_users_path do
               %h3.text-center
                 Users:
-                = approximate_count_with_delimiters(User)
+                = approximate_count_with_delimiters(@counts, User)
             = render_if_exists 'users_statistics'
             %hr
             = link_to 'New user', new_admin_user_path, class: "btn btn-new"
@@ -31,7 +31,7 @@
             = link_to admin_groups_path do
               %h3.text-center
                 Groups:
-                = approximate_count_with_delimiters(Group)
+                = approximate_count_with_delimiters(@counts, Group)
             %hr
             = link_to 'New group', new_admin_group_path, class: "btn btn-new"
     .row
@@ -42,31 +42,31 @@
             %p
               Forks
               %span.light.float-right
-                = approximate_count_with_delimiters(ForkedProjectLink)
+                = approximate_count_with_delimiters(@counts, ForkedProjectLink)
             %p
               Issues
               %span.light.float-right
-                = approximate_count_with_delimiters(Issue)
+                = approximate_count_with_delimiters(@counts, Issue)
             %p
               Merge Requests
               %span.light.float-right
-                = approximate_count_with_delimiters(MergeRequest)
+                = approximate_count_with_delimiters(@counts, MergeRequest)
             %p
               Notes
               %span.light.float-right
-                = approximate_count_with_delimiters(Note)
+                = approximate_count_with_delimiters(@counts, Note)
             %p
               Snippets
               %span.light.float-right
-                = approximate_count_with_delimiters(Snippet)
+                = approximate_count_with_delimiters(@counts, Snippet)
             %p
               SSH Keys
               %span.light.float-right
-                = approximate_count_with_delimiters(Key)
+                = approximate_count_with_delimiters(@counts, Key)
             %p
               Milestones
               %span.light.float-right
-                = approximate_count_with_delimiters(Milestone)
+                = approximate_count_with_delimiters(@counts, Milestone)
             %p
               Active Users
               %span.light.float-right
diff --git a/changelogs/unreleased/sh-fix-admin-page-counts-take-2.yml b/changelogs/unreleased/sh-fix-admin-page-counts-take-2.yml
new file mode 100644
index 00000000000..d9bd1af9380
--- /dev/null
+++ b/changelogs/unreleased/sh-fix-admin-page-counts-take-2.yml
@@ -0,0 +1,5 @@
+---
+title: Fix admin counters not working when PostgreSQL has secondaries
+merge_request:
+author:
+type: fixed
diff --git a/lib/gitlab/database/count.rb b/lib/gitlab/database/count.rb
index 3374203960e..5f549ed2b3c 100644
--- a/lib/gitlab/database/count.rb
+++ b/lib/gitlab/database/count.rb
@@ -17,31 +17,69 @@ module Gitlab
           ].freeze
         end
 
-      def self.approximate_count(model)
-        return model.count unless Gitlab::Database.postgresql?
+      # Takes in an array of models and returns a Hash for the approximate
+      # counts for them.  If the model's table has not been vacuumed or
+      # analyzed recently, simply run the Model.count to get the data.
+      #
+      # @param [Array]
+      # @return [Hash] of Model -> count mapping
+      def self.approximate_counts(models)
+        table_to_model_map = models.each_with_object({}) do |model, hash|
+          hash[model.table_name] = model
+        end
 
-        execute_estimate_if_updated_recently(model) || model.count
-      end
+        table_names = table_to_model_map.keys
+        counts_by_table_name = Gitlab::Database.postgresql? ? reltuples_from_recently_updated(table_names) : {}
 
-      def self.execute_estimate_if_updated_recently(model)
-        ActiveRecord::Base.connection.select_value(postgresql_estimate_query(model)).to_i if reltuples_updated_recently?(model)
-      rescue *CONNECTION_ERRORS
+        # Convert table -> count to Model -> count
+        counts_by_model = counts_by_table_name.each_with_object({}) do |pair, hash|
+          model = table_to_model_map[pair.first]
+          hash[model] = pair.second
+        end
+
+        missing_tables = table_names - counts_by_table_name.keys
+
+        missing_tables.each do |table|
+          model = table_to_model_map[table]
+          counts_by_model[model] = model.count
+        end
+
+        counts_by_model
       end
 
-      def self.reltuples_updated_recently?(model)
-        time = "to_timestamp(#{1.hour.ago.to_i})"
-        query = <<~SQL
-          SELECT 1 FROM pg_stat_user_tables WHERE relname = '#{model.table_name}' AND
-          (last_vacuum > #{time} OR last_autovacuum > #{time} OR last_analyze > #{time} OR last_autoanalyze > #{time})
-        SQL
+      # Returns a hash of the table names that have recently updated tuples.
+      #
+      # @param [Array] table names
+      # @returns [Hash] Table name to count mapping (e.g. { 'projects' => 5, 'users' => 100 })
+      def self.reltuples_from_recently_updated(table_names)
+        query = postgresql_estimate_query(table_names)
+        rows = []
 
-        ActiveRecord::Base.connection.select_all(query).count > 0
+        # Querying tuple stats only works on the primary. Due to load
+        # balancing, we need to ensure this query hits the load balancer.  The
+        # easiest way to do this is to start a transaction.
+        ActiveRecord::Base.transaction do
+          rows = ActiveRecord::Base.connection.select_all(query)
+        end
+
+        rows.each_with_object({}) { |row, data| data[row['table_name']] = row['estimate'].to_i }
       rescue *CONNECTION_ERRORS
-        false
+        {}
       end
 
-      def self.postgresql_estimate_query(model)
-        "SELECT reltuples::bigint AS estimate FROM pg_class where relname = '#{model.table_name}'"
+      # Generates the PostgreSQL query to return the tuples for tables
+      # that have been vacuumed or analyzed in the last hour.
+      #
+      # @param [Array] table names
+      # @returns [Hash] Table name to count mapping (e.g. { 'projects' => 5, 'users' => 100 })
+      def self.postgresql_estimate_query(table_names)
+        time = "to_timestamp(#{1.hour.ago.to_i})"
+        <<~SQL
+          SELECT pg_class.relname AS table_name, reltuples::bigint AS estimate FROM pg_class
+          LEFT JOIN pg_stat_user_tables ON pg_class.relname = pg_stat_user_tables.relname
+          WHERE pg_class.relname IN (#{table_names.map { |table| "'#{table}'" }.join(',')})
+          AND (last_vacuum > #{time} OR last_autovacuum > #{time} OR last_analyze > #{time} OR last_autoanalyze > #{time})
+        SQL
       end
     end
   end
diff --git a/spec/lib/gitlab/database/count_spec.rb b/spec/lib/gitlab/database/count_spec.rb
index 9d9caaabe16..407d9470785 100644
--- a/spec/lib/gitlab/database/count_spec.rb
+++ b/spec/lib/gitlab/database/count_spec.rb
@@ -3,59 +3,68 @@ require 'spec_helper'
 describe Gitlab::Database::Count do
   before do
     create_list(:project, 3)
+    create(:identity)
   end
 
-  describe '.execute_estimate_if_updated_recently', :postgresql do
-    context 'when reltuples have not been updated' do
-      before do
-        expect(described_class).to receive(:reltuples_updated_recently?).and_return(false)
-      end
+  let(:models) { [Project, Identity] }
 
-      it 'returns nil' do
-        expect(described_class.execute_estimate_if_updated_recently(Project)).to be nil
-      end
-    end
+  describe '.approximate_counts' do
+    context 'with MySQL' do
+      context 'when reltuples have not been updated' do
+        it 'counts all models the normal way' do
+          expect(Gitlab::Database).to receive(:postgresql?).and_return(false)
 
-    context 'when reltuples have been updated' do
-      before do
-        ActiveRecord::Base.connection.execute('ANALYZE projects')
-      end
+          expect(Project).to receive(:count).and_call_original
+          expect(Identity).to receive(:count).and_call_original
 
-      it 'calls postgresql_estimate_query' do
-        expect(described_class).to receive(:postgresql_estimate_query).with(Project).and_call_original
-        expect(described_class.execute_estimate_if_updated_recently(Project)).to eq(3)
+          expect(described_class.approximate_counts(models)).to eq({ Project => 3, Identity => 1 })
+        end
       end
     end
-  end
 
-  describe '.approximate_count' do
-    context 'when reltuples have not been updated' do
-      it 'counts all projects the normal way' do
-        allow(described_class).to receive(:reltuples_updated_recently?).and_return(false)
+    context 'with PostgreSQL', :postgresql do
+      describe 'when reltuples have not been updated' do
+        it 'counts all models the normal way' do
+          expect(described_class).to receive(:reltuples_from_recently_updated).with(%w(projects identities)).and_return({})
 
-        expect(Project).to receive(:count).and_call_original
-        expect(described_class.approximate_count(Project)).to eq(3)
+          expect(Project).to receive(:count).and_call_original
+          expect(Identity).to receive(:count).and_call_original
+          expect(described_class.approximate_counts(models)).to eq({ Project => 3, Identity => 1 })
+        end
       end
-    end
 
-    context 'no permission' do
-      it 'falls back to standard query' do
-        allow(described_class).to receive(:reltuples_updated_recently?).and_raise(PG::InsufficientPrivilege)
+      describe 'no permission' do
+        it 'falls back to standard query' do
+          allow(described_class).to receive(:postgresql_estimate_query).and_raise(PG::InsufficientPrivilege)
 
-        expect(Project).to receive(:count).and_call_original
-        expect(described_class.approximate_count(Project)).to eq(3)
+          expect(Project).to receive(:count).and_call_original
+          expect(Identity).to receive(:count).and_call_original
+          expect(described_class.approximate_counts(models)).to eq({ Project => 3, Identity => 1 })
+        end
       end
-    end
 
-    describe 'when reltuples have been updated', :postgresql do
-      before do
-        ActiveRecord::Base.connection.execute('ANALYZE projects')
+      describe 'when some reltuples have been updated' do
+        it 'counts projects in the fast way' do
+          expect(described_class).to receive(:reltuples_from_recently_updated).with(%w(projects identities)).and_return({ 'projects' => 3 })
+
+          expect(Project).not_to receive(:count).and_call_original
+          expect(Identity).to receive(:count).and_call_original
+          expect(described_class.approximate_counts(models)).to eq({ Project => 3, Identity => 1 })
+        end
       end
 
-      it 'counts all projects in the fast way' do
-        expect(described_class).to receive(:postgresql_estimate_query).with(Project).and_call_original
+      describe 'when all reltuples have been updated' do
+        before do
+          ActiveRecord::Base.connection.execute('ANALYZE projects')
+          ActiveRecord::Base.connection.execute('ANALYZE identities')
+        end
+
+        it 'counts models with the standard way' do
+          expect(Project).not_to receive(:count)
+          expect(Identity).not_to receive(:count)
 
-        expect(described_class.approximate_count(Project)).to eq(3)
+          expect(described_class.approximate_counts(models)).to eq({ Project => 3, Identity => 1 })
+        end
       end
     end
   end
diff --git a/spec/views/admin/dashboard/index.html.haml_spec.rb b/spec/views/admin/dashboard/index.html.haml_spec.rb
index 59c777ea338..0e8b7c82d3a 100644
--- a/spec/views/admin/dashboard/index.html.haml_spec.rb
+++ b/spec/views/admin/dashboard/index.html.haml_spec.rb
@@ -4,6 +4,11 @@ describe 'admin/dashboard/index.html.haml' do
   include Devise::Test::ControllerHelpers
 
   before do
+    counts = Admin::DashboardController::COUNTED_ITEMS.each_with_object({}) do |item, hash|
+      hash[item] = 100
+    end
+
+    assign(:counts, counts)
     assign(:projects, create_list(:project, 1))
     assign(:users, create_list(:user, 1))
     assign(:groups, create_list(:group, 1))