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)
-
-          send_minutes_email_notification
         end
-      end
 
-      private
+        send_minutes_email_notification
+      end
 
-      def preload_minutes_usage_data!
-        project_usage
-        namespace_usage
+      def idempotency_cache_key
+        "ci_minutes_usage:#{@project_id}:#{@build_id}:updated"
       end
 
+      private
+
       def send_minutes_email_notification
         # `perform reset` on `project` because `Namespace#namespace_statistics` will otherwise return stale data.
         # TODO(issue 335885): Remove @project
@@ -46,8 +51,14 @@ module Ci
       end
 
       def track_usage_of_monthly_minutes(consumption)
-        ::Ci::Minutes::NamespaceMonthlyUsage.increase_usage(namespace_usage, consumption) if namespace_usage
-        ::Ci::Minutes::ProjectMonthlyUsage.increase_usage(project_usage, consumption) if project_usage
+        # 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)
@@ -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/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
-    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)
+    shared_examples 'executes the update' do
+      it 'executes UpdateProjectAndNamespaceUsageService' do
+        service_instance = double
+        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 monthly usage' do
+        subject
+
+        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
 
-    it 'updates statistics and usage' do
-      worker.perform(consumption, project.id, namespace.id)
+    shared_examples 'skips the update' do
+      it 'does not execute UpdateProjectAndNamespaceUsageService' do
+        expect(::Ci::Minutes::UpdateProjectAndNamespaceUsageService).not_to receive(:new)
 
-      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)
+        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) }
+
+      it_behaves_like 'executes the update'
+
+      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)
+      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
 
-      expect { worker.perform(consumption, project.id, namespace.id) }.to raise_error(StandardError)
+        subject
 
-      expect(project.reload.statistics.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)
+        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