Prepare CI minutes consumption update to be idempotent

ee/app/services/ci/minutes/update_project_and_namespace_usage_service.rb

@@ -5,9 +5,12 @@ module Ci
     class UpdateProjectAndNamespaceUsageService
       include Gitlab::Utils::StrongMemoize
 
-      def initialize(project_id, namespace_id)
+      IDEMPOTENCY_CACHE_TTL = 12.hours
+
+      def initialize(project_id, namespace_id, build_id = nil)
         @project_id = project_id
         @namespace_id = namespace_id
+        @build_id = build_id
         # TODO(issue 335885): Use project_id only and don't query for projects which may be deleted
         @project = Project.find_by_id(project_id)
       end
@@ -16,22 +19,24 @@ module Ci
       def execute(consumption)
         legacy_track_usage_of_monthly_minutes(consumption)
 
-        preload_minutes_usage_data!
-
-        ApplicationRecord.transaction do
+        # TODO: fix this condition after the next deployment when `build_id`
+        # is made a mandatory argument.
+        # https://gitlab.com/gitlab-org/gitlab/-/issues/331785
+        if @build_id
+          ensure_idempotency { track_usage_of_monthly_minutes(consumption) }
+        else
           track_usage_of_monthly_minutes(consumption)
+        end
 
         send_minutes_email_notification
       end
+
+      def idempotency_cache_key
+        "ci_minutes_usage:#{@project_id}:#{@build_id}:updated"
       end
 
       private
 
-      def preload_minutes_usage_data!
-        project_usage
-        namespace_usage
-      end
-
       def send_minutes_email_notification
         # `perform reset` on `project` because `Namespace#namespace_statistics` will otherwise return stale data.
         # TODO(issue 335885): Remove @project
@@ -46,9 +51,15 @@ module Ci
       end
 
       def track_usage_of_monthly_minutes(consumption)
+        # preload minutes usage data outside of transaction
+        project_usage
+        namespace_usage
+
+        ::Ci::Minutes::NamespaceMonthlyUsage.transaction do
           ::Ci::Minutes::NamespaceMonthlyUsage.increase_usage(namespace_usage, consumption) if namespace_usage
           ::Ci::Minutes::ProjectMonthlyUsage.increase_usage(project_usage, consumption) if project_usage
         end
+      end
 
       def update_legacy_project_minutes(consumption_in_seconds)
         if project_statistics
@@ -88,6 +99,28 @@ module Ci
         rescue ActiveRecord::NotNullViolation, ActiveRecord::RecordInvalid
         end
       end
+
+      # Ensure we only add the CI minutes consumption once for the given build
+      # even if the worker is retried.
+      def ensure_idempotency
+        return if already_completed?
+
+        yield
+
+        mark_as_completed!
+      end
+
+      def mark_as_completed!
+        Gitlab::Redis::SharedState.with do |redis|
+          redis.set(idempotency_cache_key, 1, ex: IDEMPOTENCY_CACHE_TTL)
+        end
+      end
+
+      def already_completed?
+        Gitlab::Redis::SharedState.with do |redis|
+          redis.exists(idempotency_cache_key)
+        end
+      end
     end
   end
 end

ee/app/workers/ci/minutes/update_project_and_namespace_usage_worker.rb

@@ -9,8 +9,10 @@ module Ci
       urgency :low
       data_consistency :always # primarily performs writes
 
-      def perform(consumption, project_id, namespace_id)
-        ::Ci::Minutes::UpdateProjectAndNamespaceUsageService.new(project_id, namespace_id).execute(consumption)
+      def perform(consumption, project_id, namespace_id, build_id = nil)
+        ::Ci::Minutes::UpdateProjectAndNamespaceUsageService
+          .new(project_id, namespace_id, build_id)
+          .execute(consumption)
       end
     end
   end

ee/spec/services/ci/minutes/update_project_and_namespace_usage_service_spec.rb

@@ -5,32 +5,50 @@ require 'spec_helper'
 RSpec.describe Ci::Minutes::UpdateProjectAndNamespaceUsageService do
   let(:project) { create(:project, :private) }
   let(:namespace) { project.namespace }
+  let(:build) { create(:ci_build) }
   let(:consumption_minutes) { 120 }
   let(:consumption_seconds) { consumption_minutes * 60 }
   let(:namespace_amount_used) { Ci::Minutes::NamespaceMonthlyUsage.find_or_create_current(namespace_id: namespace.id).amount_used }
   let(:project_amount_used) { Ci::Minutes::ProjectMonthlyUsage.find_or_create_current(project_id: project.id).amount_used }
+  let(:service) { described_class.new(project.id, namespace.id, build.id) }
 
-  describe '#execute' do
-    subject { described_class.new(project.id, namespace.id) }
+  describe '#execute', :clean_gitlab_redis_shared_state do
+    subject { service.execute(consumption_minutes) }
 
-    context 'with shared runner' do
-      context 'when statistics and usage do not have existing values' do
+    shared_examples 'updates legacy consumption' do
       it 'updates legacy statistics with consumption seconds' do
-          subject.execute(consumption_minutes)
+        expect { subject }
+          .to change { project.statistics.reload.shared_runners_seconds }.by(consumption_seconds)
+          .and change { namespace.reload.namespace_statistics&.shared_runners_seconds || 0 }.by(consumption_seconds)
+      end
+    end
 
-          expect(project.statistics.reload.shared_runners_seconds)
-            .to eq(consumption_seconds)
+    shared_examples 'updates monthly consumption' do
+      it 'updates monthly usage with consumption minutes' do
+        expect { subject }
+          .to change { Ci::Minutes::NamespaceMonthlyUsage.find_or_create_current(namespace_id: namespace.id).amount_used }.by(consumption_minutes)
+          .and change { Ci::Minutes::ProjectMonthlyUsage.find_or_create_current(project_id: project.id).amount_used }.by(consumption_minutes)
+      end
+    end
 
-          expect(namespace.namespace_statistics.reload.shared_runners_seconds)
-            .to eq(consumption_seconds)
+    shared_examples 'does not update monthly consumption' do
+      it 'does not update the usage on a monthly basis' do
+        expect { subject }
+          .to change { Ci::Minutes::NamespaceMonthlyUsage.find_or_create_current(namespace_id: namespace.id).amount_used }.by(0)
+          .and change { Ci::Minutes::ProjectMonthlyUsage.find_or_create_current(project_id: project.id).amount_used }.by(0)
+      end
     end
 
+    context 'when statistics and usage do not have existing values' do
+      it_behaves_like 'updates legacy consumption'
+      it_behaves_like 'updates monthly consumption'
+
       context 'when project deleted' do
         let(:project) { double(id: non_existing_record_id) }
         let(:namespace) { create(:namespace) }
 
         it 'will complete successfully and increment namespace statistics' do
-            subject.execute(consumption_minutes)
+          subject
 
           expect(ProjectStatistics.find_by_project_id(project.id)).to be_nil
           expect(NamespaceStatistics.find_by_namespace_id(namespace.id).shared_runners_seconds).to eq(consumption_seconds)
@@ -43,7 +61,7 @@ RSpec.describe Ci::Minutes::UpdateProjectAndNamespaceUsageService do
         let(:namespace) { double(id: non_existing_record_id) }
 
         it 'will complete successfully' do
-            subject.execute(consumption_minutes)
+          subject
 
           expect(ProjectStatistics.find_by_project_id(project.id).shared_runners_seconds).to eq(consumption_seconds)
           expect(NamespaceStatistics.find_by_namespace_id(namespace.id)).to be_nil
@@ -57,7 +75,7 @@ RSpec.describe Ci::Minutes::UpdateProjectAndNamespaceUsageService do
         let(:namespace) { double(id: non_existing_record_id) }
 
         it 'will complete successfully' do
-            subject.execute(consumption_minutes)
+          subject
 
           expect(ProjectStatistics.find_by_project_id(project.id)).to be_nil
           expect(NamespaceStatistics.find_by_namespace_id(namespace.id)).to be_nil
@@ -66,13 +84,6 @@ RSpec.describe Ci::Minutes::UpdateProjectAndNamespaceUsageService do
         end
       end
 
-        it 'updates monthly usage with consumption minutes' do
-          subject.execute(consumption_minutes)
-
-          expect(namespace_amount_used).to eq(consumption_minutes)
-          expect(project_amount_used).to eq(consumption_minutes)
-        end
-
       context 'when on .com' do
         before do
           allow(Gitlab).to receive(:com?).and_return(true)
@@ -83,7 +94,7 @@ RSpec.describe Ci::Minutes::UpdateProjectAndNamespaceUsageService do
             expect(service).to receive(:execute)
           end
 
-            subject.execute(consumption_minutes)
+          subject
         end
       end
 
@@ -95,7 +106,7 @@ RSpec.describe Ci::Minutes::UpdateProjectAndNamespaceUsageService do
         it 'does not send a minute notification email' do
           expect(Ci::Minutes::EmailNotificationService).not_to receive(:new)
 
-            subject.execute(consumption_minutes)
+          subject
         end
       end
     end
@@ -119,7 +130,7 @@ RSpec.describe Ci::Minutes::UpdateProjectAndNamespaceUsageService do
         service = described_class.new(project.id, namespace.id)
 
         queries = ActiveRecord::QueryRecorder.new do
-            service.execute(consumption_minutes)
+          service
         end
 
         savepoints = queries.occurrences.keys.flatten.select do |query|
@@ -129,21 +140,41 @@ RSpec.describe Ci::Minutes::UpdateProjectAndNamespaceUsageService do
         expect(savepoints).to be_empty
       end
 
-        it 'updates legacy statistics with consumption seconds' do
-          subject.execute(consumption_minutes)
+      context 'behaves idempotently' do
+        let(:cache_key) { service.idempotency_cache_key }
 
-          expect(project.statistics.reload.shared_runners_seconds)
-            .to eq(existing_usage_in_seconds + consumption_seconds)
+        context 'when update has not been performed yet' do
+          it_behaves_like 'updates legacy consumption'
+          it_behaves_like 'updates monthly consumption'
 
-          expect(namespace.namespace_statistics.reload.shared_runners_seconds)
-            .to eq(existing_usage_in_seconds + consumption_seconds)
+          it 'tracks that the update is done' do
+            Gitlab::Redis::SharedState.with do |redis|
+              expect(redis.get(cache_key)).not_to be_present
             end
 
-        it 'updates monthly usage with consumption minutes' do
-          subject.execute(consumption_minutes)
+            subject
 
-          expect(namespace_amount_used).to eq(existing_usage_in_minutes + consumption_minutes)
-          expect(project_amount_used).to eq(existing_usage_in_minutes + consumption_minutes)
+            Gitlab::Redis::SharedState.with do |redis|
+              expect(redis.get(cache_key)).to be_present
+            end
+          end
+        end
+
+        context 'when update has previously been performed' do
+          before do
+            Gitlab::Redis::SharedState.with do |redis|
+              redis.set(cache_key, 1)
+            end
+          end
+
+          it_behaves_like 'does not update monthly consumption'
+          it_behaves_like 'updates legacy consumption' # not idempotent / to be removed
+
+          context 'when build_id is not provided as parameter' do
+            let(:service) { described_class.new(project.id, namespace.id) }
+
+            it_behaves_like 'updates monthly consumption' # not idempotent / temporary backwards compatibility
+          end
         end
       end
     end

ee/spec/workers/ci/minutes/update_project_and_namespace_usage_worker_spec.rb

@@ -5,39 +5,81 @@ require 'spec_helper'
 RSpec.describe Ci::Minutes::UpdateProjectAndNamespaceUsageWorker do
   let_it_be(:project) { create(:project) }
   let_it_be(:namespace) { project.namespace }
+  let_it_be(:build) { create(:ci_build, project: project) }
 
   let(:consumption) { 100 }
   let(:consumption_seconds) { consumption * 60 }
   let(:worker) { described_class.new }
 
   describe '#perform' do
+    shared_examples 'executes the update' do
       it 'executes UpdateProjectAndNamespaceUsageService' do
         service_instance = double
-      expect(::Ci::Minutes::UpdateProjectAndNamespaceUsageService).to receive(:new).with(project.id, namespace.id).and_return(service_instance)
-      expect(service_instance).to receive(:execute).with(consumption)
+        expect(::Ci::Minutes::UpdateProjectAndNamespaceUsageService).to receive(:new).at_least(:once).and_return(service_instance)
+        expect(service_instance).to receive(:execute).at_least(:once).with(consumption)
 
-      worker.perform(consumption, project.id, namespace.id)
+        subject
       end
 
-    it 'updates statistics and usage' do
-      worker.perform(consumption, project.id, namespace.id)
+      it 'updates monthly usage' do
+        subject
 
-      expect(project.statistics.reload.shared_runners_seconds).to eq(consumption_seconds)
-      expect(namespace.namespace_statistics.reload.shared_runners_seconds).to eq(consumption_seconds)
         expect(Ci::Minutes::NamespaceMonthlyUsage.find_by(namespace: namespace).amount_used).to eq(consumption)
         expect(Ci::Minutes::ProjectMonthlyUsage.find_by(project: project).amount_used).to eq(consumption)
       end
+    end
+
+    shared_examples 'skips the update' do
+      it 'does not execute UpdateProjectAndNamespaceUsageService' do
+        expect(::Ci::Minutes::UpdateProjectAndNamespaceUsageService).not_to receive(:new)
+
+        subject
+      end
+    end
 
-    it 'accumulates only legacy statistics on failure (behaves transactionally)' do
-      allow(Ci::Minutes::ProjectMonthlyUsage).to receive(:new).and_raise(StandardError)
+    context 'when build_id is not passed as parameter (param backward compatibility)' do
+      subject { worker.perform(consumption, project.id, namespace.id) }
 
-      expect { worker.perform(consumption, project.id, namespace.id) }.to raise_error(StandardError)
+      it_behaves_like 'executes the update'
 
-      expect(project.reload.statistics.shared_runners_seconds).to eq(consumption_seconds)
+      it 'updates legacy statistics' do
+        subject
+
+        expect(project.statistics.reload.shared_runners_seconds).to eq(consumption_seconds)
         expect(namespace.reload.namespace_statistics.shared_runners_seconds).to eq(consumption_seconds)
-      expect(Ci::Minutes::NamespaceMonthlyUsage.find_by(namespace: namespace)).to eq(nil)
-      expect(Ci::Minutes::ProjectMonthlyUsage.find_by(project: project)).to eq(nil)
-      expect(::Ci::Minutes::EmailNotificationService).not_to receive(:new)
+      end
+
+      context 'does not behave idempotently' do
+        subject { perform_multiple([consumption, project.id, namespace.id], worker: worker) }
+
+        it 'executes the operation multiple times' do
+          expect(::Ci::Minutes::UpdateProjectAndNamespaceUsageService).to receive(:new).twice.and_call_original
+
+          subject
+
+          expect(project.statistics.reload.shared_runners_seconds).to eq(2 * consumption_seconds)
+          expect(namespace.reload.namespace_statistics.shared_runners_seconds).to eq(2 * consumption_seconds)
+          expect(Ci::Minutes::NamespaceMonthlyUsage.find_by(namespace: namespace).amount_used).to eq(2 * consumption)
+          expect(Ci::Minutes::ProjectMonthlyUsage.find_by(project: project).amount_used).to eq(2 * consumption)
+        end
+      end
+    end
+
+    context 'when build_id is passed as parameter', :clean_gitlab_redis_shared_state do
+      subject { perform_multiple([consumption, project.id, namespace.id, build.id]) }
+
+      context 'behaves idempotently for monthly usage update' do
+        it_behaves_like 'executes the update'
+      end
+
+      it 'does not behave idempotently for legacy statistics update' do
+        expect(::Ci::Minutes::UpdateProjectAndNamespaceUsageService).to receive(:new).twice.and_call_original
+
+        subject
+
+        expect(project.statistics.reload.shared_runners_seconds).to eq(2 * consumption_seconds)
+        expect(namespace.reload.namespace_statistics.shared_runners_seconds).to eq(2 * consumption_seconds)
+      end
     end
   end
 end