Fix Error 500 viewing admin page due to statement timeouts

Uses PostgreSQL tuple estimates to provide a much faster yet approximate
count. See https://wiki.postgresql.org/wiki/Slow_Counting for more details.
We only use this fast method if the table has been analyzed or vacuumed
within the last hour.

Closes #46255

5 files changed, 127 insertions, 10 deletions

diff --git a/app/controllers/admin/dashboard_controller.rb b/app/controllers/admin/dashboard_controller.rb
index e85cdcb8db7..d6a6bc7d4a1 100644
--- a/app/controllers/admin/dashboard_controller.rb
+++ b/app/controllers/admin/dashboard_controller.rb
@@ -1,4 +1,6 @@
 class Admin::DashboardController < Admin::ApplicationController
+  include CountHelper
+
   def index
     @projects = Project.order_id_desc.without_deleted.with_route.limit(10)
     @users = User.order_id_desc.limit(10)
diff --git a/app/helpers/count_helper.rb b/app/helpers/count_helper.rb
new file mode 100644
index 00000000000..24ee62e68ba
--- /dev/null
+++ b/app/helpers/count_helper.rb
@@ -0,0 +1,5 @@
+module CountHelper
+  def approximate_count_with_delimiters(model)
+    number_with_delimiter(Gitlab::Database::Count.approximate_count(model))
+  end
+end
diff --git a/app/views/admin/dashboard/index.html.haml b/app/views/admin/dashboard/index.html.haml
index bbf0e0fb95c..41ef646fc0e 100644
--- a/app/views/admin/dashboard/index.html.haml
+++ b/app/views/admin/dashboard/index.html.haml
@@ -10,7 +10,7 @@
             = link_to admin_projects_path do
               %h3.text-center
                 Projects:
-                = number_with_delimiter(Project.cached_count)
+                = approximate_count_with_delimiters(Project)
             %hr
             = link_to('New project', new_project_path, class: "btn btn-new")
       .col-sm-4
@@ -19,7 +19,7 @@
             = link_to admin_users_path do
               %h3.text-center
                 Users:
-                = number_with_delimiter(User.count)
+                = approximate_count_with_delimiters(User)
             %hr
             = link_to 'New user', new_admin_user_path, class: "btn btn-new"
       .col-sm-4
@@ -28,7 +28,7 @@
             = link_to admin_groups_path do
               %h3.text-center
                 Groups:
-                = number_with_delimiter(Group.count)
+                = approximate_count_with_delimiters(Group)
             %hr
             = link_to 'New group', new_admin_group_path, class: "btn btn-new"
     .row
@@ -39,31 +39,31 @@
             %p
               Forks
               %span.light.pull-right
-                = number_with_delimiter(ForkedProjectLink.count)
+                = approximate_count_with_delimiters(ForkedProjectLink)
             %p
               Issues
               %span.light.pull-right
-                = number_with_delimiter(Issue.count)
+                = approximate_count_with_delimiters(Issue)
             %p
               Merge Requests
               %span.light.pull-right
-                = number_with_delimiter(MergeRequest.count)
+                = approximate_count_with_delimiters(MergeRequest)
             %p
               Notes
               %span.light.pull-right
-                = number_with_delimiter(Note.count)
+                = approximate_count_with_delimiters(Note)
             %p
               Snippets
               %span.light.pull-right
-                = number_with_delimiter(Snippet.count)
+                = approximate_count_with_delimiters(Snippet)
             %p
               SSH Keys
               %span.light.pull-right
-                = number_with_delimiter(Key.count)
+                = approximate_count_with_delimiters(Key)
             %p
               Milestones
               %span.light.pull-right
-                = number_with_delimiter(Milestone.count)
+                = approximate_count_with_delimiters(Milestone)
             %p
               Active Users
               %span.light.pull-right
diff --git a/lib/gitlab/database/count.rb b/lib/gitlab/database/count.rb
new file mode 100644
index 00000000000..3374203960e
--- /dev/null
+++ b/lib/gitlab/database/count.rb
@@ -0,0 +1,48 @@
+# For large tables, PostgreSQL can take a long time to count rows due to MVCC.
+# We can optimize this by using the reltuples count as described in https://wiki.postgresql.org/wiki/Slow_Counting.
+module Gitlab
+  module Database
+    module Count
+      CONNECTION_ERRORS =
+        if defined?(PG)
+          [
+            ActionView::Template::Error,
+            ActiveRecord::StatementInvalid,
+            PG::Error
+          ].freeze
+        else
+          [
+            ActionView::Template::Error,
+            ActiveRecord::StatementInvalid
+          ].freeze
+        end
+
+      def self.approximate_count(model)
+        return model.count unless Gitlab::Database.postgresql?
+
+        execute_estimate_if_updated_recently(model) || model.count
+      end
+
+      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
+      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
+
+        ActiveRecord::Base.connection.select_all(query).count > 0
+      rescue *CONNECTION_ERRORS
+        false
+      end
+
+      def self.postgresql_estimate_query(model)
+        "SELECT reltuples::bigint AS estimate FROM pg_class where relname = '#{model.table_name}'"
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/database/count_spec.rb b/spec/lib/gitlab/database/count_spec.rb
new file mode 100644
index 00000000000..9d9caaabe16
--- /dev/null
+++ b/spec/lib/gitlab/database/count_spec.rb
@@ -0,0 +1,62 @@
+require 'spec_helper'
+
+describe Gitlab::Database::Count do
+  before do
+    create_list(:project, 3)
+  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
+
+      it 'returns nil' do
+        expect(described_class.execute_estimate_if_updated_recently(Project)).to be nil
+      end
+    end
+
+    context 'when reltuples have been updated' do
+      before do
+        ActiveRecord::Base.connection.execute('ANALYZE projects')
+      end
+
+      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)
+      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)
+
+        expect(Project).to receive(:count).and_call_original
+        expect(described_class.approximate_count(Project)).to eq(3)
+      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)
+
+        expect(Project).to receive(:count).and_call_original
+        expect(described_class.approximate_count(Project)).to eq(3)
+      end
+    end
+
+    describe 'when reltuples have been updated', :postgresql do
+      before do
+        ActiveRecord::Base.connection.execute('ANALYZE projects')
+      end
+
+      it 'counts all projects in the fast way' do
+        expect(described_class).to receive(:postgresql_estimate_query).with(Project).and_call_original
+
+        expect(described_class.approximate_count(Project)).to eq(3)
+      end
+    end
+  end
+end