Move overridden Load balancing calls in SidekiqMiddleware to Core

This naively moves the specs, we will rationalise them in next commit

ee/lib/ee/gitlab/sidekiq_middleware.rb

-# frozen_string_literal: true
-
-module EE
-  module Gitlab
-    # The SidekiqMiddleware class is responsible for configuring the
-    # middleware stacks used in the client and server middlewares
-    module SidekiqMiddleware
-      extend ::Gitlab::Utils::Override
-
-      override :server_configurator
-      def server_configurator(metrics: true, arguments_logger: true, memory_killer: true)
-        lambda do |chain|
-          super.call(chain)
-
-          if load_balancing_enabled?
-            chain.insert_after(::Gitlab::SidekiqMiddleware::InstrumentationLogger,
-                               ::Gitlab::Database::LoadBalancing::SidekiqServerMiddleware)
-          end
-        end
-      end
-
-      override :client_configurator
-      def client_configurator
-        lambda do |chain|
-          super.call(chain)
-
-          chain.add ::Gitlab::Database::LoadBalancing::SidekiqClientMiddleware if load_balancing_enabled?
-        end
-      end
-
-      private
-
-      def load_balancing_enabled?
-        ::Gitlab::Database::LoadBalancing.enable?
-      end
-    end
-  end
-end

ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb

-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Gitlab::SidekiqMiddleware do
-  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
-
-      include ApplicationWorker
-
-      def perform(*args)
-      end
-    end
-  end
-
-  before do
-    stub_const('TestWorker', worker_class)
-  end
-
-  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
-
-    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 }
-
-    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::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
-
-    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
-
-    context "when load balancing is disabled" do
-      let(:disabled_sidekiq_middlewares) do
-        [
-          Gitlab::Database::LoadBalancing::SidekiqServerMiddleware
-        ]
-      end
-
-      it_behaves_like "a middleware chain", false
-      it_behaves_like "a middleware chain for mailer", false
-    end
-  end
-
-  describe '.client_configurator' do
-    let(:configurator) { described_class.client_configurator }
-    let(:redis_pool) { Sidekiq.redis_pool }
-    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::Database::LoadBalancing::SidekiqClientMiddleware
-      ]
-    end
-
-    context "when load balancing is disabled" do
-      let(:disabled_sidekiq_middlewares) do
-        [
-          Gitlab::Database::LoadBalancing::SidekiqClientMiddleware
-        ]
-      end
-
-      it_behaves_like "a middleware chain", false
-      it_behaves_like "a middleware chain for mailer", false
-    end
-
-    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

lib/gitlab/sidekiq_middleware.rb

@@ -20,6 +20,7 @@ module Gitlab
         chain.add ::Gitlab::SidekiqMiddleware::BatchLoader
         chain.add ::Labkit::Middleware::Sidekiq::Server
         chain.add ::Gitlab::SidekiqMiddleware::InstrumentationLogger
+        chain.add ::Gitlab::Database::LoadBalancing::SidekiqServerMiddleware if load_balancing_enabled?
         chain.add ::Gitlab::SidekiqMiddleware::AdminMode::Server
         chain.add ::Gitlab::SidekiqVersioning::Middleware
         chain.add ::Gitlab::SidekiqStatus::ServerMiddleware
@@ -41,9 +42,13 @@ module Gitlab
         # Size limiter should be placed at the bottom, but before the metrics midleware
         chain.add ::Gitlab::SidekiqMiddleware::SizeLimiter::Client
         chain.add ::Gitlab::SidekiqMiddleware::ClientMetrics
+        chain.add ::Gitlab::Database::LoadBalancing::SidekiqClientMiddleware if load_balancing_enabled?
       end
     end
+
+    def self.load_balancing_enabled?
+      ::Gitlab::Database::LoadBalancing.enable?
+    end
+    private_class_method :load_balancing_enabled?
   end
 end
-
-Gitlab::SidekiqMiddleware.singleton_class.prepend_mod_with('Gitlab::SidekiqMiddleware')

spec/lib/gitlab/sidekiq_middleware_spec.rb

@@ -216,3 +216,130 @@ RSpec.describe Gitlab::SidekiqMiddleware do
     end
   end
 end
+
+RSpec.describe Gitlab::SidekiqMiddleware, 'load balancing' do
+  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
+
+      include ApplicationWorker
+
+      def perform(*args)
+      end
+    end
+  end
+
+  before do
+    stub_const('TestWorker', worker_class)
+  end
+
+  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
+
+    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 }
+
+    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::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
+
+    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
+
+    context "when load balancing is disabled" do
+      let(:disabled_sidekiq_middlewares) do
+        [
+          Gitlab::Database::LoadBalancing::SidekiqServerMiddleware
+        ]
+      end
+
+      it_behaves_like "a middleware chain", false
+      it_behaves_like "a middleware chain for mailer", false
+    end
+  end
+
+  describe '.client_configurator' do
+    let(:configurator) { described_class.client_configurator }
+    let(:redis_pool) { Sidekiq.redis_pool }
+    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::Database::LoadBalancing::SidekiqClientMiddleware
+      ]
+    end
+
+    context "when load balancing is disabled" do
+      let(:disabled_sidekiq_middlewares) do
+        [
+          Gitlab::Database::LoadBalancing::SidekiqClientMiddleware
+        ]
+      end
+
+      it_behaves_like "a middleware chain", false
+      it_behaves_like "a middleware chain for mailer", false
+    end
+
+    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