1 files changed, 41 insertions, 3 deletions

diff --git a/app/models/projects/build_artifacts_size_refresh.rb b/app/models/projects/build_artifacts_size_refresh.rb
index 2ffc7478178..b791cb1254c 100644
--- a/app/models/projects/build_artifacts_size_refresh.rb
+++ b/app/models/projects/build_artifacts_size_refresh.rb
@@ -7,16 +7,34 @@ module Projects
 
     STALE_WINDOW = 2.hours
 
+    # This delay is set to 10 minutes to accommodate any ongoing
+    # deletion that might have happened.
+    # The delete on the database may have been committed before
+    # the refresh completed its batching. If the resulting decrement is
+    # pushed into Redis after the refresh has ended, it would result in net negative value.
+    # The delay is needed to ensure this negative value is ignored.
+    FINALIZE_DELAY = 10.minutes
+
     self.table_name = 'project_build_artifacts_size_refreshes'
 
+    COUNTER_ATTRIBUTE_NAME = :build_artifacts_size
+
     belongs_to :project
 
     validates :project, presence: true
 
+    # The refresh of the project statistics counter is performed in 4 stages:
+    # 1. created - The refresh is on the queue to be processed by Projects::RefreshBuildArtifactsSizeStatisticsWorker
+    # 2. running - The refresh is ongoing. The project statistics counter switches to the temporary refresh counter key.
+    #    Counter increments are deduplicated.
+    # 3. pending - The refresh is pending to be picked up by Projects::RefreshBuildArtifactsSizeStatisticsWorker again.
+    # 4. finalizing - The refresh has finished summing existing job artifact size into the refresh counter key.
+    #    The sum will need to be moved into the counter key.
     STATES = {
       created: 1,
       running: 2,
-      pending: 3
+      pending: 3,
+      finalizing: 4
     }.freeze
 
     state_machine :state, initial: :created do
@@ -24,6 +42,7 @@ module Projects
       state :created, value: STATES[:created]
       state :running, value: STATES[:running]
       state :pending, value: STATES[:pending]
+      state :finalizing, value: STATES[:finalizing]
 
       event :process do
         transition [:created, :pending, :running] => :running
@@ -33,7 +52,10 @@ module Projects
         transition running: :pending
       end
 
-      # set it only the first time we execute the refresh
+      event :schedule_finalize do
+        transition running: :finalizing
+      end
+
       before_transition created: :running do |refresh|
         refresh.reset_project_statistics!
         refresh.refresh_started_at = Time.zone.now
@@ -47,6 +69,10 @@ module Projects
       before_transition running: :pending do |refresh, transition|
         refresh.last_job_artifact_id = transition.args.first
       end
+
+      before_transition running: :finalizing do |refresh, transition|
+        refresh.schedule_finalize_worker
+      end
     end
 
     scope :stale, -> { with_state(:running).where('updated_at < ?', STALE_WINDOW.ago) }
@@ -80,7 +106,7 @@ module Projects
     end
 
     def reset_project_statistics!
-      project.statistics.reset_counter!(:build_artifacts_size)
+      project.statistics.initiate_refresh!(COUNTER_ATTRIBUTE_NAME)
     end
 
     def next_batch(limit:)
@@ -95,6 +121,18 @@ module Projects
       !created?
     end
 
+    def finalize!
+      project.statistics.finalize_refresh(COUNTER_ATTRIBUTE_NAME)
+
+      destroy!
+    end
+
+    def schedule_finalize_worker
+      run_after_commit do
+        Projects::FinalizeProjectStatisticsRefreshWorker.perform_in(FINALIZE_DELAY, self.class.to_s, id)
+      end
+    end
+
     private
 
     def schedule_namespace_aggregation_worker