Geo: Improve Geo Status API performance with cached counters in SiteStatistic

app/models/project.rb

@@ -34,6 +34,7 @@ class Project < ActiveRecord::Base
 
   BoardLimitExceeded = Class.new(StandardError)
 
+  STATISTICS_ATTRIBUTE = 'repositories_count'.freeze
   NUMBER_OF_PERMITTED_BOARDS = 1
   UNKNOWN_IMPORT_URL = 'http://unknown.git'.freeze
   # Hashed Storage versions handle rolling out new storage to project and dependents models:
@@ -83,6 +84,10 @@ class Project < ActiveRecord::Base
 
   after_create :create_project_feature, unless: :project_feature
 
+  after_create -> { SiteStatistic.track(STATISTICS_ATTRIBUTE) }
+  before_destroy ->(project) { project.project_feature.untrack_statistics_for_deletion! }
+  after_destroy -> { SiteStatistic.untrack(STATISTICS_ATTRIBUTE) }
+
   after_create :create_ci_cd_settings,
     unless: :ci_cd_settings,
     if: proc { ProjectCiCdSetting.available? }

app/models/project_feature.rb

@@ -19,6 +19,7 @@ class ProjectFeature < ActiveRecord::Base
   ENABLED  = 20
 
   FEATURES = %i(issues merge_requests wiki snippets builds repository).freeze
+  STATISTICS_ATTRIBUTE = 'wikis_count'.freeze
 
   class << self
     def access_level_attribute(feature)
@@ -58,6 +59,9 @@ class ProjectFeature < ActiveRecord::Base
     end
   end
 
+  after_create ->(model) { SiteStatistic.track(STATISTICS_ATTRIBUTE) if model.wiki_enabled? }
+  after_update :update_site_statistics
+
   def feature_available?(feature, user)
     get_permission(user, access_level(feature))
   end
@@ -82,8 +86,30 @@ class ProjectFeature < ActiveRecord::Base
     issues_access_level > DISABLED
   end
 
+  # This is a workaround for the removal hooks not been triggered when removing a Project.
+  #
+  # ProjectFeature is removed using database cascade index rule.
+  # This method is called by Project model when deletion starts.
+  def untrack_statistics_for_deletion!
+    return unless wiki_enabled?
+
+    SiteStatistic.untrack(STATISTICS_ATTRIBUTE)
+  end
+
   private
 
+  def update_site_statistics
+    return unless wiki_access_level_changed?
+
+    if self.wiki_access_level_was == DISABLED
+      # possible new states are PRIVATE / ENABLED, both should be tracked
+      SiteStatistic.track(STATISTICS_ATTRIBUTE)
+    elsif self.wiki_access_level == DISABLED
+      # old state was either PRIVATE / ENABLED, only untrack if new state is DISABLED
+      SiteStatistic.untrack(STATISTICS_ATTRIBUTE)
+    end
+  end
+
   # Validates builds and merge requests access level
   # which cannot be higher than repository access level
   def repository_children_level

app/models/site_statistic.rb

+class SiteStatistic < ActiveRecord::Base
+  # prevents the creation of multiple rows
+  default_value_for :id, 1
+
+  COUNTER_ATTRIBUTES = %w(repositories_count wikis_count).freeze
+  REQUIRED_SCHEMA_VERSION = 20180629153018
+
+  # Tracks specific attribute
+  #
+  # @param [String] raw_attribute must be one of the values listed in COUNTER_ATTRIBUTES
+  def self.track(raw_attribute)
+    with_statistics_available(raw_attribute) do |attribute|
+      SiteStatistic.update_all(["#{attribute} = #{attribute}+1"])
+    end
+  end
+
+  # Untracks specific attribute
+  #
+  # @param [String] raw_attribute must be one of the values listed in COUNTER_ATTRIBUTES
+  def self.untrack(raw_attribute)
+    with_statistics_available(raw_attribute) do |attribute|
+      SiteStatistic.update_all(["#{attribute} = #{attribute}-1 WHERE #{attribute} > 0"])
+    end
+  end
+
+  # Wrapper for track/untrack operations with basic validations and enforced requirements
+  #
+  # @param [String] raw_attribute must be one of the values listed in COUNTER_ATTRIBUTES
+  # @yield [String] attribute quoted to be used inside SQL / Arel query
+  def self.with_statistics_available(raw_attribute)
+    unless raw_attribute.in?(COUNTER_ATTRIBUTES)
+      raise ArgumentError, "Invalid attribute: '#{raw_attribute}' to '#{caller_locations(1, 1)[0].label}' method. " \
+                           "Valid attributes are: #{COUNTER_ATTRIBUTES.join(', ')}"
+    end
+
+    return unless available?
+
+    self.fetch # make sure record exists
+
+    attribute = self.connection.quote_column_name(raw_attribute)
+
+    # will be running on its own transaction context
+    yield(attribute)
+  end
+
+  # Returns a site statistic record with tracked information
+  #
+  # @return [SiteStatistic] record with tracked information
+  def self.fetch
+    SiteStatistic.transaction(requires_new: true) do
+      SiteStatistic.first_or_create!
+    end
+  rescue ActiveRecord::RecordNotUnique
+    retry
+  end
+
+  # Return whether required schema change is available
+  #
+  # This is needed in order to degrade gracefully when testing schema migrations
+  #
+  # @return [Boolean] whether schema is available
+  def self.available?
+    @available_flag ||= ActiveRecord::Migrator.current_version >= REQUIRED_SCHEMA_VERSION
+  end
+
+  # Resets cached column information
+  #
+  # This is called during schema migration specs, in order to reset internal cache state
+  def self.reset_column_information
+    @available_flag = nil
+
+    super
+  end
+end

db/migrate/20180629153018_create_site_statistics.rb

+class CreateSiteStatistics < ActiveRecord::Migration
+  include Gitlab::Database::MigrationHelpers
+
+  DOWNTIME = false
+
+  def up
+    create_table :site_statistics do |t|
+      t.integer :repositories_count, default: 0, null: false
+      t.integer :wikis_count, default: 0, null: false
+    end
+
+    execute('INSERT INTO site_statistics (id) VALUES(1)')
+  end
+
+  def down
+    drop_table :site_statistics
+  end
+end

db/post_migrate/20180706223200_populate_site_statistics.rb

+class PopulateSiteStatistics < ActiveRecord::Migration
+  include Gitlab::Database::MigrationHelpers
+
+  DOWNTIME = false
+
+  disable_ddl_transaction!
+
+  def up
+    transaction do
+      execute('SET LOCAL statement_timeout TO 0') if Gitlab::Database.postgresql? # see https://gitlab.com/gitlab-org/gitlab-ce/issues/48967
+
+      execute("UPDATE site_statistics SET repositories_count = (SELECT COUNT(*) FROM projects)")
+    end
+
+    transaction do
+      execute('SET LOCAL statement_timeout TO 0') if Gitlab::Database.postgresql? # see https://gitlab.com/gitlab-org/gitlab-ce/issues/48967
+
+      execute("UPDATE site_statistics SET wikis_count = (SELECT COUNT(*) FROM project_features WHERE wiki_access_level != 0)")
+    end
+  end
+
+  def down
+    # No downside in keeping the counter up-to-date
+  end
+end

db/schema.rb

@@ -2406,6 +2406,11 @@ ActiveRecord::Schema.define(version: 20180722103201) do
   add_index "services", ["project_id"], name: "index_services_on_project_id", using: :btree
   add_index "services", ["template"], name: "index_services_on_template", using: :btree
 
+  create_table "site_statistics", force: :cascade do |t|
+    t.integer "repositories_count", default: 0, null: false
+    t.integer "wikis_count", default: 0, null: false
+  end
+
   create_table "slack_integrations", force: :cascade do |t|
     t.integer "service_id", null: false
     t.string "team_id", null: false

ee/app/finders/geo/project_registry_finder.rb

 module Geo
   class ProjectRegistryFinder < RegistryFinder
     def count_repositories
-      current_node.projects.count
+      if selective_sync?
+        # We need to count only the selected projects according to the sync rule
+        current_node.projects.count
+      else
+        # Counting whole table can be expensive, so we use a counter cache instead
+        SiteStatistic.fetch.repositories_count
+      end
     end
 
     def count_wikis
-      current_node.projects.with_wiki_enabled.count
+      if selective_sync?
+        # We need to count only withing the selected projects according to the sync rule
+        current_node.projects.with_wiki_enabled.count
+      else
+        # Counting whole table can be expensive, so we use a counter cache instead
+        SiteStatistic.fetch.wikis_count
+      end
     end
 
     def count_synced_repositories

ee/changelogs/unreleased/6064-geo-sql-query-for-counting-projects-with-wikis-is-very-slow.yml

+---
+title: 'Geo: Improve Geo Status API performance with cached counters in SiteStatistic'
+merge_request: 6328
+author:
+type: performance

ee/spec/finders/geo/project_registry_finder_spec.rb

@@ -25,6 +25,54 @@ describe Geo::ProjectRegistryFinder, :geo do
   end
 
   shared_examples 'counts all the things' do
+    describe '#count_repositories' do
+      context 'without selective sync' do
+        it 'returns cached count values' do
+          SiteStatistic.fetch.update(repositories_count: 222)
+
+          expect(subject.count_repositories).to eq(222)
+        end
+      end
+
+      context 'with selective sync' do
+        before do
+          secondary.update!(selective_sync_type: 'namespaces', namespaces: [synced_group])
+        end
+
+        it 'returns only slice of selected repositories' do
+          create(:project, group: synced_group)
+          create(:project, group: synced_group)
+          create(:project)
+
+          expect(subject.count_repositories).to eq(2)
+        end
+      end
+    end
+
+    describe '#count_wikis' do
+      context 'without selective sync' do
+        it 'returns cached count values' do
+          SiteStatistic.fetch.update(wikis_count: 333)
+
+          expect(subject.count_wikis).to eq(333)
+        end
+      end
+
+      context 'with selective sync' do
+        before do
+          secondary.update!(selective_sync_type: 'namespaces', namespaces: [synced_group])
+        end
+
+        it 'returns only slice of selected repositories' do
+          create(:project, group: synced_group)
+          create(:project, :wiki_disabled, group: synced_group)
+          create(:project)
+
+          expect(subject.count_wikis).to eq(1)
+        end
+      end
+    end
+
     describe '#count_synced_repositories' do
       it 'delegates to #find_synced_repositories' do
         expect(subject).to receive(:find_synced_repositories).and_call_original

spec/factories/projects.rb

@@ -34,7 +34,7 @@ FactoryBot.define do
       builds_access_level = [evaluator.builds_access_level, evaluator.repository_access_level].min
       merge_requests_access_level = [evaluator.merge_requests_access_level, evaluator.repository_access_level].min
 
-      project.project_feature.update_columns(
+      project.project_feature.update(
         wiki_access_level: evaluator.wiki_access_level,
         builds_access_level: builds_access_level,
         snippets_access_level: evaluator.snippets_access_level,

spec/factories/site_statistics.rb

+FactoryBot.define do
+  factory :site_statistics, class: 'SiteStatistic' do
+    id 1
+    repositories_count 999
+    wikis_count 555
+  end
+end

spec/models/project_feature_spec.rb

@@ -128,4 +128,46 @@ describe ProjectFeature do
       end
     end
   end
+
+  context 'Site Statistics' do
+    set(:project_with_wiki) { create(:project, :wiki_enabled) }
+    set(:project_without_wiki) { create(:project, :wiki_disabled) }
+
+    context 'when creating a project' do
+      it 'tracks wiki availability when wikis are enabled by default' do
+        expect { create(:project) }.to change { SiteStatistic.fetch.wikis_count }.by(1)
+      end
+
+      it 'does not track wiki availability when wikis are disabled by default' do
+        expect { create(:project, :wiki_disabled) }.not_to change { SiteStatistic.fetch.wikis_count }
+      end
+    end
+
+    context 'when updating a project_feature' do
+      it 'untracks wiki availability when disabling wiki access' do
+        expect { project_with_wiki.project_feature.update_attribute(:wiki_access_level, ProjectFeature::DISABLED) }
+          .to change { SiteStatistic.fetch.wikis_count }.by(-1)
+      end
+
+      it 'tracks again wiki availability when re-enabling wiki access as public' do
+        expect { project_without_wiki.project_feature.update_attribute(:wiki_access_level, ProjectFeature::ENABLED) }
+          .to change { SiteStatistic.fetch.wikis_count }.by(1)
+      end
+
+      it 'tracks again wiki availability when re-enabling wiki access as private' do
+        expect { project_without_wiki.project_feature.update_attribute(:wiki_access_level, ProjectFeature::PRIVATE) }
+          .to change { SiteStatistic.fetch.wikis_count }.by(1)
+      end
+    end
+
+    context 'when removing a project' do
+      it 'untracks wiki availability when removing a project with previous wiki access' do
+        expect { project_with_wiki.destroy }.to change { SiteStatistic.fetch.wikis_count }.by(-1)
+      end
+
+      it 'does not untrack wiki availability when removing a project without wiki access' do
+        expect { project_without_wiki.destroy }.not_to change { SiteStatistic.fetch.wikis_count }
+      end
+    end
+  end
 end

spec/models/project_spec.rb

@@ -104,6 +104,22 @@ describe Project do
       end
     end
 
+    context 'Site Statistics' do
+      context 'when creating a new project' do
+        it 'tracks project in SiteStatistic' do
+          expect { create(:project) }.to change { SiteStatistic.fetch.repositories_count }.by(1)
+        end
+      end
+
+      context 'when deleting a project' do
+        it 'untracks project in SiteStatistic' do
+          project = create(:project)
+
+          expect { project.destroy }.to change { SiteStatistic.fetch.repositories_count }.by(-1)
+        end
+      end
+    end
+
     context 'updating cd_cd_settings' do
       it 'does not raise an error' do
         project = create(:project)

spec/models/site_statistic_spec.rb

+require 'spec_helper'
+
+describe SiteStatistic do
+  describe '.fetch' do
+    context 'existing record' do
+      it 'returns existing SiteStatistic model' do
+        statistics = create(:site_statistics)
+
+        expect(described_class.fetch).to be_a(described_class)
+        expect(described_class.fetch).to eq(statistics)
+      end
+    end
+
+    context 'non existing record' do
+      it 'creates a new SiteStatistic model' do
+        expect(described_class.first).to be_nil
+        expect(described_class.fetch).to be_a(described_class)
+      end
+    end
+  end
+
+  describe '.track' do
+    context 'with allowed attributes' do
+      let(:statistics) { create(:site_statistics) }
+
+      it 'increases the attribute counter' do
+        expect { described_class.track('repositories_count') }.to change { statistics.reload.repositories_count }.by(1)
+        expect { described_class.track('wikis_count') }.to change { statistics.reload.wikis_count }.by(1)
+      end
+
+      it 'doesnt increase the attribute counter when an exception happens during transaction' do
+        expect do
+          begin
+            described_class.transaction do
+              described_class.track('repositories_count')
+
+              raise StandardError
+            end
+          rescue StandardError
+            # no-op
+          end
+        end.not_to change { statistics.reload.repositories_count }
+      end
+    end
+
+    context 'with not allowed attributes' do
+      it 'returns error' do
+        expect { described_class.track('something_else') }.to raise_error(ArgumentError).with_message(/Invalid attribute: \'something_else\' to \'track\' method/)
+      end
+    end
+  end
+
+  describe '.untrack' do
+    context 'with allowed attributes' do
+      let(:statistics) { create(:site_statistics) }
+
+      it 'decreases the attribute counter' do
+        expect { described_class.untrack('repositories_count') }.to change { statistics.reload.repositories_count }.by(-1)
+        expect { described_class.untrack('wikis_count') }.to change { statistics.reload.wikis_count }.by(-1)
+      end
+
+      it 'doesnt decrease the attribute counter when an exception happens during transaction' do
+        expect do
+          begin
+            described_class.transaction do
+              described_class.track('repositories_count')
+
+              raise StandardError
+            end
+          rescue StandardError
+            # no-op
+          end
+        end.not_to change { described_class.fetch.repositories_count }
+      end
+    end
+
+    context 'with not allowed attributes' do
+      it 'returns error' do
+        expect { described_class.untrack('something_else') }.to raise_error(ArgumentError).with_message(/Invalid attribute: \'something_else\' to \'untrack\' method/)
+      end
+    end
+  end
+end