Add Sidekiq middleware stack integration tests

Before this change, Sidekiq middleware components were only tested
through unit tests. This add basic integration tests for our server-side
and client-side middleware stacks, to ensure that messages pass through
the stack without causing exceptions.

The change also ensures that each middleware yields control as expected.

config/initializers/sidekiq.rb

@@ -31,21 +31,17 @@ enable_json_logs = Gitlab.config.sidekiq.log_format == 'json'
 enable_sidekiq_memory_killer = ENV['SIDEKIQ_MEMORY_KILLER_MAX_RSS'].to_i.nonzero?
 use_sidekiq_daemon_memory_killer = ENV["SIDEKIQ_DAEMON_MEMORY_KILLER"].to_i.nonzero?
 use_sidekiq_legacy_memory_killer = !use_sidekiq_daemon_memory_killer
+use_request_store = ENV['SIDEKIQ_REQUEST_STORE'].to_i.nonzero?
 
 Sidekiq.configure_server do |config|
   config.redis = queues_config_hash
 
-  config.server_middleware do |chain|
-    chain.add Gitlab::SidekiqMiddleware::Monitor
-    chain.add Gitlab::SidekiqMiddleware::Metrics if Settings.monitoring.sidekiq_exporter
-    chain.add Gitlab::SidekiqMiddleware::ArgumentsLogger if ENV['SIDEKIQ_LOG_ARGUMENTS'] && !enable_json_logs
-    chain.add Gitlab::SidekiqMiddleware::MemoryKiller if enable_sidekiq_memory_killer && use_sidekiq_legacy_memory_killer
-    chain.add Gitlab::SidekiqMiddleware::RequestStoreMiddleware unless ENV['SIDEKIQ_REQUEST_STORE'] == '0'
-    chain.add Gitlab::SidekiqMiddleware::BatchLoader
-    chain.add Gitlab::SidekiqMiddleware::CorrelationLogger
-    chain.add Gitlab::SidekiqMiddleware::InstrumentationLogger
-    chain.add Gitlab::SidekiqStatus::ServerMiddleware
-  end
+  config.server_middleware(&Gitlab::SidekiqMiddleware.server_configurator({
+    metrics: Settings.monitoring.sidekiq_exporter,
+    arguments_logger: ENV['SIDEKIQ_LOG_ARGUMENTS'] && !enable_json_logs,
+    memory_killer: enable_sidekiq_memory_killer && use_sidekiq_legacy_memory_killer,
+    request_store: use_request_store
+  }))
 
   if enable_json_logs
     Sidekiq.logger.formatter = Gitlab::SidekiqLogging::JSONFormatter.new
@@ -56,10 +52,7 @@ Sidekiq.configure_server do |config|
     config.error_handlers << Gitlab::SidekiqLogging::ExceptionHandler.new
   end
 
-  config.client_middleware do |chain|
-    chain.add Gitlab::SidekiqStatus::ClientMiddleware
-    chain.add Gitlab::SidekiqMiddleware::CorrelationInjector
-  end
+  config.client_middleware(&Gitlab::SidekiqMiddleware.client_configurator)
 
   config.on :startup do
     # Clear any connections that might have been obtained before starting

lib/gitlab/sidekiq_middleware.rb

+# frozen_string_literal: true
+
+module Gitlab
+  # The SidekiqMiddleware class is responsible for configuring the
+  # middleware stacks used in the client and server middlewares
+  module SidekiqMiddleware
+    # The result of this method should be passed to
+    # Sidekiq's `config.server_middleware` method
+    # eg: `config.server_middleware(&Gitlab::SidekiqMiddleware.server_configurator)`
+    def self.server_configurator(metrics: true, arguments_logger: true, memory_killer: true, request_store: true)
+      lambda do |chain|
+        chain.add Gitlab::SidekiqMiddleware::Monitor
+        chain.add Gitlab::SidekiqMiddleware::Metrics if metrics
+        chain.add Gitlab::SidekiqMiddleware::ArgumentsLogger if arguments_logger
+        chain.add Gitlab::SidekiqMiddleware::MemoryKiller if memory_killer
+        chain.add Gitlab::SidekiqMiddleware::RequestStoreMiddleware if request_store
+        chain.add Gitlab::SidekiqMiddleware::BatchLoader
+        chain.add Gitlab::SidekiqMiddleware::CorrelationLogger
+        chain.add Gitlab::SidekiqMiddleware::InstrumentationLogger
+        chain.add Gitlab::SidekiqStatus::ServerMiddleware
+      end
+    end
+
+    # The result of this method should be passed to
+    # Sidekiq's `config.client_middleware` method
+    # eg: `config.client_middleware(&Gitlab::SidekiqMiddleware.client_configurator)`
+    def self.client_configurator
+      lambda do |chain|
+        chain.add Gitlab::SidekiqStatus::ClientMiddleware
+        chain.add Gitlab::SidekiqMiddleware::CorrelationInjector
+      end
+    end
+  end
+end

spec/lib/gitlab/sidekiq_middleware_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+require 'sidekiq/testing'
+
+describe Gitlab::SidekiqMiddleware do
+  class TestWorker
+    include Sidekiq::Worker
+
+    def perform(_arg)
+    end
+  end
+
+  around do |example|
+    Sidekiq::Testing.inline! { example.run }
+  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
+    let(:middleware_expected_args) { [a_kind_of(worker_class), hash_including({ 'args' => job_args }), anything] }
+    let(:all_sidekiq_middlewares) do
+      [
+       Gitlab::SidekiqMiddleware::Monitor,
+       Gitlab::SidekiqMiddleware::BatchLoader,
+       Gitlab::SidekiqMiddleware::CorrelationLogger,
+       Gitlab::SidekiqMiddleware::InstrumentationLogger,
+       Gitlab::SidekiqStatus::ServerMiddleware,
+       Gitlab::SidekiqMiddleware::Metrics,
+       Gitlab::SidekiqMiddleware::ArgumentsLogger,
+       Gitlab::SidekiqMiddleware::MemoryKiller,
+       Gitlab::SidekiqMiddleware::RequestStoreMiddleware
+      ]
+    end
+    let(:enabled_sidekiq_middlewares) { all_sidekiq_middlewares - disabled_sidekiq_middlewares }
+
+    before do
+      Sidekiq::Testing.server_middleware.clear
+      Sidekiq::Testing.server_middleware(&described_class.server_configurator(
+        metrics: metrics,
+        arguments_logger: arguments_logger,
+        memory_killer: memory_killer,
+        request_store: request_store
+      ))
+
+      enabled_sidekiq_middlewares.each do |middleware|
+        expect_any_instance_of(middleware).to receive(:call).with(*middleware_expected_args).once.and_call_original
+      end
+
+      disabled_sidekiq_middlewares.each do |middleware|
+        expect_any_instance_of(Gitlab::SidekiqMiddleware::ArgumentsLogger).not_to receive(:call)
+      end
+    end
+
+    context "all optional middlewares off" do
+      let(:metrics) { false }
+      let(:arguments_logger) { false }
+      let(:memory_killer) { false }
+      let(:request_store) { false }
+      let(:disabled_sidekiq_middlewares) do
+        [
+          Gitlab::SidekiqMiddleware::Metrics,
+          Gitlab::SidekiqMiddleware::ArgumentsLogger,
+          Gitlab::SidekiqMiddleware::MemoryKiller,
+          Gitlab::SidekiqMiddleware::RequestStoreMiddleware
+        ]
+      end
+
+      it "passes through server middlewares" do
+        worker_class.perform_async(*job_args)
+      end
+    end
+
+    context "all optional middlewares on" do
+      let(:metrics) { true }
+      let(:arguments_logger) { true }
+      let(:memory_killer) { true }
+      let(:request_store) { true }
+      let(:disabled_sidekiq_middlewares) { [] }
+
+      it "passes through server middlewares" do
+        worker_class.perform_async(*job_args)
+      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(:redis_pool) { Sidekiq.redis_pool }
+    let(:middleware_expected_args) { [worker_class_arg, job, queue, redis_pool] }
+
+    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
+        expect_any_instance_of(Gitlab::SidekiqStatus::ClientMiddleware).to receive(:call).with(*middleware_expected_args).once.and_call_original
+        expect_any_instance_of(Gitlab::SidekiqMiddleware::CorrelationInjector).to receive(:call).with(*middleware_expected_args).once.and_call_original
+
+        expect { |b| chain.invoke(worker_class_arg, job, queue, redis_pool, &b) }.to yield_control.once
+      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 client middleware chain"
+    end
+
+    context "handles string worker_class values" do
+      let(:worker_class_arg) { worker_class }
+
+      it_behaves_like "a client middleware chain"
+    end
+  end
+end