Merge branch '415584-add-on-assignment-with-locking' into 'master'

Add resource locking

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/127677



Merged-by: Stan Hu <stanhu@gmail.com>
Approved-by: Stan Hu <stanhu@gmail.com>
Reviewed-by: Abdul Wadood <awadood@gitlab.com>
Reviewed-by: Bishwa Hang Rai <bhrai@gitlab.com>
Co-authored-by: Bishwa Hang Rai <bhrai@gitlab.com>
Co-authored-by: Abdul Wadood <awadood@gitlab.com>

ee/app/services/gitlab_subscriptions/user_add_on_assignments/create_service.rb

@@ -3,9 +3,17 @@
 module GitlabSubscriptions
   module UserAddOnAssignments
     class CreateService < BaseService
+      include Gitlab::Utils::StrongMemoize
+
       ERROR_NO_SEATS_AVAILABLE = 'NO_SEATS_AVAILABLE'
       ERROR_INVALID_USER_MEMBERSHIP = 'INVALID_USER_MEMBERSHIP'
 
+      NoSeatsAvailableError = Class.new(StandardError) do
+        def initialize(message = ERROR_NO_SEATS_AVAILABLE)
+          super(message)
+        end
+      end
+
       def initialize(add_on_purchase:, user:)
         @add_on_purchase = add_on_purchase
         @user = user
@@ -17,14 +25,18 @@ def execute
         errors = validate
 
         if errors.blank?
-          # TODO: implement resource locking to avoid race condition
-          # https://gitlab.com/gitlab-org/gitlab/-/issues/415584#race-condition
-          add_on_purchase.assigned_users.create!(user: user)
+          add_on_purchase.with_lock do
+            raise NoSeatsAvailableError unless seats_available?
+
+            add_on_purchase.assigned_users.create!(user: user)
+          end
 
           ServiceResponse.success
         else
           ServiceResponse.error(message: errors)
         end
+      rescue NoSeatsAvailableError => error
+        ServiceResponse.error(message: error.message)
       end
 
       private
@@ -41,7 +53,7 @@ def seats_available?
       end
 
       def assigned_seats
-        @assigned_seats ||= add_on_purchase.assigned_users.count
+        add_on_purchase.assigned_users.count
       end
 
       def user_already_assigned?
@@ -54,6 +66,7 @@ def billed_member_of_namespace?
           namespace.billed_shared_group_user?(user, exclude_guests: true) ||
           namespace.billed_shared_project_user?(user, exclude_guests: true)
       end
+      strong_memoize_attr :billed_member_of_namespace?
 
       def namespace
         @namespace ||= add_on_purchase.namespace

ee/spec/services/gitlab_subscriptions/user_add_on_assignments/create_service_spec.rb

@@ -100,5 +100,48 @@
 
       it_behaves_like 'success response'
     end
+
+    context 'with resource locking' do
+      before do
+        add_on_purchase.update!(quantity: 1)
+      end
+
+      it 'prevents from double booking assignment' do
+        users = create_list(:user, 3)
+
+        expect(add_on_purchase.assigned_users.count).to eq(0)
+
+        users.map do |user|
+          namespace.add_developer(user)
+
+          Thread.new do
+            described_class.new(
+              add_on_purchase: add_on_purchase,
+              user: user
+            ).execute
+          end
+        end.each(&:join)
+
+        expect(add_on_purchase.assigned_users.count).to eq(1)
+      end
+
+      context 'when NoSeatsAvailableError is raised' do
+        let(:service_instance) { described_class.new(add_on_purchase: add_on_purchase, user: user) }
+
+        subject(:response) { service_instance.execute }
+
+        it 'handes the error correctly' do
+          # fill up the available seats
+          create(:gitlab_subscription_user_add_on_assignment, add_on_purchase: add_on_purchase)
+
+          # Mock first call to return true to pass the validate phase
+          expect(service_instance).to receive(:seats_available?).and_return(true)
+          expect(service_instance).to receive(:seats_available?).and_call_original
+
+          expect { subject }.not_to raise_error
+          expect(response.errors).to include('NO_SEATS_AVAILABLE')
+        end
+      end
+    end
   end
 end