Schedule namespace aggregation in other contexts62214-follow-up-schedule-aggregation-in-other-contexts

Schedules a Namespace::AggregationSchedule worker if some of the project
statistics are refreshed.

The worker is only executed if the feature flag is enabled.

3 files changed, 80 insertions, 2 deletions

diff --git a/app/models/project_statistics.rb b/app/models/project_statistics.rb
index 8a179b4d56d..3802d258664 100644
--- a/app/models/project_statistics.rb
+++ b/app/models/project_statistics.rb
@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 class ProjectStatistics < ApplicationRecord
+  include AfterCommitQueue
+
   belongs_to :project
   belongs_to :namespace
 
@@ -15,6 +17,7 @@ class ProjectStatistics < ApplicationRecord
 
   COLUMNS_TO_REFRESH = [:repository_size, :wiki_size, :lfs_objects_size, :commit_count].freeze
   INCREMENTABLE_COLUMNS = { build_artifacts_size: %i[storage_size], packages_size: %i[storage_size] }.freeze
+  NAMESPACE_RELATABLE_COLUMNS = [:repository_size, :wiki_size, :lfs_objects_size].freeze
 
   scope :for_project_ids, ->(project_ids) { where(project_id: project_ids) }
 
@@ -22,13 +25,17 @@ class ProjectStatistics < ApplicationRecord
     repository_size + lfs_objects_size
   end
 
-  def refresh!(only: nil)
+  def refresh!(only: [])
     COLUMNS_TO_REFRESH.each do |column, generator|
-      if only.blank? || only.include?(column)
+      if only.empty? || only.include?(column)
         public_send("update_#{column}") # rubocop:disable GitlabSecurity/PublicSend
       end
     end
 
+    if only.empty? || only.any? { |column| NAMESPACE_RELATABLE_COLUMNS.include?(column) }
+      schedule_namespace_aggregation_worker
+    end
+
     save!
   end
 
@@ -81,4 +88,18 @@ class ProjectStatistics < ApplicationRecord
 
     update_all(updates.join(', '))
   end
+
+  private
+
+  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?
+    Feature.enabled?(:update_statistics_namespace, project&.root_ancestor)
+  end
 end
diff --git a/spec/models/project_statistics_spec.rb b/spec/models/project_statistics_spec.rb
index 1cb49d83ffa..db3e4902c64 100644
--- a/spec/models/project_statistics_spec.rb
+++ b/spec/models/project_statistics_spec.rb
@@ -135,6 +135,49 @@ describe ProjectStatistics do
         expect(statistics.wiki_size).to eq(0)
       end
     end
+
+    context 'when the column is namespace relatable' do
+      let(:namespace) { create(:group) }
+      let(:project) { create(:project, namespace: namespace) }
+
+      context 'when the feature flag is off' do
+        it 'does not schedule the aggregation worker' do
+          stub_feature_flags(update_statistics_namespace: false, namespace: namespace)
+
+          expect(Namespaces::ScheduleAggregationWorker)
+            .not_to receive(:perform_async)
+
+          statistics.refresh!(only: [:lfs_objects_size])
+        end
+      end
+
+      context 'when the feature flag is on' do
+        it 'schedules the aggregation worker' do
+          expect(Namespaces::ScheduleAggregationWorker)
+            .to receive(:perform_async)
+
+          statistics.refresh!(only: [:lfs_objects_size])
+        end
+      end
+
+      context 'when no argument is passed' do
+        it 'schedules the aggregation worker' do
+          expect(Namespaces::ScheduleAggregationWorker)
+            .to receive(:perform_async)
+
+          statistics.refresh!
+        end
+      end
+    end
+
+    context 'when the column is not namespace relatable' do
+      it 'does not schedules an aggregation worker' do
+        expect(Namespaces::ScheduleAggregationWorker)
+          .not_to receive(:perform_async)
+
+        statistics.refresh!(only: [:commit_count])
+      end
+    end
   end
 
   describe '#update_commit_count' do
diff --git a/spec/workers/project_cache_worker_spec.rb b/spec/workers/project_cache_worker_spec.rb
index 51afb076da1..edc55920b8e 100644
--- a/spec/workers/project_cache_worker_spec.rb
+++ b/spec/workers/project_cache_worker_spec.rb
@@ -36,6 +36,11 @@ describe ProjectCacheWorker do
     end
 
     context 'with an existing project' do
+      before do
+        lease_key = "namespace:namespaces_root_statistics:#{project.namespace_id}"
+        stub_exclusive_lease_taken(lease_key, timeout: Namespace::AggregationSchedule::DEFAULT_LEASE_TIMEOUT)
+      end
+
       it 'refreshes the method caches' do
         expect_any_instance_of(Repository).to receive(:refresh_method_caches)
           .with(%i(readme))
@@ -81,6 +86,10 @@ describe ProjectCacheWorker do
 
         expect(UpdateProjectStatisticsWorker).not_to receive(:perform_in)
 
+        expect(Namespaces::ScheduleAggregationWorker)
+          .not_to receive(:perform_async)
+          .with(project.namespace_id)
+
         worker.update_statistics(project, statistics)
       end
     end
@@ -98,6 +107,11 @@ describe ProjectCacheWorker do
           .with(lease_timeout, project.id, statistics)
           .and_call_original
 
+        expect(Namespaces::ScheduleAggregationWorker)
+          .to receive(:perform_async)
+          .with(project.namespace_id)
+          .twice
+
         worker.update_statistics(project, statistics)
       end
     end