Includes logic to persist namespace statistics

- Add two new ActiveRecord models:
  - RootNamespaceStoragestatistics will persist root namespace statistics
  - NamespaceAggregationSchedule will save information when a new update
to the namespace statistics needs to be scheduled
- Inject into UpdateProjectStatistics concern a new callback that will
call an async job to insert a new row onto NamespaceAggregationSchedule
table
- When a new row is inserted a new job is scheduled. This job will
update call an specific service to update the statistics and after that
it will delete thee aggregated scheduled row
- The RefresherServices makes heavy use of arel to build composable
queries to update Namespace::RootStorageStatistics attributes.
- Add an extra worker to traverse pending rows on
NAmespace::AggregationSchedule table and schedule a worker for each one
of this rows.
- Add an extra worker to traverse pending rows on
NAmespace::AggregationSchedule table and schedule a worker for each one
of this rows

app/models/concerns/update_project_statistics.rb

@@ -19,9 +19,9 @@
 #
 # - `statistic_attribute` must be an ActiveRecord attribute
 # - The model must implement `project` and `project_id`. i.e. direct Project relationship or delegation
-#
 module UpdateProjectStatistics
   extend ActiveSupport::Concern
+  include AfterCommitQueue
 
   class_methods do
     attr_reader :project_statistics_name, :statistic_attribute
@@ -31,7 +31,6 @@ module UpdateProjectStatistics
     #
     # - project_statistics_name: A column of `ProjectStatistics` to update
     # - statistic_attribute: An attribute of the current model, default to `size`
-    #
     def update_project_statistics(project_statistics_name:, statistic_attribute: :size)
       @project_statistics_name = project_statistics_name
       @statistic_attribute = statistic_attribute
@@ -51,6 +50,7 @@ module UpdateProjectStatistics
       delta = read_attribute(attr).to_i - attribute_before_last_save(attr).to_i
 
       update_project_statistics(delta)
+      schedule_namespace_aggregation_worker
     end
 
     def update_project_statistics_attribute_changed?
@@ -59,6 +59,8 @@ module UpdateProjectStatistics
 
     def update_project_statistics_after_destroy
       update_project_statistics(-read_attribute(self.class.statistic_attribute).to_i)
+
+      schedule_namespace_aggregation_worker
     end
 
     def project_destroyed?
@@ -68,5 +70,18 @@ module UpdateProjectStatistics
     def update_project_statistics(delta)
       ProjectStatistics.increment_statistic(project_id, self.class.project_statistics_name, delta)
     end
+
+    def schedule_namespace_aggregation_worker
+      run_after_commit do
+        next unless schedule_aggregation_worker?
+
+        Namespaces::ScheduleAggregationWorker.perform_async(project.namespace_id)
+      end
+    end
+
+    def schedule_aggregation_worker?
+      !project.nil? &&
+        Feature.enabled?(:update_statistics_namespace, project.root_ancestor)
+    end
   end
 end

app/models/namespace.rb

@@ -293,6 +293,10 @@ class Namespace < ApplicationRecord
     end
   end
 
+  def aggregation_scheduled?
+    aggregation_schedule.present?
+  end
+
   private
 
   def parent_changed?

app/models/namespace/aggregation_schedule.rb

 # frozen_string_literal: true
 
 class Namespace::AggregationSchedule < ApplicationRecord
+  include AfterCommitQueue
+  include ExclusiveLeaseGuard
+
   self.primary_key = :namespace_id
 
+  DEFAULT_LEASE_TIMEOUT = 3.hours
+  REDIS_SHARED_KEY = 'gitlab:update_namespace_statistics_delay'.freeze
+
   belongs_to :namespace
+
+  after_create :schedule_root_storage_statistics
+
+  def self.delay_timeout
+    redis_timeout = Gitlab::Redis::SharedState.with do |redis|
+      redis.get(REDIS_SHARED_KEY)
+    end
+
+    redis_timeout.nil? ? DEFAULT_LEASE_TIMEOUT : redis_timeout.to_i
+  end
+
+  def schedule_root_storage_statistics
+    run_after_commit_or_now do
+      try_obtain_lease do
+        Namespaces::RootStatisticsWorker
+          .perform_async(namespace_id)
+
+        Namespaces::RootStatisticsWorker
+          .perform_in(self.class.delay_timeout, namespace_id)
+      end
+    end
+  end
+
+  private
+
+  # Used by ExclusiveLeaseGuard
+  def lease_timeout
+    self.class.delay_timeout
+  end
+
+  # Used by ExclusiveLeaseGuard
+  def lease_key
+    "namespace:namespaces_root_statistics:#{namespace_id}"
+  end
 end

app/models/namespace/root_storage_statistics.rb

 # frozen_string_literal: true
 
 class Namespace::RootStorageStatistics < ApplicationRecord
+  STATISTICS_ATTRIBUTES = %w(storage_size repository_size wiki_size lfs_objects_size build_artifacts_size packages_size).freeze
+
   self.primary_key = :namespace_id
 
   belongs_to :namespace
   has_one :route, through: :namespace
 
   delegate :all_projects, to: :namespace
+
+  def recalculate!
+    update!(attributes_from_project_statistics)
+  end
+
+  private
+
+  def attributes_from_project_statistics
+    from_project_statistics
+      .take
+      .attributes
+      .slice(*STATISTICS_ATTRIBUTES)
+  end
+
+  def from_project_statistics
+    all_projects
+      .joins('INNER JOIN project_statistics ps ON ps.project_id  = projects.id')
+      .select(
+        'COALESCE(SUM(ps.storage_size), 0) AS storage_size',
+        'COALESCE(SUM(ps.repository_size), 0) AS repository_size',
+        'COALESCE(SUM(ps.wiki_size), 0) AS wiki_size',
+        'COALESCE(SUM(ps.lfs_objects_size), 0) AS lfs_objects_size',
+        'COALESCE(SUM(ps.build_artifacts_size), 0) AS build_artifacts_size',
+        'COALESCE(SUM(ps.packages_size), 0) AS packages_size'
+      )
+  end
 end

app/services/namespaces/statistics_refresher_service.rb

+# frozen_string_literal: true
+
+module Namespaces
+  class StatisticsRefresherService
+    RefresherError = Class.new(StandardError)
+
+    def execute(root_namespace)
+      root_storage_statistics = find_or_create_root_storage_statistics(root_namespace.id)
+
+      root_storage_statistics.recalculate!
+    rescue ActiveRecord::ActiveRecordError => e
+      raise RefresherError.new(e.message)
+    end
+
+    private
+
+    def find_or_create_root_storage_statistics(root_namespace_id)
+      Namespace::RootStorageStatistics
+        .safe_find_or_create_by!(namespace_id: root_namespace_id)
+    end
+  end
+end

app/workers/all_queues.yml

@@ -26,6 +26,7 @@
 - cronjob:issue_due_scheduler
 - cronjob:prune_web_hook_logs
 - cronjob:schedule_migrate_external_diffs
+- cronjob:namespaces_prune_aggregation_schedules
 
 - gcp_cluster:cluster_install_app
 - gcp_cluster:cluster_patch_app
@@ -101,6 +102,9 @@
 - todos_destroyer:todos_destroyer_project_private
 - todos_destroyer:todos_destroyer_private_features
 
+- update_namespace_statistics:namespaces_schedule_aggregation
+- update_namespace_statistics:namespaces_root_statistics
+
 - object_pool:object_pool_create
 - object_pool:object_pool_schedule_join
 - object_pool:object_pool_join

app/workers/namespaces/prune_aggregation_schedules_worker.rb

+# frozen_string_literal: true
+
+module Namespaces
+  class PruneAggregationSchedulesWorker
+    include ApplicationWorker
+    include CronjobQueue
+
+    # Worker to prune pending rows on Namespace::AggregationSchedule
+    # It's scheduled to run once a day at 1:05am.
+    def perform
+      aggregation_schedules.find_each do |aggregation_schedule|
+        aggregation_schedule.schedule_root_storage_statistics
+      end
+    end
+
+    private
+
+    def aggregation_schedules
+      Namespace::AggregationSchedule.all
+    end
+  end
+end

app/workers/namespaces/root_statistics_worker.rb

+# frozen_string_literal: true
+
+module Namespaces
+  class RootStatisticsWorker
+    include ApplicationWorker
+
+    queue_namespace :update_namespace_statistics
+
+    def perform(namespace_id)
+      namespace = Namespace.find(namespace_id)
+
+      return unless update_statistics_enabled_for?(namespace) && namespace.aggregation_scheduled?
+
+      Namespaces::StatisticsRefresherService.new.execute(namespace)
+
+      namespace.aggregation_schedule.destroy
+    rescue ::Namespaces::StatisticsRefresherService::RefresherError, ActiveRecord::RecordNotFound => ex
+      log_error(namespace.full_path, ex.message) if namespace
+    end
+
+    private
+
+    def log_error(namespace_path, error_message)
+      Gitlab::SidekiqLogger.error("Namespace statistics can't be updated for #{namespace_path}: #{error_message}")
+    end
+
+    def update_statistics_enabled_for?(namespace)
+      Feature.enabled?(:update_statistics_namespace, namespace)
+    end
+  end
+end

app/workers/namespaces/schedule_aggregation_worker.rb

+# frozen_string_literal: true
+
+module Namespaces
+  class ScheduleAggregationWorker
+    include ApplicationWorker
+
+    queue_namespace :update_namespace_statistics
+
+    def perform(namespace_id)
+      return unless aggregation_schedules_table_exists?
+
+      namespace = Namespace.find(namespace_id)
+      root_ancestor = namespace.root_ancestor
+
+      return unless update_statistics_enabled_for?(root_ancestor) && !root_ancestor.aggregation_scheduled?
+
+      Namespace::AggregationSchedule.safe_find_or_create_by!(namespace_id: root_ancestor.id)
+    rescue ActiveRecord::RecordNotFound
+      log_error(namespace_id)
+    end
+
+    private
+
+    # On db/post_migrate/20180529152628_schedule_to_archive_legacy_traces.rb
+    # traces are archived through build.trace.archive, which in consequence
+    # calls UpdateProjectStatistics#schedule_namespace_statistics_worker.
+    #
+    # The migration and specs fails since NamespaceAggregationSchedule table
+    # does not exist at that point.
+    # https://gitlab.com/gitlab-org/gitlab-ce/issues/50712
+    def aggregation_schedules_table_exists?
+      return true unless Rails.env.test?
+
+      Namespace::AggregationSchedule.table_exists?
+    end
+
+    def log_error(root_ancestor_id)
+      Gitlab::SidekiqLogger.error("Namespace can't be scheduled for aggregation: #{root_ancestor_id} does not exist")
+    end
+
+    def update_statistics_enabled_for?(root_ancestor)
+      Feature.enabled?(:update_statistics_namespace, root_ancestor)
+    end
+  end
+end

config/initializers/1_settings.rb

@@ -441,6 +441,9 @@ Settings.cron_jobs['prune_web_hook_logs_worker']['job_class'] = 'PruneWebHookLog
 Settings.cron_jobs['schedule_migrate_external_diffs_worker'] ||= Settingslogic.new({})
 Settings.cron_jobs['schedule_migrate_external_diffs_worker']['cron'] ||= '15 * * * *'
 Settings.cron_jobs['schedule_migrate_external_diffs_worker']['job_class'] = 'ScheduleMigrateExternalDiffsWorker'
+Settings.cron_jobs['namespaces_prune_aggregation_schedules_worker'] ||= Settingslogic.new({})
+Settings.cron_jobs['namespaces_prune_aggregation_schedules_worker']['cron'] ||= '5 1 * * *'
+Settings.cron_jobs['namespaces_prune_aggregation_schedules_worker']['job_class'] = 'Namespaces::PruneAggregationSchedulesWorker'
 
 Gitlab.ee do
   Settings.cron_jobs['clear_shared_runners_minutes_worker'] ||= Settingslogic.new({})

config/sidekiq_queues.yml

@@ -94,6 +94,7 @@
   - [migrate_external_diffs, 1]
   - [update_project_statistics, 1]
   - [phabricator_import_import_tasks, 1]
+  - [update_namespace_statistics, 1]
 
   # EE-specific queues
   - [ldap_group_sync, 2]

spec/factories/project_statistics.rb

@@ -6,5 +6,20 @@ FactoryBot.define do
       # statistics are automatically created when a project is created
       project&.statistics || new
     end
+
+    transient do
+      with_data { false }
+      size_multiplier { 1 }
+    end
+
+    after(:build) do |project_statistics, evaluator|
+      if evaluator.with_data
+        project_statistics.repository_size = evaluator.size_multiplier
+        project_statistics.wiki_size = evaluator.size_multiplier * 2
+        project_statistics.lfs_objects_size = evaluator.size_multiplier * 3
+        project_statistics.build_artifacts_size = evaluator.size_multiplier * 4
+        project_statistics.packages_size = evaluator.size_multiplier * 5
+      end
+    end
   end
 end

spec/models/namespace/aggregation_schedule_spec.rb

@@ -2,6 +2,77 @@
 
 require 'spec_helper'
 
-RSpec.describe Namespace::AggregationSchedule, type: :model do
+RSpec.describe Namespace::AggregationSchedule, :clean_gitlab_redis_shared_state, type: :model do
+  include ExclusiveLeaseHelpers
+
   it { is_expected.to belong_to :namespace }
+
+  describe '.delay_timeout' do
+    context 'when timeout is set on redis' do
+      it 'uses personalized timeout' do
+        Gitlab::Redis::SharedState.with do |redis|
+          redis.set(described_class::REDIS_SHARED_KEY, 1.hour)
+        end
+
+        expect(described_class.delay_timeout).to eq(1.hour)
+      end
+    end
+
+    context 'when timeout is not set on redis' do
+      it 'uses default timeout' do
+        expect(described_class.delay_timeout).to eq(3.hours)
+      end
+    end
+  end
+
+  describe '#schedule_root_storage_statistics' do
+    let(:namespace) { create(:namespace) }
+    let(:aggregation_schedule) { namespace.build_aggregation_schedule }
+    let(:lease_key) { "namespace:namespaces_root_statistics:#{namespace.id}" }
+
+    context "when we can't obtain the lease" do
+      it 'does not schedule the workers' do
+        stub_exclusive_lease_taken(lease_key, timeout: described_class::DEFAULT_LEASE_TIMEOUT)
+
+        expect(Namespaces::RootStatisticsWorker)
+          .not_to receive(:perform_async)
+
+        expect(Namespaces::RootStatisticsWorker)
+          .not_to receive(:perform_in)
+
+        aggregation_schedule.save!
+      end
+    end
+
+    context 'when we can obtain the lease' do
+      it 'schedules a root storage statistics after create' do
+        stub_exclusive_lease(lease_key, timeout: described_class::DEFAULT_LEASE_TIMEOUT)
+
+        expect(Namespaces::RootStatisticsWorker)
+          .to receive(:perform_async).once
+
+        expect(Namespaces::RootStatisticsWorker)
+          .to receive(:perform_in).once
+          .with(described_class::DEFAULT_LEASE_TIMEOUT, aggregation_schedule.namespace_id )
+
+        aggregation_schedule.save!
+      end
+    end
+
+    context 'with a personalized lease timeout' do
+      before do
+        Gitlab::Redis::SharedState.with do |redis|
+          redis.set(described_class::REDIS_SHARED_KEY, 1.hour)
+        end
+      end
+
+      it 'uses a personalized time' do
+        expect(Namespaces::RootStatisticsWorker)
+          .to receive(:perform_in)
+          .with(1.hour, aggregation_schedule.namespace_id)
+
+        aggregation_schedule.save!
+      end
+    end
+  end
 end

spec/models/namespace/root_storage_statistics_spec.rb

@@ -7,4 +7,69 @@ RSpec.describe Namespace::RootStorageStatistics, type: :model do
   it { is_expected.to have_one(:route).through(:namespace) }
 
   it { is_expected.to delegate_method(:all_projects).to(:namespace) }
+
+  describe '#recalculate!' do
+    let(:namespace) { create(:group) }
+    let(:root_storage_statistics) { create(:namespace_root_storage_statistics, namespace: namespace) }
+
+    let(:project1) { create(:project, namespace: namespace) }
+    let(:project2) { create(:project, namespace: namespace) }
+
+    let!(:stat1) { create(:project_statistics, project: project1, with_data: true, size_multiplier: 100) }
+    let!(:stat2) { create(:project_statistics, project: project2, with_data: true, size_multiplier: 200) }
+
+    shared_examples 'data refresh' do
+      it 'aggregates project statistics' do
+        root_storage_statistics.recalculate!
+
+        root_storage_statistics.reload
+
+        total_repository_size = stat1.repository_size + stat2.repository_size
+        total_wiki_size = stat1.wiki_size + stat2.wiki_size
+        total_lfs_objects_size = stat1.lfs_objects_size + stat2.lfs_objects_size
+        total_build_artifacts_size = stat1.build_artifacts_size + stat2.build_artifacts_size
+        total_packages_size = stat1.packages_size + stat2.packages_size
+        total_storage_size = stat1.storage_size + stat2.storage_size
+
+        expect(root_storage_statistics.repository_size).to eq(total_repository_size)
+        expect(root_storage_statistics.wiki_size).to eq(total_wiki_size)
+        expect(root_storage_statistics.lfs_objects_size).to eq(total_lfs_objects_size)
+        expect(root_storage_statistics.build_artifacts_size).to eq(total_build_artifacts_size)
+        expect(root_storage_statistics.packages_size).to eq(total_packages_size)
+        expect(root_storage_statistics.storage_size).to eq(total_storage_size)
+      end
+
+      it 'works when there are no projects' do
+        Project.delete_all
+
+        root_storage_statistics.recalculate!
+
+        root_storage_statistics.reload
+        expect(root_storage_statistics.repository_size).to eq(0)
+        expect(root_storage_statistics.wiki_size).to eq(0)
+        expect(root_storage_statistics.lfs_objects_size).to eq(0)
+        expect(root_storage_statistics.build_artifacts_size).to eq(0)
+        expect(root_storage_statistics.packages_size).to eq(0)
+        expect(root_storage_statistics.storage_size).to eq(0)
+      end
+    end
+
+    it_behaves_like 'data refresh'
+
+    context 'with subgroups', :nested_groups do
+      let(:subgroup1) { create(:group, parent: namespace)}
+      let(:subgroup2) { create(:group, parent: subgroup1)}
+
+      let(:project1) { create(:project, namespace: subgroup1) }
+      let(:project2) { create(:project, namespace: subgroup2) }
+
+      it_behaves_like 'data refresh'
+    end
+
+    context 'with a personal namespace' do
+      let(:namespace) { create(:user).namespace }
+
+      it_behaves_like 'data refresh'
+    end
+  end
 end

spec/models/namespace_spec.rb

@@ -837,4 +837,20 @@ describe Namespace do
       it { is_expected.to be_falsy }
     end
   end
+
+  describe '#aggregation_scheduled?' do
+    let(:namespace) { create(:namespace) }
+
+    subject { namespace.aggregation_scheduled? }
+
+    context 'with an aggregation scheduled association' do
+      let(:namespace) { create(:namespace, :with_aggregation_schedule) }
+
+      it { is_expected.to be_truthy }
+    end
+
+    context 'without an aggregation scheduled association' do
+      it { is_expected.to be_falsy }
+    end
+  end
 end

spec/services/namespaces/statistics_refresher_service_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Namespaces::StatisticsRefresherService, '#execute' do
+  let(:group) { create(:group) }
+  let(:projects) { create_list(:project, 5, namespace: group) }
+  let(:service) { described_class.new }
+
+  context 'without a root storage statistics relation' do
+    it 'creates one' do
+      expect do
+        service.execute(group)
+      end.to change(Namespace::RootStorageStatistics, :count).by(1)
+
+      expect(group.reload.root_storage_statistics).to be_present
+    end
+
+    it 'recalculate the namespace statistics' do
+      expect_any_instance_of(Namespace::RootStorageStatistics).to receive(:recalculate!).once
+
+      service.execute(group)
+    end
+  end
+
+  context 'with a root storage statistics relation' do
+    before do
+      Namespace::AggregationSchedule.safe_find_or_create_by!(namespace_id: group.id)
+    end
+
+    it 'does not create one' do
+      expect do
+        service.execute(group)
+      end.not_to change(Namespace::RootStorageStatistics, :count)
+    end
+
+    it 'recalculate the namespace statistics' do
+      expect(Namespace::RootStorageStatistics)
+        .to receive(:safe_find_or_create_by!).with({ namespace_id: group.id })
+        .and_return(group.root_storage_statistics)
+
+      service.execute(group)
+    end
+  end
+
+  context 'when something goes wrong' do
+    before do
+      allow_any_instance_of(Namespace::RootStorageStatistics)
+        .to receive(:recalculate!).and_raise(ActiveRecord::ActiveRecordError)
+    end
+
+    it 'raises RefreshError' do
+      expect do
+        service.execute(group)
+      end.to raise_error(Namespaces::StatisticsRefresherService::RefresherError)
+    end
+  end
+end

spec/support/shared_examples/models/update_project_statistics_shared_examples.rb

@@ -25,16 +25,36 @@ shared_examples_for 'UpdateProjectStatistics' do
         .to change { reload_stat }
         .by(delta)
     end
+
+    it 'schedules a namespace statistics worker' do
+      expect(Namespaces::ScheduleAggregationWorker)
+        .to receive(:perform_async).once
+
+      subject.save!
+    end
+
+    context 'when feature flag is disabled for the namespace' do
+      it 'does not schedules a namespace statistics worker' do
+        namespace = subject.project.root_ancestor
+
+        stub_feature_flags(update_statistics_namespace: false, namespace: namespace)
+
+        expect(Namespaces::ScheduleAggregationWorker)
+          .not_to receive(:perform_async)
+
+        subject.save!
+      end
+    end
   end
 
   context 'when updating' do
+    let(:delta) { 42 }
+
     before do
       subject.save!
     end
 
     it 'updates project statistics' do
-      delta = 42
-
       expect(ProjectStatistics)
         .to receive(:increment_statistic)
         .and_call_original
@@ -45,6 +65,42 @@ shared_examples_for 'UpdateProjectStatistics' do
         .to change { reload_stat }
         .by(delta)
     end
+
+    it 'schedules a namespace statistics worker' do
+      expect(Namespaces::ScheduleAggregationWorker)
+        .to receive(:perform_async).once
+
+      subject.write_attribute(statistic_attribute, read_attribute + delta)
+      subject.save!
+    end
+
+    it 'avoids N + 1 queries' do
+      subject.write_attribute(statistic_attribute, read_attribute + delta)
+
+      control_count = ActiveRecord::QueryRecorder.new do
+        subject.save!
+      end
+
+      subject.write_attribute(statistic_attribute, read_attribute + delta)
+
+      expect do
+        subject.save!
+      end.not_to exceed_query_limit(control_count)
+    end
+
+    context 'when the feature flag is disabled for the namespace' do
+      it 'does not schedule a namespace statistics worker' do
+        namespace = subject.project.root_ancestor
+
+        stub_feature_flags(update_statistics_namespace: false, namespace: namespace)
+
+        expect(Namespaces::ScheduleAggregationWorker)
+          .not_to receive(:perform_async)
+
+        subject.write_attribute(statistic_attribute, read_attribute + delta)
+        subject.save!
+      end
+    end
   end
 
   context 'when destroying' do
@@ -59,11 +115,18 @@ shared_examples_for 'UpdateProjectStatistics' do
         .to receive(:increment_statistic)
         .and_call_original
 
-      expect { subject.destroy }
+      expect { subject.destroy! }
         .to change { reload_stat }
         .by(delta)
     end
 
+    it 'schedules a namespace statistics worker' do
+      expect(Namespaces::ScheduleAggregationWorker)
+        .to receive(:perform_async).once
+
+      subject.destroy!
+    end
+
     context 'when it is destroyed from the project level' do
       it 'does not update the project statistics' do
         expect(ProjectStatistics)
@@ -72,6 +135,27 @@ shared_examples_for 'UpdateProjectStatistics' do
         project.update(pending_delete: true)
         project.destroy!
       end
+
+      it 'does not schedule a namespace statistics worker' do
+        expect(Namespaces::ScheduleAggregationWorker)
+          .not_to receive(:perform_async)
+
+        project.update(pending_delete: true)
+        project.destroy!
+      end
+    end
+
+    context 'when feature flag is disabled for the namespace' do
+      it 'does not schedule a namespace statistics worker' do
+        namespace = subject.project.root_ancestor
+
+        stub_feature_flags(update_statistics_namespace: false, namespace: namespace)
+
+        expect(Namespaces::ScheduleAggregationWorker)
+          .not_to receive(:perform_async)
+
+        subject.destroy!
+      end
     end
   end
 end

spec/workers/namespaces/prune_aggregation_schedules_worker_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Namespaces::PruneAggregationSchedulesWorker, '#perform', :clean_gitlab_redis_shared_state do
+  include ExclusiveLeaseHelpers
+
+  let(:namespaces) { create_list(:namespace, 5, :with_aggregation_schedule) }
+  let(:timeout) { Namespace::AggregationSchedule::DEFAULT_LEASE_TIMEOUT }
+
+  subject(:worker) { described_class.new }
+
+  before do
+    allow(Namespaces::RootStatisticsWorker)
+      .to receive(:perform_async).and_return(nil)
+
+    allow(Namespaces::RootStatisticsWorker)
+      .to receive(:perform_in).and_return(nil)
+
+    namespaces.each do |namespace|
+      lease_key = "namespace:namespaces_root_statistics:#{namespace.id}"
+      stub_exclusive_lease(lease_key, timeout: timeout)
+    end
+  end
+
+  it 'schedules a worker per pending aggregation' do
+    expect(Namespaces::RootStatisticsWorker)
+      .to receive(:perform_async).exactly(5).times
+
+    expect(Namespaces::RootStatisticsWorker)
+      .to receive(:perform_in).exactly(5).times
+
+    worker.perform
+  end
+end

spec/workers/namespaces/root_statistics_worker_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Namespaces::RootStatisticsWorker, '#perform' do
+  let(:group) { create(:group, :with_aggregation_schedule) }
+
+  subject(:worker) { described_class.new }
+
+  context 'with a namespace' do
+    it 'executes refresher service' do
+      expect_any_instance_of(Namespaces::StatisticsRefresherService)
+        .to receive(:execute)
+
+      worker.perform(group.id)
+    end
+
+    it 'deletes namespace aggregated schedule row' do
+      worker.perform(group.id)
+
+      expect(group.reload.aggregation_schedule).to be_nil
+    end
+
+    context 'when something goes wrong when updating' do
+      before do
+        allow_any_instance_of(Namespaces::StatisticsRefresherService)
+          .to receive(:execute)
+          .and_raise(Namespaces::StatisticsRefresherService::RefresherError, 'error')
+      end
+
+      it 'does not delete the aggregation schedule' do
+        worker.perform(group.id)
+
+        expect(group.reload.aggregation_schedule).to be_present
+      end
+
+      it 'logs the error' do
+        # A Namespace::RootStatisticsWorker is scheduled when
+        # a Namespace::AggregationSchedule is created, so having
+        # create(:group, :with_aggregation_schedule), will execute
+        # another worker
+        allow_any_instance_of(Namespace::AggregationSchedule)
+          .to receive(:schedule_root_storage_statistics).and_return(nil)
+
+        expect(Gitlab::SidekiqLogger).to receive(:error).once
+
+        worker.perform(group.id)
+      end
+    end
+  end
+
+  context 'with no namespace' do
+    before do
+      group.destroy
+    end
+
+    it 'does not execute the refresher service' do
+      expect_any_instance_of(Namespaces::StatisticsRefresherService)
+        .not_to receive(:execute)
+
+      worker.perform(group.id)
+    end
+  end
+
+  context 'with a namespace with no aggregation scheduled' do
+    before do
+      group.aggregation_schedule.destroy
+    end
+
+    it 'does not execute the refresher service' do
+      expect_any_instance_of(Namespaces::StatisticsRefresherService)
+        .not_to receive(:execute)
+
+      worker.perform(group.id)
+    end
+  end
+
+  context 'when update_statistics_namespace is off' do
+    it 'does not create a new one' do
+      stub_feature_flags(update_statistics_namespace: false, namespace: group)
+
+      expect_any_instance_of(Namespaces::StatisticsRefresherService)
+        .not_to receive(:execute)
+
+      worker.perform(group.id)
+    end
+  end
+end

spec/workers/namespaces/schedule_aggregation_worker_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Namespaces::ScheduleAggregationWorker, '#perform' do
+  let(:group) { create(:group) }
+
+  subject(:worker) { described_class.new }
+
+  context 'when group is the root ancestor' do
+    context 'when aggregation schedule exists' do
+      it 'does not create a new one' do
+        Namespace::AggregationSchedule.safe_find_or_create_by!(namespace_id: group.id)
+
+        expect do
+          worker.perform(group.id)
+        end.not_to change(Namespace::AggregationSchedule, :count)
+      end
+    end
+
+    context 'when update_statistics_namespace is off' do
+      it 'does not create a new one' do
+        stub_feature_flags(update_statistics_namespace: false, namespace: group)
+
+        expect do
+          worker.perform(group.id)
+        end.not_to change(Namespace::AggregationSchedule, :count)
+      end
+    end
+
+    context 'when aggregation schedule does not exist' do
+      it 'creates one' do
+        allow_any_instance_of(Namespace::AggregationSchedule)
+          .to receive(:schedule_root_storage_statistics).and_return(nil)
+
+        expect do
+          worker.perform(group.id)
+        end.to change(Namespace::AggregationSchedule, :count).by(1)
+
+        expect(group.aggregation_schedule).to be_present
+      end
+    end
+  end
+
+  context 'when group is not the root ancestor' do
+    let(:parent_group) { create(:group) }
+    let(:group) { create(:group, parent: parent_group) }
+
+    it 'creates an aggregation schedule for the root' do
+      allow_any_instance_of(Namespace::AggregationSchedule)
+        .to receive(:schedule_root_storage_statistics).and_return(nil)
+
+      worker.perform(group.id)
+
+      expect(parent_group.aggregation_schedule).to be_present
+    end
+  end
+
+  context 'when namespace does not exist' do
+    it 'logs the error' do
+      expect(Gitlab::SidekiqLogger).to receive(:error).once
+
+      worker.perform(12345)
+    end
+  end
+end