From 44dc70051a5550d8032555469406b46321cee9cc Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Mon, 27 May 2024 22:46:18 +1200
Subject: [PATCH 1/3] Add QueryLimiting to Sidekiq workers as well

For now, we only add while in test environment, as development
environment Sidekiq is not visible
---
 .../query_limiting/sidekiq_middleware.rb      | 28 +++++++++++
 .../query_limiting/sidekiq_middleware_spec.rb | 48 +++++++++++++++++++
 spec/spec_helper.rb                           |  2 +
 spec/support/sidekiq_middleware.rb            |  2 +-
 4 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 lib/gitlab/query_limiting/sidekiq_middleware.rb
 create mode 100644 spec/lib/gitlab/query_limiting/sidekiq_middleware_spec.rb

diff --git a/lib/gitlab/query_limiting/sidekiq_middleware.rb b/lib/gitlab/query_limiting/sidekiq_middleware.rb
new file mode 100644
index 00000000000000..dc215f3894650d
--- /dev/null
+++ b/lib/gitlab/query_limiting/sidekiq_middleware.rb
@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module QueryLimiting
+    # Middleware for reporting (or raising) when a Sidekiq worker performs more than a
+    # certain amount of database queries.
+    class SidekiqMiddleware
+      def call(worker, _job, _queue)
+        #return yield if Sidekiq::Testing.inline?
+
+        transaction, retval = ::Gitlab::QueryLimiting::Transaction.run do
+          yield
+        end
+
+        transaction.action = action_name(worker)
+        transaction.act_upon_results
+
+        retval
+      end
+
+      private
+
+      def action_name(worker)
+        worker.class.name
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/query_limiting/sidekiq_middleware_spec.rb b/spec/lib/gitlab/query_limiting/sidekiq_middleware_spec.rb
new file mode 100644
index 00000000000000..3a1c8bdac94425
--- /dev/null
+++ b/spec/lib/gitlab/query_limiting/sidekiq_middleware_spec.rb
@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::QueryLimiting::SidekiqMiddleware, feature_category: :database do
+  describe '#call' do
+    let(:worker_class) do
+      Class.new do
+        def self.name
+          'TestWorker'
+        end
+
+        include ApplicationWorker
+      end
+    end
+
+    let(:worker) { worker_class.new }
+    let(:job) { {} }
+    let(:queue) { :test }
+
+    it 'runs the middleware with query limiting in place' do
+      expect_next_instance_of(Gitlab::QueryLimiting::Transaction) do |instance|
+        expect(instance).to receive(:action=).with('TestWorker')
+        expect(instance).to receive(:act_upon_results)
+      end
+
+      middleware = described_class.new
+
+      middleware.call(worker, job, queue) { nil }
+    end
+
+    it 'yields block' do
+      middleware = described_class.new
+
+      expect { |b| middleware.call(worker, job, queue, &b) }.to yield_control.once
+    end
+
+    it 'returns value of block' do
+      middleware = described_class.new
+
+      return_value = middleware.call(worker, job, queue) do
+        { value: 11 }
+      end
+
+      expect(return_value).to eq({ value: 11 })
+    end
+  end
+end
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 63dbf36b9e92c3..0e41b25a577eda 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -436,7 +436,9 @@
         arguments_logger: false, # We're not logging the regular messages for inline jobs
         skip_jobs: false # We're not skipping jobs for inline tests
       ).call(chain)
+
       chain.add DisableQueryLimit
+      chain.insert_after ::Gitlab::SidekiqMiddleware::RequestStoreMiddleware, ::Gitlab::QueryLimiting::SidekiqMiddleware
       chain.insert_after ::Gitlab::SidekiqMiddleware::RequestStoreMiddleware, IsolatedRequestStore
 
       example.run
diff --git a/spec/support/sidekiq_middleware.rb b/spec/support/sidekiq_middleware.rb
index cbd6163d46b25a..1419da55d2c8be 100644
--- a/spec/support/sidekiq_middleware.rb
+++ b/spec/support/sidekiq_middleware.rb
@@ -19,7 +19,7 @@ def with_sidekiq_server_middleware(&block)
 # SQL limit counter, query limiting is disabled during Sidekiq block
 class DisableQueryLimit
   def call(worker_instance, msg, queue)
-    ::Gitlab::QueryLimiting.disable!('https://mock-issue')
+    ::Gitlab::QueryLimiting.disable!('https://mock-issue') if Sidekiq::Testing.inline?
     yield
   ensure
     ::Gitlab::QueryLimiting.enable!
-- 
GitLab


From 673941f371ccd7f1c8b18b5bc379b592622aeaaa Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Tue, 28 May 2024 14:48:14 +1200
Subject: [PATCH 2/3] Fix inline sidekiq case correctly

Rather than disabling QueryLimiting completely, we put aside the
previous transaction, and restore that after completion.
---
 lib/gitlab/query_limiting/sidekiq_middleware.rb    |  2 --
 lib/gitlab/query_limiting/transaction.rb           |  4 +++-
 spec/lib/gitlab/query_limiting/transaction_spec.rb | 12 ++++++++++++
 spec/spec_helper.rb                                |  1 -
 spec/support/sidekiq_middleware.rb                 | 13 -------------
 5 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/lib/gitlab/query_limiting/sidekiq_middleware.rb b/lib/gitlab/query_limiting/sidekiq_middleware.rb
index dc215f3894650d..cfeeecce3b7a33 100644
--- a/lib/gitlab/query_limiting/sidekiq_middleware.rb
+++ b/lib/gitlab/query_limiting/sidekiq_middleware.rb
@@ -6,8 +6,6 @@ module QueryLimiting
     # certain amount of database queries.
     class SidekiqMiddleware
       def call(worker, _job, _queue)
-        #return yield if Sidekiq::Testing.inline?
-
         transaction, retval = ::Gitlab::QueryLimiting::Transaction.run do
           yield
         end
diff --git a/lib/gitlab/query_limiting/transaction.rb b/lib/gitlab/query_limiting/transaction.rb
index 728b0da5049b00..2b410f37d2a8b7 100644
--- a/lib/gitlab/query_limiting/transaction.rb
+++ b/lib/gitlab/query_limiting/transaction.rb
@@ -39,12 +39,14 @@ def self.current
       #
       #     retval # => 10
       def self.run
+        previous_transaction = current
+
         transaction = new
         Thread.current[THREAD_KEY] = transaction
 
         [transaction, yield]
       ensure
-        Thread.current[THREAD_KEY] = nil
+        Thread.current[THREAD_KEY] = previous_transaction
       end
 
       def initialize
diff --git a/spec/lib/gitlab/query_limiting/transaction_spec.rb b/spec/lib/gitlab/query_limiting/transaction_spec.rb
index 3ce8b08c77dc6a..4c653a261fe6d2 100644
--- a/spec/lib/gitlab/query_limiting/transaction_spec.rb
+++ b/spec/lib/gitlab/query_limiting/transaction_spec.rb
@@ -36,6 +36,18 @@
 
       expect(Thread.current[described_class::THREAD_KEY]).to be_nil
     end
+
+    it 'restores the previous transaction upon completion' do
+      original_transaction = described_class.new
+      Thread.current[described_class::THREAD_KEY] = original_transaction
+
+      new_transaction, _ret = described_class.run do
+        10
+      end
+
+      expect(new_transaction).not_to eq(original_transaction)
+      expect(Thread.current[described_class::THREAD_KEY]).to eq(original_transaction)
+    end
   end
 
   describe '#act_upon_results' do
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 0e41b25a577eda..9e1b7983121cca 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -437,7 +437,6 @@
         skip_jobs: false # We're not skipping jobs for inline tests
       ).call(chain)
 
-      chain.add DisableQueryLimit
       chain.insert_after ::Gitlab::SidekiqMiddleware::RequestStoreMiddleware, ::Gitlab::QueryLimiting::SidekiqMiddleware
       chain.insert_after ::Gitlab::SidekiqMiddleware::RequestStoreMiddleware, IsolatedRequestStore
 
diff --git a/spec/support/sidekiq_middleware.rb b/spec/support/sidekiq_middleware.rb
index 1419da55d2c8be..ae3680b11d533d 100644
--- a/spec/support/sidekiq_middleware.rb
+++ b/spec/support/sidekiq_middleware.rb
@@ -13,19 +13,6 @@ def with_sidekiq_server_middleware(&block)
 end
 # rubocop:enable RSpec/ModifySidekiqMiddleware
 
-# If Sidekiq::Testing.inline! is used, SQL transactions done inside
-# Sidekiq worker are included in the SQL query limit (in a real
-# deployment sidekiq worker is executed separately). To avoid increasing
-# SQL limit counter, query limiting is disabled during Sidekiq block
-class DisableQueryLimit
-  def call(worker_instance, msg, queue)
-    ::Gitlab::QueryLimiting.disable!('https://mock-issue') if Sidekiq::Testing.inline?
-    yield
-  ensure
-    ::Gitlab::QueryLimiting.enable!
-  end
-end
-
 # When running `Sidekiq::Testing.inline!` each job is using a request-store.
 # This middleware makes sure the values don't leak into eachother.
 class IsolatedRequestStore
-- 
GitLab


From 5b8172494deafe73549e2d47b2e0e6a3f2abee6c Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Tue, 28 May 2024 16:27:32 +1200
Subject: [PATCH 3/3] Allowlist workers which exceeds 100 queries

---
 app/workers/ci/create_downstream_pipeline_worker.rb       | 2 ++
 app/workers/create_pipeline_worker.rb                     | 2 ++
 app/workers/delete_user_worker.rb                         | 2 ++
 app/workers/group_destroy_worker.rb                       | 2 ++
 app/workers/group_import_worker.rb                        | 2 ++
 app/workers/mail_scheduler/notification_service_worker.rb | 2 ++
 app/workers/merge_requests/create_pipeline_worker.rb      | 2 ++
 app/workers/merge_worker.rb                               | 2 ++
 app/workers/new_merge_request_worker.rb                   | 2 ++
 app/workers/project_destroy_worker.rb                     | 2 ++
 app/workers/repository_import_worker.rb                   | 2 ++
 app/workers/update_merge_requests_worker.rb               | 2 ++
 ee/app/workers/merge_trains/refresh_worker.rb             | 2 ++
 13 files changed, 26 insertions(+)

diff --git a/app/workers/ci/create_downstream_pipeline_worker.rb b/app/workers/ci/create_downstream_pipeline_worker.rb
index 9f5ff45b8a6597..fcfb077d4b9c4a 100644
--- a/app/workers/ci/create_downstream_pipeline_worker.rb
+++ b/app/workers/ci/create_downstream_pipeline_worker.rb
@@ -10,6 +10,8 @@ class CreateDownstreamPipelineWorker # rubocop:disable Scalability/IdempotentWor
     urgency :high
 
     def perform(bridge_id)
+      Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/464668')
+
       ::Ci::Bridge.find_by_id(bridge_id).try do |bridge|
         result = ::Ci::CreateDownstreamPipelineService
           .new(bridge.project, bridge.user)
diff --git a/app/workers/create_pipeline_worker.rb b/app/workers/create_pipeline_worker.rb
index 3c442760d3830f..f183a222138269 100644
--- a/app/workers/create_pipeline_worker.rb
+++ b/app/workers/create_pipeline_worker.rb
@@ -15,6 +15,8 @@ class CreatePipelineWorker # rubocop:disable Scalability/IdempotentWorker
   loggable_arguments 2, 3, 4
 
   def perform(project_id, user_id, ref, source, params = {})
+    Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/464671')
+
     project = Project.find(project_id)
     user = User.find(user_id)
     params = params.deep_symbolize_keys
diff --git a/app/workers/delete_user_worker.rb b/app/workers/delete_user_worker.rb
index 4634ea8ff4ffb4..097dfc060d9a6a 100644
--- a/app/workers/delete_user_worker.rb
+++ b/app/workers/delete_user_worker.rb
@@ -11,6 +11,8 @@ class DeleteUserWorker # rubocop:disable Scalability/IdempotentWorker
   loggable_arguments 2
 
   def perform(current_user_id, delete_user_id, options = {})
+    Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/464672')
+
     delete_user = User.find_by_id(delete_user_id)
     return unless delete_user.present?
 
diff --git a/app/workers/group_destroy_worker.rb b/app/workers/group_destroy_worker.rb
index d8dd0b6d7b3bd7..28e0735f5d6a16 100644
--- a/app/workers/group_destroy_worker.rb
+++ b/app/workers/group_destroy_worker.rb
@@ -14,6 +14,8 @@ class GroupDestroyWorker
   deduplicate :until_executed, ttl: 2.hours
 
   def perform(group_id, user_id)
+    Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/464673')
+
     begin
       group = Group.find(group_id)
     rescue ActiveRecord::RecordNotFound
diff --git a/app/workers/group_import_worker.rb b/app/workers/group_import_worker.rb
index 198c6274166ee3..a6521e7804be28 100644
--- a/app/workers/group_import_worker.rb
+++ b/app/workers/group_import_worker.rb
@@ -9,6 +9,8 @@ class GroupImportWorker # rubocop:disable Scalability/IdempotentWorker
   feature_category :importers
 
   def perform(user_id, group_id)
+    Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/464675')
+
     current_user = User.find(user_id)
     group = Group.find(group_id)
     group_import_state = group.import_state
diff --git a/app/workers/mail_scheduler/notification_service_worker.rb b/app/workers/mail_scheduler/notification_service_worker.rb
index 8206a17021bbc1..7f3066f883ef9e 100644
--- a/app/workers/mail_scheduler/notification_service_worker.rb
+++ b/app/workers/mail_scheduler/notification_service_worker.rb
@@ -16,6 +16,8 @@ class NotificationServiceWorker # rubocop:disable Scalability/IdempotentWorker
     loggable_arguments 0
 
     def perform(meth, *args)
+      Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/464670')
+
       check_arguments!(args)
 
       if NotificationService.permitted_actions.exclude?(meth.to_sym)
diff --git a/app/workers/merge_requests/create_pipeline_worker.rb b/app/workers/merge_requests/create_pipeline_worker.rb
index e92e557284c070..12f114f49691fd 100644
--- a/app/workers/merge_requests/create_pipeline_worker.rb
+++ b/app/workers/merge_requests/create_pipeline_worker.rb
@@ -16,6 +16,8 @@ class CreatePipelineWorker
     idempotent!
 
     def perform(project_id, user_id, merge_request_id, params = {})
+      Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/464679')
+
       project = Project.find_by_id(project_id)
       return unless project
 
diff --git a/app/workers/merge_worker.rb b/app/workers/merge_worker.rb
index 29f0c0bbbf4fa5..2979a77e19f277 100644
--- a/app/workers/merge_worker.rb
+++ b/app/workers/merge_worker.rb
@@ -16,6 +16,8 @@ class MergeWorker # rubocop:disable Scalability/IdempotentWorker
   deduplicate :until_executed, including_scheduled: true
 
   def perform(merge_request_id, current_user_id, params)
+    Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/464676')
+
     begin
       current_user = User.find(current_user_id)
       merge_request = MergeRequest.find(merge_request_id)
diff --git a/app/workers/new_merge_request_worker.rb b/app/workers/new_merge_request_worker.rb
index e2e738f79a5792..a11f26088b38e0 100644
--- a/app/workers/new_merge_request_worker.rb
+++ b/app/workers/new_merge_request_worker.rb
@@ -17,6 +17,8 @@ class NewMergeRequestWorker # rubocop:disable Scalability/IdempotentWorker
   weight 2
 
   def perform(merge_request_id, user_id)
+    Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/337182')
+
     return unless objects_found?(merge_request_id, user_id)
     return if issuable.prepared?
 
diff --git a/app/workers/project_destroy_worker.rb b/app/workers/project_destroy_worker.rb
index 181eebe56e83e8..565dfab9930f51 100644
--- a/app/workers/project_destroy_worker.rb
+++ b/app/workers/project_destroy_worker.rb
@@ -14,6 +14,8 @@ class ProjectDestroyWorker
   deduplicate :until_executed, ttl: 2.hours
 
   def perform(project_id, user_id, params)
+    Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/333366')
+
     project = Project.find(project_id)
     user = User.find(user_id)
 
diff --git a/app/workers/repository_import_worker.rb b/app/workers/repository_import_worker.rb
index c16071e8e31c05..94bbd812b68764 100644
--- a/app/workers/repository_import_worker.rb
+++ b/app/workers/repository_import_worker.rb
@@ -15,6 +15,8 @@ class RepositoryImportWorker # rubocop:disable Scalability/IdempotentWorker
   worker_resource_boundary :memory
 
   def perform(project_id)
+    Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/464677')
+
     @project = Project.find_by_id(project_id)
     return if project.nil? || !start_import?
 
diff --git a/app/workers/update_merge_requests_worker.rb b/app/workers/update_merge_requests_worker.rb
index caf46c1ac4ef2d..d61b6e601b87a7 100644
--- a/app/workers/update_merge_requests_worker.rb
+++ b/app/workers/update_merge_requests_worker.rb
@@ -14,6 +14,8 @@ class UpdateMergeRequestsWorker # rubocop:disable Scalability/IdempotentWorker
   loggable_arguments 2, 3, 4
 
   def perform(project_id, user_id, oldrev, newrev, ref, params = {})
+    Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/24907')
+
     project = Project.find_by_id(project_id)
     return unless project
 
diff --git a/ee/app/workers/merge_trains/refresh_worker.rb b/ee/app/workers/merge_trains/refresh_worker.rb
index 521701c8b32dbf..f1442d192df378 100644
--- a/ee/app/workers/merge_trains/refresh_worker.rb
+++ b/ee/app/workers/merge_trains/refresh_worker.rb
@@ -17,6 +17,8 @@ class RefreshWorker
     idempotent!
 
     def perform(target_project_id, target_branch)
+      Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/-/issues/464678')
+
       ::MergeTrains::RefreshService
         .new(target_project_id, target_branch)
         .execute
-- 
GitLab