1 files changed, 150 insertions, 153 deletions

diff --git a/spec/lib/gitlab/sidekiq_middleware_spec.rb b/spec/lib/gitlab/sidekiq_middleware_spec.rb
index 0efdef0c999..5e4e79e818e 100644
--- a/spec/lib/gitlab/sidekiq_middleware_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware_spec.rb
@@ -4,215 +4,212 @@ require 'spec_helper'
 require 'sidekiq/testing'
 
 RSpec.describe Gitlab::SidekiqMiddleware do
-  before do
-    stub_const('TestWorker', Class.new)
+  let(:job_args) { [0.01] }
+  let(:disabled_sidekiq_middlewares) { [] }
+  let(:chain) { Sidekiq::Middleware::Chain.new }
+  let(:queue) { 'test' }
+  let(:enabled_sidekiq_middlewares) { all_sidekiq_middlewares - disabled_sidekiq_middlewares }
+  let(:worker_class) do
+    Class.new do
+      def self.name
+        'TestWorker'
+      end
 
-    TestWorker.class_eval do
-      include Sidekiq::Worker
       include ApplicationWorker
 
-      def perform(_arg)
+      def perform(*args)
         Gitlab::SafeRequestStore['gitaly_call_actual'] = 1
         Gitlab::SafeRequestStore[:gitaly_query_time] = 5
       end
     end
   end
 
-  around do |example|
-    Sidekiq::Testing.inline! { example.run }
+  before do
+    stub_const('TestWorker', worker_class)
   end
 
-  let(:worker_class) { TestWorker }
-  let(:job_args) { [0.01] }
-
-  # The test sets up a new server middleware stack, ensuring that the
-  # appropriate middlewares, as passed into server_configurator,
-  # are invoked.
-  # Additionally the test ensure that each middleware is
-  # 1) not failing
-  # 2) yielding exactly once
-  describe '.server_configurator' do
-    around do |example|
-      with_sidekiq_server_middleware do |chain|
-        described_class.server_configurator(
-          metrics: metrics,
-          arguments_logger: arguments_logger,
-          memory_killer: memory_killer
-        ).call(chain)
+  shared_examples "a middleware chain" do |load_balancing_enabled|
+    before do
+      allow(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(load_balancing_enabled)
+      configurator.call(chain)
+    end
 
-        example.run
+    it "passes through the right middlewares", :aggregate_failures do
+      enabled_sidekiq_middlewares.each do |middleware|
+        expect_next_instances_of(middleware, 1, true) do |middleware_instance|
+          expect(middleware_instance).to receive(:call).with(*middleware_expected_args).once.and_call_original
+        end
       end
+
+      expect { |b| chain.invoke(*worker_args, &b) }.to yield_control.once
     end
+  end
+
+  shared_examples "a middleware chain for mailer" do |load_balancing_enabled|
+    let(:worker_class) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper }
 
-    let(:middleware_expected_args) { [a_kind_of(worker_class), hash_including({ 'args' => job_args }), anything] }
+    it_behaves_like "a middleware chain", load_balancing_enabled
+  end
+
+  describe '.server_configurator' do
+    let(:configurator) { described_class.server_configurator }
+    let(:worker_args) { [worker_class.new, { 'args' => job_args }, queue] }
+    let(:middleware_expected_args) { [a_kind_of(worker_class), hash_including({ 'args' => job_args }), queue] }
     let(:all_sidekiq_middlewares) do
       [
-       Gitlab::SidekiqMiddleware::Monitor,
-       Gitlab::SidekiqMiddleware::BatchLoader,
-       Labkit::Middleware::Sidekiq::Server,
-       Gitlab::SidekiqMiddleware::InstrumentationLogger,
-       Gitlab::SidekiqVersioning::Middleware,
-       Gitlab::SidekiqStatus::ServerMiddleware,
-       Gitlab::SidekiqMiddleware::ServerMetrics,
-       Gitlab::SidekiqMiddleware::ArgumentsLogger,
-       Gitlab::SidekiqMiddleware::MemoryKiller,
-       Gitlab::SidekiqMiddleware::RequestStoreMiddleware,
-       Gitlab::SidekiqMiddleware::ExtraDoneLogMetadata,
-       Gitlab::SidekiqMiddleware::WorkerContext::Server,
-       Gitlab::SidekiqMiddleware::AdminMode::Server,
-       Gitlab::SidekiqMiddleware::DuplicateJobs::Server
+        ::Gitlab::SidekiqMiddleware::Monitor,
+        ::Gitlab::SidekiqMiddleware::ServerMetrics,
+        ::Gitlab::SidekiqMiddleware::ArgumentsLogger,
+        ::Gitlab::SidekiqMiddleware::MemoryKiller,
+        ::Gitlab::SidekiqMiddleware::RequestStoreMiddleware,
+        ::Gitlab::SidekiqMiddleware::ExtraDoneLogMetadata,
+        ::Gitlab::SidekiqMiddleware::BatchLoader,
+        ::Labkit::Middleware::Sidekiq::Server,
+        ::Gitlab::SidekiqMiddleware::InstrumentationLogger,
+        ::Gitlab::Database::LoadBalancing::SidekiqServerMiddleware,
+        ::Gitlab::SidekiqMiddleware::AdminMode::Server,
+        ::Gitlab::SidekiqVersioning::Middleware,
+        ::Gitlab::SidekiqStatus::ServerMiddleware,
+        ::Gitlab::SidekiqMiddleware::WorkerContext::Server,
+        ::Gitlab::SidekiqMiddleware::DuplicateJobs::Server
       ]
     end
 
-    let(:enabled_sidekiq_middlewares) { all_sidekiq_middlewares - disabled_sidekiq_middlewares }
+    describe "server metrics" do
+      around do |example|
+        with_sidekiq_server_middleware do |chain|
+          described_class.server_configurator(
+            metrics: true,
+            arguments_logger: true,
+            memory_killer: true
+          ).call(chain)
 
-    shared_examples "a server middleware chain" do
-      it "passes through the right server middlewares" do
-        enabled_sidekiq_middlewares.each do |middleware|
-          expect_next_instance_of(middleware) do |middleware_instance|
-            expect(middleware_instance).to receive(:call).with(*middleware_expected_args).once.and_call_original
-          end
+          Sidekiq::Testing.inline! { example.run }
         end
+      end
+      let(:gitaly_histogram) { double(:gitaly_histogram) }
 
-        disabled_sidekiq_middlewares.each do |middleware|
-          expect(middleware).not_to receive(:new)
-        end
+      before do
+        allow(Gitlab::Metrics).to receive(:histogram).and_call_original
+
+        allow(Gitlab::Metrics).to receive(:histogram)
+          .with(:sidekiq_jobs_gitaly_seconds, anything, anything, anything)
+          .and_return(gitaly_histogram)
+      end
+
+      it "records correct Gitaly duration" do
+        expect(gitaly_histogram).to receive(:observe).with(anything, 5.0)
 
         worker_class.perform_async(*job_args)
       end
     end
 
-    shared_examples "a server middleware chain for mailer" do
-      let(:worker_class) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper }
-      let(:job_args) do
-        [
-          {
-            "job_class" => "ActionMailer::MailDeliveryJob",
-            "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e",
-            "provider_job_id" => nil,
-            "queue_name" => "mailers",
-            "priority" => nil,
-            "arguments" => [
-              "Notify",
-              "test_email",
-              "deliver_now",
-              {
-                "args" => [
-                  "test@example.com",
-                  "subject",
-                  "body"
-                ],
-                ActiveJob::Arguments.const_get('RUBY2_KEYWORDS_KEY', false) => ["args"]
-              }
-            ],
-            "executions" => 0,
-            "exception_executions" => {},
-            "locale" => "en",
-            "timezone" => "UTC",
-            "enqueued_at" => "2020-07-27T07:43:31Z"
-          }
-        ]
+    context "all optional middlewares on" do
+      context "when load balancing is enabled" do
+        before do
+          allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
+        end
+
+        it_behaves_like "a middleware chain", true
+        it_behaves_like "a middleware chain for mailer", true
       end
 
-      it_behaves_like "a server middleware chain"
-    end
+      context "when load balancing is disabled" do
+        let(:disabled_sidekiq_middlewares) do
+          [
+            Gitlab::Database::LoadBalancing::SidekiqServerMiddleware
+          ]
+        end
 
-    context "all optional middlewares off" do
-      let(:metrics) { false }
-      let(:arguments_logger) { false }
-      let(:memory_killer) { false }
-      let(:disabled_sidekiq_middlewares) do
-        [
-          Gitlab::SidekiqMiddleware::ServerMetrics,
-          Gitlab::SidekiqMiddleware::ArgumentsLogger,
-          Gitlab::SidekiqMiddleware::MemoryKiller
-        ]
+        it_behaves_like "a middleware chain", false
+        it_behaves_like "a middleware chain for mailer", false
       end
-
-      it_behaves_like "a server middleware chain"
-      it_behaves_like "a server middleware chain for mailer"
     end
 
-    context "all optional middlewares on" do
-      let(:metrics) { true }
-      let(:arguments_logger) { true }
-      let(:memory_killer) { true }
-      let(:disabled_sidekiq_middlewares) { [] }
-
-      it_behaves_like "a server middleware chain"
-      it_behaves_like "a server middleware chain for mailer"
+    context "all optional middlewares off" do
+      let(:configurator) do
+        described_class.server_configurator(
+          metrics: false,
+          arguments_logger: false,
+          memory_killer: false
+        )
+      end
 
-      context "server metrics" do
-        let(:gitaly_histogram) { double(:gitaly_histogram) }
+      context "when load balancing is enabled" do
+        let(:disabled_sidekiq_middlewares) do
+          [
+            Gitlab::SidekiqMiddleware::ServerMetrics,
+            Gitlab::SidekiqMiddleware::ArgumentsLogger,
+            Gitlab::SidekiqMiddleware::MemoryKiller
+          ]
+        end
 
         before do
-          allow(Gitlab::Metrics).to receive(:histogram).and_call_original
-
-          allow(Gitlab::Metrics).to receive(:histogram)
-                                      .with(:sidekiq_jobs_gitaly_seconds, anything, anything, anything)
-                                      .and_return(gitaly_histogram)
+          allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
         end
 
-        it "records correct Gitaly duration" do
-          expect(gitaly_histogram).to receive(:observe).with(anything, 5.0)
+        it_behaves_like "a middleware chain", true
+        it_behaves_like "a middleware chain for mailer", true
+      end
 
-          worker_class.perform_async(*job_args)
+      context "when load balancing is disabled" do
+        let(:disabled_sidekiq_middlewares) do
+          [
+            Gitlab::SidekiqMiddleware::ServerMetrics,
+            Gitlab::SidekiqMiddleware::ArgumentsLogger,
+            Gitlab::SidekiqMiddleware::MemoryKiller,
+            Gitlab::Database::LoadBalancing::SidekiqServerMiddleware
+          ]
         end
+
+        it_behaves_like "a middleware chain", false
+        it_behaves_like "a middleware chain for mailer", false
       end
     end
   end
 
-  # The test sets up a new client middleware stack. The test ensures
-  # that each middleware is:
-  # 1) not failing
-  # 2) yielding exactly once
   describe '.client_configurator' do
-    let(:chain) { Sidekiq::Middleware::Chain.new }
-    let(:job) { { 'args' => job_args } }
-    let(:queue) { 'default' }
+    let(:configurator) { described_class.client_configurator }
     let(:redis_pool) { Sidekiq.redis_pool }
-    let(:middleware_expected_args) { [worker_class_arg, job, queue, redis_pool] }
-    let(:expected_middlewares) do
+    let(:middleware_expected_args) { [worker_class, hash_including({ 'args' => job_args }), queue, redis_pool] }
+    let(:worker_args) { [worker_class, { 'args' => job_args }, queue, redis_pool] }
+    let(:all_sidekiq_middlewares) do
       [
-         ::Gitlab::SidekiqMiddleware::WorkerContext::Client,
-         ::Labkit::Middleware::Sidekiq::Client,
-         ::Gitlab::SidekiqMiddleware::DuplicateJobs::Client,
-         ::Gitlab::SidekiqStatus::ClientMiddleware,
-         ::Gitlab::SidekiqMiddleware::AdminMode::Client,
-         ::Gitlab::SidekiqMiddleware::SizeLimiter::Client,
-         ::Gitlab::SidekiqMiddleware::ClientMetrics
+        ::Gitlab::SidekiqMiddleware::WorkerContext::Client,
+        ::Labkit::Middleware::Sidekiq::Client,
+        ::Gitlab::SidekiqMiddleware::DuplicateJobs::Client,
+        ::Gitlab::SidekiqStatus::ClientMiddleware,
+        ::Gitlab::SidekiqMiddleware::AdminMode::Client,
+        ::Gitlab::SidekiqMiddleware::SizeLimiter::Client,
+        ::Gitlab::SidekiqMiddleware::ClientMetrics,
+        ::Gitlab::Database::LoadBalancing::SidekiqClientMiddleware
       ]
     end
 
-    before do
-      described_class.client_configurator.call(chain)
-    end
-
-    shared_examples "a client middleware chain" do
-      # Its possible that a middleware could accidentally omit a yield call
-      # this will prevent the full middleware chain from being executed.
-      # This test ensures that this does not happen
-      it "invokes the chain" do
-        expected_middlewares do |middleware|
-          expect_any_instance_of(middleware).to receive(:call).with(*middleware_expected_args).once.ordered.and_call_original
-        end
-
-        expect { |b| chain.invoke(worker_class_arg, job, queue, redis_pool, &b) }.to yield_control.once
+    context "when load balancing is disabled" do
+      let(:disabled_sidekiq_middlewares) do
+        [
+          Gitlab::Database::LoadBalancing::SidekiqClientMiddleware
+        ]
       end
-    end
 
-    # Sidekiq documentation states that the worker class could be a string
-    # or a class reference. We should test for both
-    context "handles string worker_class values" do
-      let(:worker_class_arg) { worker_class.to_s }
+      it_behaves_like "a middleware chain", false
+      it_behaves_like "a middleware chain for mailer", false
 
-      it_behaves_like "a client middleware chain"
-    end
+      # Sidekiq documentation states that the worker class could be a string
+      # or a class reference. We should test for both
+      context "worker_class as string value" do
+        let(:worker_args) { [worker_class.to_s, { 'args' => job_args }, queue, redis_pool] }
+        let(:middleware_expected_args) { [worker_class.to_s, hash_including({ 'args' => job_args }), queue, redis_pool] }
 
-    context "handles string worker_class values" do
-      let(:worker_class_arg) { worker_class }
+        it_behaves_like "a middleware chain", false
+        it_behaves_like "a middleware chain for mailer", false
+      end
+    end
 
-      it_behaves_like "a client middleware chain"
+    context "when load balancing is enabled" do
+      it_behaves_like "a middleware chain", true
+      it_behaves_like "a middleware chain for mailer", true
     end
   end
 end