Merge branch 'sh-fast-admin-counts-ee' into 'master'

Fix Error 500 viewing admin page due to statement timeouts (EE port)

See merge request gitlab-org/gitlab-ee!5733

app/controllers/admin/dashboard_controller.rb

 class Admin::DashboardController < Admin::ApplicationController
   prepend ::EE::Admin::DashboardController
 
+  include CountHelper
+
   def index
     @projects = Project.order_id_desc.without_deleted.with_route.limit(10)
     @users = User.order_id_desc.limit(10)

app/helpers/count_helper.rb

+module CountHelper
+  def approximate_count_with_delimiters(model)
+    number_with_delimiter(Gitlab::Database::Count.approximate_count(model))
+  end
+end

app/views/admin/dashboard/index.html.haml

@@ -13,7 +13,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
@@ -22,7 +22,7 @@
             = link_to admin_users_path do
               %h3.text-center
                 Users:
-                = number_with_delimiter(User.count)
+                = approximate_count_with_delimiters(User)
             -# EE specific
             .text-center
               = link_to 'Users statistics', admin_dashboard_stats_path
@@ -34,7 +34,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
@@ -45,31 +45,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

lib/gitlab/database/count.rb

+# 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

spec/lib/gitlab/database/count_spec.rb

+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