Merge branch 'store-seats_currently_in_use-3' into 'master'

Persist seats in use and owed during cron update

See merge request gitlab-org/gitlab!40831

ee/app/models/gitlab_subscription.rb

@@ -24,6 +24,8 @@ class GitlabSubscription < ApplicationRecord
     with_hosted_plan(Plan::PAID_HOSTED_PLANS)
   end
 
+  scope :preload_for_refresh_seat, -> { preload([{ namespace: :route }, :hosted_plan]) }
+
   DAYS_AFTER_EXPIRATION_BEFORE_REMOVING_FROM_INDEX = 7
 
   # We set a 7 days as the threshold for expiration before removing them from
@@ -56,6 +58,13 @@ class GitlabSubscription < ApplicationRecord
     [0, max_seats_used - seats].max
   end
 
+  # Refresh seat related attribute (without persisting them)
+  def refresh_seat_attributes!
+    self.seats_in_use = calculate_seats_in_use
+    self.max_seats_used = [max_seats_used, seats_in_use].max
+    self.seats_owed = calculate_seats_owed
+  end
+
   def has_a_paid_hosted_plan?(include_trials: false)
     (include_trials || !trial?) &&
       hosted? &&

ee/app/workers/update_max_seats_used_for_gitlab_com_subscriptions_worker.rb

@@ -12,21 +12,22 @@ class UpdateMaxSeatsUsedForGitlabComSubscriptionsWorker # rubocop:disable Scalab
     return if ::Gitlab::Database.read_only?
     return unless ::Gitlab::CurrentSettings.should_check_namespace_plan?
 
-    GitlabSubscription.with_a_paid_hosted_plan.find_in_batches(batch_size: 100) do |subscriptions|
+    GitlabSubscription.with_a_paid_hosted_plan.preload_for_refresh_seat.find_in_batches(batch_size: 100) do |subscriptions|
       tuples = []
 
       subscriptions.each do |subscription|
-        seats_in_use = subscription.calculate_seats_in_use
+        subscription.refresh_seat_attributes!
 
-        next if subscription.max_seats_used >= seats_in_use
-
-        tuples << [subscription.id, seats_in_use]
+        tuples << [subscription.id, subscription.max_seats_used, subscription.seats_in_use, subscription.seats_owed]
       end
 
       if tuples.present?
         GitlabSubscription.connection.execute <<-EOF
-          UPDATE gitlab_subscriptions AS s SET max_seats_used = v.max_seats_used
-          FROM (VALUES #{tuples.map { |tuple| "(#{tuple.join(', ')})" }.join(', ')}) AS v(id, max_seats_used)
+          UPDATE gitlab_subscriptions AS s
+          SET max_seats_used = v.max_seats_used,
+              seats_in_use = v.seats_in_use,
+              seats_owed = v.seats_owed
+          FROM (VALUES #{tuples.map { |tuple| "(#{tuple.join(', ')})" }.join(', ')}) AS v(id, max_seats_used, seats_in_use, seats_owed)
           WHERE s.id = v.id
         EOF
       end

ee/spec/models/gitlab_subscription_spec.rb

@@ -176,6 +176,38 @@ RSpec.describe GitlabSubscription do
     end
   end
 
+  describe '#refresh_seat_attributes!' do
+    subject { create(:gitlab_subscription, seats: 3, max_seats_used: 2) }
+
+    before do
+      expect(subject).to receive(:calculate_seats_in_use).and_return(calculate_seats_in_use)
+    end
+
+    context 'when current seats in use is lower than recorded max_seats_used' do
+      let(:calculate_seats_in_use) { 1 }
+
+      it 'does not increase max_seats_used' do
+        expect do
+          subject.refresh_seat_attributes!
+        end.to change(subject, :seats_in_use).from(0).to(1)
+          .and not_change(subject, :max_seats_used)
+          .and not_change(subject, :seats_owed)
+      end
+    end
+
+    context 'when current seats in use is higher than seats and max_seats_used' do
+      let(:calculate_seats_in_use) { 4 }
+
+      it 'increases seats and max_seats_used' do
+        expect do
+          subject.refresh_seat_attributes!
+        end.to change(subject, :seats_in_use).from(0).to(4)
+          .and change(subject, :max_seats_used).from(2).to(4)
+          .and change(subject, :seats_owed).from(0).to(1)
+      end
+    end
+  end
+
   describe '#expired?' do
     let(:gitlab_subscription) { create(:gitlab_subscription, end_date: end_date) }
 

ee/spec/workers/update_max_seats_used_for_gitlab_com_subscriptions_worker_spec.rb

@@ -5,11 +5,10 @@ require 'spec_helper'
 RSpec.describe UpdateMaxSeatsUsedForGitlabComSubscriptionsWorker do
   subject { described_class.new }
 
-  let!(:user)                { create(:user) }
-  let!(:group)               { create(:group) }
-  let!(:bronze_plan)         { create(:bronze_plan) }
-  let!(:early_adopter_plan)  { create(:early_adopter_plan) }
-  let!(:gitlab_subscription) { create(:gitlab_subscription, namespace: group) }
+  let_it_be(:bronze_plan) { create(:bronze_plan) }
+  let_it_be(:early_adopter_plan) { create(:early_adopter_plan) }
+  let_it_be(:gitlab_subscription, refind: true) { create(:gitlab_subscription, seats: 1) }
+  let_it_be(:gitlab_subscription_2, refind: true) { create(:gitlab_subscription, seats: 11) }
 
   let(:db_is_read_only) { false }
   let(:subscription_attrs) { nil }
@@ -18,22 +17,54 @@ RSpec.describe UpdateMaxSeatsUsedForGitlabComSubscriptionsWorker do
     allow(Gitlab::Database).to receive(:read_only?) { db_is_read_only }
     allow(Gitlab::CurrentSettings).to receive(:should_check_namespace_plan?) { true }
 
-    group.add_developer(user)
+    allow_next_found_instance_of(GitlabSubscription) do |subscription|
+      allow(subscription).to receive(:refresh_seat_attributes!) do
+        subscription.max_seats_used = subscription.seats + 3
+        subscription.seats_in_use = subscription.seats + 2
+        subscription.seats_owed = subscription.seats + 1
+      end
+    end
+  end
+
+  def perform_and_reload
+    subject.perform
+    gitlab_subscription.reload
+    gitlab_subscription_2.reload
+  end
+
+  shared_examples 'updates nothing' do
+    it 'does not update seat columns' do
+      expect do
+        perform_and_reload
+      end.to not_change(gitlab_subscription, :max_seats_used)
+        .and not_change(gitlab_subscription, :seats_in_use)
+        .and not_change(gitlab_subscription, :seats_owed)
+        .and not_change(gitlab_subscription_2, :max_seats_used)
+        .and not_change(gitlab_subscription_2, :seats_in_use)
+        .and not_change(gitlab_subscription_2, :seats_owed)
+    end
   end
 
-  shared_examples 'keeps original max_seats_used value' do
-    it 'does not update max_seats_used' do
-      expect { subject.perform }.not_to change { gitlab_subscription.reload.max_seats_used }
+  shared_examples 'updates only paid plans' do
+    it "persists seat attributes after refresh_seat_attributes! for only paid plans" do
+      expect do
+        perform_and_reload
+      end.to not_change(gitlab_subscription, :max_seats_used)
+        .and not_change(gitlab_subscription, :seats_in_use)
+        .and not_change(gitlab_subscription, :seats_owed)
+        .and change(gitlab_subscription_2, :max_seats_used).from(0).to(14)
+        .and change(gitlab_subscription_2, :seats_in_use).from(0).to(13)
+        .and change(gitlab_subscription_2, :seats_owed).from(0).to(12)
     end
   end
 
   context 'where the DB is read only' do
     let(:db_is_read_only) { true }
 
-    include_examples 'keeps original max_seats_used value'
+    include_examples 'updates nothing'
   end
 
-  context 'when the DB PostgreSQK AND is not read only' do
+  context 'when the DB is not read only' do
     before do
       gitlab_subscription.update!(subscription_attrs) if subscription_attrs
     end
@@ -41,34 +72,36 @@ RSpec.describe UpdateMaxSeatsUsedForGitlabComSubscriptionsWorker do
     context 'with a free plan' do
       let(:subscription_attrs) { { hosted_plan: nil } }
 
-      include_examples 'keeps original max_seats_used value'
+      include_examples 'updates only paid plans'
     end
 
     context 'with a trial plan' do
       let(:subscription_attrs) { { hosted_plan: bronze_plan, trial: true } }
 
-      include_examples 'keeps original max_seats_used value'
+      include_examples 'updates only paid plans'
     end
 
     context 'with an early adopter plan' do
       let(:subscription_attrs) { { hosted_plan: early_adopter_plan } }
 
-      include_examples 'keeps original max_seats_used value'
+      include_examples 'updates only paid plans'
     end
 
-    context 'with a paid plan' do
+    context 'with a paid plan', :aggregate_failures do
       before do
         gitlab_subscription.update!(hosted_plan: bronze_plan)
+        gitlab_subscription_2.update!(hosted_plan: bronze_plan)
       end
 
-      it 'only updates max_seats_used if active users count is greater than it' do
-        expect { subject.perform }.to change { gitlab_subscription.reload.max_seats_used }.to(1)
-      end
-
-      it 'does not update max_seats_used if active users count is lower than it' do
-        gitlab_subscription.update_attribute(:max_seats_used, 5)
-
-        expect { subject.perform }.not_to change { gitlab_subscription.reload.max_seats_used }
+      it 'persists seat attributes after refresh_seat_attributes' do
+        expect do
+          perform_and_reload
+        end.to change(gitlab_subscription, :max_seats_used).from(0).to(4)
+          .and change(gitlab_subscription, :seats_in_use).from(0).to(3)
+          .and change(gitlab_subscription, :seats_owed).from(0).to(2)
+          .and change(gitlab_subscription_2, :max_seats_used).from(0).to(14)
+          .and change(gitlab_subscription_2, :seats_in_use).from(0).to(13)
+          .and change(gitlab_subscription_2, :seats_owed).from(0).to(12)
       end
     end
   end