From bf71ed22cd77ee288d4e1cf617ec678888f67012 Mon Sep 17 00:00:00 2001
From: Doug Stull <dstull@gitlab.com>
Date: Wed, 7 Jun 2023 15:35:26 -0400
Subject: [PATCH 1/2] Fine tune over limit email notification worker

- remove locking mechanism as we no longer use it.

Changelog: other
EE: true
---
 ee/app/models/ee/namespace/detail.rb          |  3 +-
 .../over_limit_notification_worker.rb         | 35 +++++++++++--------
 ee/spec/models/ee/namespace/detail_spec.rb    |  4 +--
 .../over_limit_notification_worker_spec.rb    | 30 +++++++++++-----
 4 files changed, 45 insertions(+), 27 deletions(-)

diff --git a/ee/app/models/ee/namespace/detail.rb b/ee/app/models/ee/namespace/detail.rb
index c0e117f28a70b37f..a4ce7d57c50b450b 100644
--- a/ee/app/models/ee/namespace/detail.rb
+++ b/ee/app/models/ee/namespace/detail.rb
@@ -19,11 +19,10 @@ module Detail
           order(arel_table[:next_over_limit_check_at].asc.nulls_first)
         end
 
-        scope :lock_for_over_limit_check, ->(limit, namespace_ids) do
+        scope :for_over_limit_check, ->(limit, namespace_ids) do
           scheduled_for_over_limit_check
             .where(namespace_id: namespace_ids)
             .limit(limit)
-            .lock('FOR UPDATE SKIP LOCKED')
         end
       end
     end
diff --git a/ee/app/workers/namespaces/free_user_cap/over_limit_notification_worker.rb b/ee/app/workers/namespaces/free_user_cap/over_limit_notification_worker.rb
index 1b7b6f3695f754e5..5dd415f82a42ff8c 100644
--- a/ee/app/workers/namespaces/free_user_cap/over_limit_notification_worker.rb
+++ b/ee/app/workers/namespaces/free_user_cap/over_limit_notification_worker.rb
@@ -15,13 +15,18 @@ class OverLimitNotificationWorker
       BATCH_SIZE = 100
       SCHEDULE_BUFFER_IN_HOURS = 24
 
-      def perform_work(*args)
-        next_batch.map(&:namespace_id).each do |namespace_id|
-          notify namespace_id: namespace_id
+      def perform_work(*_args)
+        ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
+          batch_results = next_batch
+          next if batch_results.blank?
+
+          batch_results.map(&:namespace_id).each do |namespace_id|
+            notify namespace_id: namespace_id
+          end
         end
       end
 
-      def remaining_work_count(*args)
+      def remaining_work_count(*_args)
         (Namespace::Detail.scheduled_for_over_limit_check.limit(MAX_RUNNING_JOBS * BATCH_SIZE).count / BATCH_SIZE).ceil
       end
 
@@ -43,28 +48,28 @@ def notify(namespace_id:)
 
       def next_batch
         Namespace::Detail.transaction do
-          Namespace::Detail.find_by_sql next_batch_sql # rubocop: disable CodeReuse/ActiveRecord
+          ids = item_ids
+          Namespace::Detail.find_by_sql(perform_update_sql(ids)) if ids.present? # rubocop: disable CodeReuse/ActiveRecord
         end
       end
 
-      def next_batch_sql
-        # The alterntive to setting the timestamp would be using something like
+      def perform_update_sql(ids)
+        # The alternative to setting the timestamp would be using something like
         # SET "next_over_limit_check_at" = NOW() + INTERVAL '#{SCHEDULE_BUFFER_IN_HOURS} hours'
         # that makes it quite hared to spec, since the DB doesn't freeze time along with ruby
         <<~SQL.squish
-        UPDATE "namespace_details"
-        SET "next_over_limit_check_at" = to_timestamp(#{schedule.to_i})
-          WHERE "namespace_details"."namespace_id" IN (#{locked_item_ids_sql})
-          RETURNING *
+          UPDATE "namespace_details"
+          SET "next_over_limit_check_at" = to_timestamp(#{schedule.to_i})
+            WHERE "namespace_details"."namespace_id" IN (#{ids.join(', ')})
+            RETURNING *
         SQL
       end
 
-      def locked_item_ids_sql
+      def item_ids
         Namespace::Detail
           .not_over_limit_notified
-          .lock_for_over_limit_check(BATCH_SIZE, namespace_ids)
-          .select(:namespace_id)
-          .to_sql
+          .for_over_limit_check(BATCH_SIZE, namespace_ids)
+          .pluck(:namespace_id) # rubocop: disable CodeReuse/ActiveRecord
       end
 
       def namespace_ids
diff --git a/ee/spec/models/ee/namespace/detail_spec.rb b/ee/spec/models/ee/namespace/detail_spec.rb
index cec9acc5a04884e0..e7041722d6c0794c 100644
--- a/ee/spec/models/ee/namespace/detail_spec.rb
+++ b/ee/spec/models/ee/namespace/detail_spec.rb
@@ -55,11 +55,11 @@ def set_schedule(group, time, notified_at = nil)
       end
     end
 
-    context 'for lock_for_over_limit_check' do
+    context 'for for_over_limit_check' do
       let(:limit) { 2 }
       let(:scheduled_items) { [group_1.namespace_details, group_2.namespace_details] }
 
-      subject(:lock_scope_result) { described_class.lock_for_over_limit_check(limit, namespace_ids) }
+      subject(:lock_scope_result) { described_class.for_over_limit_check(limit, namespace_ids) }
 
       it 'only returns scheduled entries up to the limit' do
         expect(lock_scope_result).to eq(scheduled_items)
diff --git a/ee/spec/workers/namespaces/free_user_cap/over_limit_notification_worker_spec.rb b/ee/spec/workers/namespaces/free_user_cap/over_limit_notification_worker_spec.rb
index 9110d348f37071d3..89457b8008dfbc09 100644
--- a/ee/spec/workers/namespaces/free_user_cap/over_limit_notification_worker_spec.rb
+++ b/ee/spec/workers/namespaces/free_user_cap/over_limit_notification_worker_spec.rb
@@ -2,9 +2,7 @@
 
 require 'spec_helper'
 
-# Interim feature category experimentation_activation used here while waiting for
-# https://gitlab.com/gitlab-com/www-gitlab-com/-/merge_requests/113300 to merge
-RSpec.describe Namespaces::FreeUserCap::OverLimitNotificationWorker, :saas, feature_category: :experimentation_activation, type: :worker do
+RSpec.describe Namespaces::FreeUserCap::OverLimitNotificationWorker, :saas, feature_category: :measurement_and_locking, type: :worker do
   using RSpec::Parameterized::TableSyntax
 
   describe '#perform' do
@@ -26,14 +24,30 @@
       stub_ee_application_setting should_check_namespace_plan: true
     end
 
-    it 'runs notify service and marks next check for the namespace' do
-      stub_ee_application_setting dashboard_limit_enabled: true
+    context 'when dashboard_limit_enabled is true' do
+      before do
+        stub_ee_application_setting dashboard_limit_enabled: true
+      end
+
+      it 'runs notify service and marks next check for the namespace' do
+        expect(::Namespaces::FreeUserCap::NotifyOverLimitService).to receive(:execute).with(root_namespace: namespace)
+
+        next_check_time = frozen_time + described_class::SCHEDULE_BUFFER_IN_HOURS.hours
+
+        expect { worker }.to change { namespace.reload.namespace_details.next_over_limit_check_at }.to(next_check_time)
+      end
 
-      expect(::Namespaces::FreeUserCap::NotifyOverLimitService).to receive(:execute).with(root_namespace: namespace)
+      context 'when there are no results for next batch' do
+        before do
+          namespace.update!(visibility_level: Gitlab::VisibilityLevel::PUBLIC)
+        end
 
-      next_check_time = frozen_time + described_class::SCHEDULE_BUFFER_IN_HOURS.hours
+        it 'gracefully handles and does not make any updates' do
+          expect(::Namespaces::FreeUserCap::NotifyOverLimitService).not_to receive(:execute)
 
-      expect { worker }.to change { namespace.reload.namespace_details.next_over_limit_check_at }.to(next_check_time)
+          expect { worker }.not_to change { namespace.reload.namespace_details.next_over_limit_check_at }
+        end
+      end
     end
 
     context 'with feature flags enabled/disabled' do
-- 
GitLab


From c4a3c0178e08d7f47b6e04390ad6498ca4cffa61 Mon Sep 17 00:00:00 2001
From: Doug Stull <dstull@gitlab.com>
Date: Thu, 8 Jun 2023 07:27:00 -0400
Subject: [PATCH 2/2] Update per review suggestion

---
 .../free_user_cap/over_limit_notification_worker.rb          | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/ee/app/workers/namespaces/free_user_cap/over_limit_notification_worker.rb b/ee/app/workers/namespaces/free_user_cap/over_limit_notification_worker.rb
index 5dd415f82a42ff8c..21703ae581a8135f 100644
--- a/ee/app/workers/namespaces/free_user_cap/over_limit_notification_worker.rb
+++ b/ee/app/workers/namespaces/free_user_cap/over_limit_notification_worker.rb
@@ -17,10 +17,7 @@ class OverLimitNotificationWorker
 
       def perform_work(*_args)
         ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
-          batch_results = next_batch
-          next if batch_results.blank?
-
-          batch_results.map(&:namespace_id).each do |namespace_id|
+          Array.wrap(next_batch).map(&:namespace_id).each do |namespace_id|
             notify namespace_id: namespace_id
           end
         end
-- 
GitLab