Refactor refresh user assignments jobs

* Call reconcile_seat_overage from refresh_user_assignments_worker
* Call reconcile_seat_overage from bulk_refresh_user_assignments_worker
* Call refresh_user_assignments_worker from update_service

Changelog: added
EE: true

ee/app/services/gitlab_subscriptions/add_on_purchases/update_service.rb

@@ -15,7 +15,13 @@ def initialize(namespace, add_on, params = {})
       def execute
         return error_response unless add_on_purchase
 
-        update_add_on_purchase ? successful_response : error_response
+        if update_add_on_purchase
+          perform_after_update_actions
+
+          successful_response
+        else
+          error_response
+        end
       end
 
       private
@@ -41,6 +47,12 @@ def update_add_on_purchase
         add_on_purchase.update(attributes)
       end
 
+      def perform_after_update_actions
+        GitlabSubscriptions::AddOnPurchases::RefreshUserAssignmentsWorker.perform_async(
+          add_on_purchase.namespace_id
+        )
+      end
+
       override :error_response
       def error_response
         if add_on_purchase.nil?

ee/app/workers/gitlab_subscriptions/add_on_purchases/bulk_refresh_user_assignments_worker.rb

@@ -19,6 +19,12 @@ def perform_work
 
         deleted_assignments_count = add_on_purchase.delete_ineligible_user_assignments_in_batches!
 
+        reconcile_response = GitlabSubscriptions::AddOnPurchases::ReconcileSeatOverageService.new(
+          add_on_purchase: add_on_purchase
+        ).execute
+
+        deleted_assignments_count += reconcile_response.payload[:removed_seats_count]
+
         log_event(deleted_assignments_count) if deleted_assignments_count > 0
       end
 

ee/app/workers/gitlab_subscriptions/add_on_purchases/refresh_user_assignments_worker.rb

@@ -6,7 +6,7 @@ class RefreshUserAssignmentsWorker
       include ::ApplicationWorker
       include Gitlab::ExclusiveLeaseHelpers
 
-      LEASE_TTL = 1.minute
+      LEASE_TTL = 3.minutes
 
       feature_category :seat_cost_management
 
@@ -19,10 +19,16 @@ class RefreshUserAssignmentsWorker
       def perform(root_namespace_id)
         @root_namespace_id = root_namespace_id
 
-        return unless root_namespace && add_on_purchase
+        return unless add_on_purchase
 
-        deleted_assignments_count = in_lock(lock_key, ttl: LEASE_TTL) do
-          add_on_purchase.delete_ineligible_user_assignments_in_batches!
+        deleted_assignments_count = 0
+        in_lock(lock_key, ttl: LEASE_TTL, retries: 0) do
+          deleted_assignments_count += add_on_purchase.delete_ineligible_user_assignments_in_batches!
+
+          reconcile_response = GitlabSubscriptions::AddOnPurchases::ReconcileSeatOverageService.new(
+            add_on_purchase: add_on_purchase.reset
+          ).execute
+          deleted_assignments_count += reconcile_response.payload[:removed_seats_count]
         end
 
         # #update_column used to skip validations and callbacks.
@@ -38,16 +44,15 @@ def perform(root_namespace_id)
 
       attr_reader :root_namespace_id
 
-      def root_namespace
-        @root_namespace ||= Group.find_by_id(root_namespace_id)
-      end
-
       def add_on_purchase
-        @add_on_purchase ||= root_namespace.subscription_add_on_purchases.for_gitlab_duo_pro.first
+        @add_on_purchase ||= ::GitlabSubscriptions::AddOnPurchase
+          .for_gitlab_duo_pro
+          .by_namespace_id(root_namespace_id)
+          .first
       end
 
       def lock_key
-        "#{self.class.name.underscore}:#{root_namespace.id}"
+        "#{self.class.name.underscore}:#{add_on_purchase.id}"
       end
 
       def log_event(deleted_count)
@@ -55,7 +60,8 @@ def log_event(deleted_count)
           message: 'AddOnPurchase user assignments refreshed in bulk',
           deleted_assignments_count: deleted_count,
           add_on: add_on_purchase.add_on.name,
-          namespace: root_namespace.path
+          add_on_purchase_id: add_on_purchase.id,
+          namespace_id: root_namespace_id
         )
       end
     end

ee/spec/services/gitlab_subscriptions/add_on_purchases/update_service_spec.rb

@@ -49,6 +49,14 @@
             .and change { add_on_purchase.trial }.from(false).to(true)
         end
 
+        it 'enqueues RefreshUserAssignmentsWorker' do
+          expect(
+            GitlabSubscriptions::AddOnPurchases::RefreshUserAssignmentsWorker
+          ).to receive(:perform_async).with(add_on_purchase.namespace_id)
+
+          result
+        end
+
         context 'when passing in the add-on purchase record' do
           let(:params) do
             super().merge(add_on_purchase: add_on_purchase)
@@ -73,6 +81,14 @@
             expect(result[:add_on_purchase]).to be_an_instance_of(GitlabSubscriptions::AddOnPurchase)
             expect(result[:add_on_purchase]).to eq(add_on_purchase)
           end
+
+          it 'does not enqueue RefreshUserAssignmentsWorker' do
+            expect(
+              GitlabSubscriptions::AddOnPurchases::RefreshUserAssignmentsWorker
+            ).not_to receive(:perform_async)
+
+            result
+          end
         end
       end
     end

ee/spec/workers/gitlab_subscriptions/add_on_purchases/bulk_refresh_user_assignments_worker_spec.rb

@@ -44,6 +44,34 @@
 
             expect(add_on_purchase_fresh.assigned_users.count).to eq(1)
           end
+
+          context 'when there is a seat overage' do
+            let_it_be(:user_1) { create(:user) }
+            let_it_be(:user_2) { create(:user) }
+            let_it_be(:user_3) { create(:user) }
+
+            before_all do
+              GitlabSubscriptions::UserAddOnAssignment.delete_all
+
+              # user_3 is not added as a group member; thus is ineligible
+              add_on_purchase_stale.namespace.add_developer(user_1)
+              add_on_purchase_stale.namespace.add_developer(user_2)
+
+              # user_2 is assigned last; thus will be prioritzed for overage cleanup
+              add_on_purchase_stale.assigned_users.create!(user: user_3)
+              add_on_purchase_stale.assigned_users.create!(user: user_1)
+              add_on_purchase_stale.assigned_users.create!(user: user_2)
+            end
+
+            it 'reconciles any seat overage' do
+              expect do
+                perform_work
+              end.to change { add_on_purchase_stale.reload.assigned_users.count }.by(-2)
+                .and change { add_on_purchase_stale.reload.last_assigned_users_refreshed_at }
+
+              expect(add_on_purchase_stale.assigned_users.map(&:user)).to eq([user_1])
+            end
+          end
         end
       end
 

ee/spec/workers/gitlab_subscriptions/add_on_purchases/refresh_user_assignments_worker_spec.rb

@@ -6,7 +6,10 @@
   describe '#perform' do
     let_it_be(:namespace) { create(:group) }
     let_it_be(:add_on) { create(:gitlab_subscription_add_on) }
-    let_it_be(:add_on_purchase) { create(:gitlab_subscription_add_on_purchase, namespace: namespace, add_on: add_on) }
+    let_it_be(:add_on_purchase) do
+      create(:gitlab_subscription_add_on_purchase, namespace: namespace, quantity: 2, add_on: add_on)
+    end
+
     let_it_be(:other_add_on_purchase) { create(:gitlab_subscription_add_on_purchase, add_on: add_on) }
 
     let(:root_namespace_id) { namespace.id }
@@ -39,12 +42,6 @@
       end
     end
 
-    context 'when root_namespace_id does not exists' do
-      let(:root_namespace_id) { nil }
-
-      it_behaves_like 'does not remove seat assignment'
-    end
-
     context 'when root namespace does not have related purchase' do
       let(:root_namespace_id) { create(:group).id }
 
@@ -78,6 +75,34 @@
             expect(add_on_purchase.assigned_users.by_user(user_1).count).to eq(1)
           end
         end
+
+        context 'when there is a seat overage' do
+          let_it_be(:user_3) { create(:user) }
+
+          before_all do
+            GitlabSubscriptions::UserAddOnAssignment.delete_all
+
+            # user_2 is not added as a group member; thus is inelligble
+            add_on_purchase.reload.namespace.add_guest(user_1)
+            add_on_purchase.namespace.add_developer(user_3)
+
+            # user_3 is assigned last; thus will be prioritzed for overage cleanup
+            add_on_purchase.assigned_users.create!(user: user_1)
+            add_on_purchase.assigned_users.create!(user: user_2)
+            add_on_purchase.assigned_users.create!(user: user_3)
+
+            # decrease the quantity to create overage
+            add_on_purchase.update!(quantity: 1)
+          end
+
+          it 'removes ineligible users and reconciles any seat overage' do
+            expect do
+              subject
+            end.to change { GitlabSubscriptions::UserAddOnAssignment.count }.by(-2)
+
+            expect(add_on_purchase.assigned_users.map(&:user)).to eq([user_1])
+          end
+        end
       end
     end
 
@@ -88,12 +113,42 @@
         message: 'AddOnPurchase user assignments refreshed in bulk',
         deleted_assignments_count: 2,
         add_on: add_on_purchase.add_on.name,
-        namespace: namespace.path
+        add_on_purchase_id: add_on_purchase.id,
+        namespace_id: namespace.id
       )
 
       subject.perform(root_namespace_id)
     end
 
+    context 'when root_namespace_id is nil' do
+      let(:root_namespace_id) { nil }
+
+      context 'when there is no associated add_on_purchase' do
+        it_behaves_like 'does not remove seat assignment'
+      end
+
+      context 'when there is an associated add_on_purchase' do
+        let(:add_on_purchase) do
+          create(:gitlab_subscription_add_on_purchase, :self_managed, add_on: add_on)
+        end
+
+        before do
+          add_on_purchase.assigned_users.create!(user: user_1)
+          add_on_purchase.assigned_users.create!(user: user_2)
+        end
+
+        it 'refreshes users seat assignments' do
+          expect do
+            subject.perform(root_namespace_id)
+          end.to change { add_on_purchase.reload.assigned_users.count }.from(2).to(1)
+
+          expect(add_on_purchase.assigned_users.map(&:user)).to eq([user_1])
+        end
+
+        it_behaves_like 'updates last_assigned_users_refreshed_at attribute'
+      end
+    end
+
     context 'when no assignments were deleted' do
       before_all do
         namespace.add_guest(user_1)
@@ -112,7 +167,7 @@
     context 'with exclusive lease' do
       include ExclusiveLeaseHelpers
 
-      let(:lock_key) { "#{described_class.name.underscore}:#{root_namespace_id}" }
+      let(:lock_key) { "#{described_class.name.underscore}:#{add_on_purchase.id}" }
       let(:timeout) { described_class::LEASE_TTL }
 
       context 'when exclusive lease has not been taken' do