From a8fbb12f5b7ca0fd7c34944c9be1f2df8acf888d Mon Sep 17 00:00:00 2001
From: nmilojevic1 <nmilojevic@gitlab.com>
Date: Fri, 5 Mar 2021 17:35:38 +0100
Subject: [PATCH 1/9] Enable load_balancing for sidekiq

- Introduce server and client sidekiq middleware
- Introduce data consistencies for workers
---
 app/workers/concerns/worker_attributes.rb     | 19 +++++++
 .../development/load_balancer_for_sidekiq.yml |  8 +++
 .../logging_for_all_caught_up_method.yml      |  8 +++
 config/initializers/load_balancing.rb         | 14 +++++
 ee/lib/gitlab/database/load_balancing.rb      |  6 ++-
 .../database/load_balancing/load_balancer.rb  | 27 +++++++++-
 .../sidekiq_client_middleware.rb              | 23 ++++++++
 .../sidekiq_server_middleware.rb              | 53 +++++++++++++++++++
 8 files changed, 156 insertions(+), 2 deletions(-)
 create mode 100644 config/feature_flags/development/load_balancer_for_sidekiq.yml
 create mode 100644 config/feature_flags/development/logging_for_all_caught_up_method.yml
 create mode 100644 ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
 create mode 100644 ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb

diff --git a/app/workers/concerns/worker_attributes.rb b/app/workers/concerns/worker_attributes.rb
index 042508d08f2183..13357c37e07190 100644
--- a/app/workers/concerns/worker_attributes.rb
+++ b/app/workers/concerns/worker_attributes.rb
@@ -11,6 +11,8 @@ module WorkerAttributes
   # Urgencies that workers can declare through the `urgencies` attribute
   VALID_URGENCIES = [:high, :low, :throttled].freeze
 
+  VALID_DATA_CONSISTENCIES = [:always, :sticky ,:delayed].freeze
+
   NAMESPACE_WEIGHTS = {
     auto_devops: 2,
     auto_merge: 3,
@@ -69,6 +71,23 @@ def get_urgency
       class_attributes[:urgency] || :low
     end
 
+    def data_consistency (data_consistency)
+      raise "Invalid data consistency: #{data_consistency}" unless VALID_DATA_CONSISTENCIES.include?(data_consistency)
+
+      class_attributes[:data_consistency] = data_consistency
+    end
+
+    VALID_DATA_CONSISTENCIES.each do |data_consistency|
+      client_method_name = "#{data_consistency}?".to_sym
+
+      define_singleton_method(client_method_name) do
+        class_attributes[:data_consistency] == data_consistency
+      end
+    end
+
+    def get_data_consistency
+      class_attributes[:data_consistency] || :always
+    end
     # Set this attribute on a job when it will call to services outside of the
     # application, such as 3rd party applications, other k8s clusters etc See
     # doc/development/sidekiq_style_guide.md#jobs-with-external-dependencies for
diff --git a/config/feature_flags/development/load_balancer_for_sidekiq.yml b/config/feature_flags/development/load_balancer_for_sidekiq.yml
new file mode 100644
index 00000000000000..047d76d53df22c
--- /dev/null
+++ b/config/feature_flags/development/load_balancer_for_sidekiq.yml
@@ -0,0 +1,8 @@
+---
+name: load_balancer_for_sidekiq
+introduced_by_url:
+rollout_issue_url:
+milestone: '13.10'
+type: development
+group: group::memory
+default_enabled: false
diff --git a/config/feature_flags/development/logging_for_all_caught_up_method.yml b/config/feature_flags/development/logging_for_all_caught_up_method.yml
new file mode 100644
index 00000000000000..53a30a3b921b82
--- /dev/null
+++ b/config/feature_flags/development/logging_for_all_caught_up_method.yml
@@ -0,0 +1,8 @@
+---
+name: logging_for_all_caught_up_method
+introduced_by_url:
+rollout_issue_url:
+milestone: '13.10'
+type: development
+group: group::memory
+default_enabled: false
diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb
index 7502a6299aeae9..8c1597ddb8122a 100644
--- a/config/initializers/load_balancing.rb
+++ b/config/initializers/load_balancing.rb
@@ -13,6 +13,20 @@
 
       Gitlab::Database::LoadBalancing.configure_proxy
 
+      if ::Feature.enabled?(:load_balancer_for_sidekiq)
+        Sidekiq.configure_server do |config|
+          config.server_middleware do |chain|
+            chain.add(Gitlab::Database::LoadBalancing::SidekiqServerMiddleware)
+          end
+        end
+
+        Sidekiq.configure_client do |config|
+          config.client_middleware do |chain|
+            chain.add(Gitlab::Database::LoadBalancing::SidekiqClientMiddleware)
+          end
+        end
+      end
+
       # This needs to be executed after fork of clustered processes
       Gitlab::Cluster::LifecycleEvents.on_worker_start do
         # Service discovery must be started after configuring the proxy, as service
diff --git a/ee/lib/gitlab/database/load_balancing.rb b/ee/lib/gitlab/database/load_balancing.rb
index 1ea5632b93e724..8ba681be3559f6 100644
--- a/ee/lib/gitlab/database/load_balancing.rb
+++ b/ee/lib/gitlab/database/load_balancing.rb
@@ -84,12 +84,16 @@ def self.pool_size
 
       # Returns true if load balancing is to be enabled.
       def self.enable?
-        return false if Gitlab::Runtime.rake? || Gitlab::Runtime.sidekiq?
+        return false if program_name == 'rake' || disabled_for_sidekiq?
         return false unless self.configured?
 
         true
       end
 
+      def self.disabled_for_sidekiq?
+         Gitlab::Runtime.sidekiq? && ::Feature.disabled?(:load_balancer_for_sidekiq)
+      end
+
       # Returns true if load balancing has been configured. Since
       # Sidekiq does not currently use load balancing, we
       # may want Web application servers to detect replication lag by
diff --git a/ee/lib/gitlab/database/load_balancing/load_balancer.rb b/ee/lib/gitlab/database/load_balancing/load_balancer.rb
index 85543e46b241ae..f4c96fe349f209 100644
--- a/ee/lib/gitlab/database/load_balancing/load_balancer.rb
+++ b/ee/lib/gitlab/database/load_balancing/load_balancer.rb
@@ -146,7 +146,32 @@ def primary_write_location
         # Returns true if all hosts have caught up to the given transaction
         # write location.
         def all_caught_up?(location)
-          @host_list.hosts.all? { |host| host.caught_up?(location) }
+          if logging_for_all_caught_up_method?
+            all_caught_up = true
+            caught_up_count = 0
+            @host_list.hosts.each do |host|
+              if host.caught_up?(location)
+                caught_up_count += 1
+              else
+                all_caught_up = false
+              end
+            end
+
+            LoadBalancing::Logger.warn(
+              event: :hosts_not_caught_up,
+              message: 'Not all Hosts have caught up to the given transaction write location.',
+              caught_up_count: caught_up_count,
+              host_list_length: @host_list.length
+            ) unless all_caught_up
+
+            all_caught_up
+          else
+            @host_list.hosts.all? { |host| host.caught_up?(location) }
+          end
+        end
+
+        def  logging_for_all_caught_up_method?
+          @logging_for_all_caught_up_method ||= ::Feature.enabled?(:logging_for_all_caught_up_method)
         end
 
         # Yields a block, retrying it upon error using an exponential backoff.
diff --git a/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb b/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
new file mode 100644
index 00000000000000..1a923675b379f5
--- /dev/null
+++ b/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module LoadBalancing
+      class SidekiqClientMiddleware
+        def call(worker_class, job, _queue, _redis_pool)
+          mark_primary_write_location(worker_class, job)
+
+          yield
+        end
+
+        def mark_primary_write_location(worker_class, job)
+          return unless LoadBalancing.enable? || !worker_class.always? || Session.current.performed_write?
+
+          Sticking.with_primary_write_location do |location|
+            job[:primary_write_location] = location
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb b/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
new file mode 100644
index 00000000000000..7dc3e611608e7c
--- /dev/null
+++ b/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module LoadBalancing
+      class SidekiqServerMiddleware
+
+        SecondariesStaleError = Class.new(StandardError)
+
+        def call(worker, job, _queue)
+          clear
+
+          if worker.class.always?
+            Session.current.use_primary!
+          else
+            retry_count = job[:retry_count].present? ? job[:retry_count].to_i + 1 : 0
+            location = job[:primary_write_location]
+            stick_or_delay_if_necessary(worker.class.sticky?, location, retry_count)
+          end
+
+          yield
+        ensure
+          clear
+        end
+
+        private
+
+        def clear
+          load_balancer.release_host
+          Session.clear_session
+        end
+
+        def stick_or_delay_if_necessary(sticky, location, retry_count)
+          unless all_caught_up?(location)
+            if sticky || retry_count > 1
+              Session.current.use_primary!
+            else
+              raise SecondariesStaleError.new("Replicas haven't caught up to the given transaction")
+            end
+          end
+        end
+
+        def load_balancer
+          LoadBalancing.proxy.load_balancer
+        end
+
+        def all_caught_up?(location)
+          load_balancer.all_caught_up?(location)
+        end
+      end
+    end
+  end
+end
-- 
GitLab


From c2dc69f4a67bb48d4f4ff07c0fc82f9ab4a17aef Mon Sep 17 00:00:00 2001
From: nmilojevic1 <nmilojevic@gitlab.com>
Date: Fri, 5 Mar 2021 17:43:28 +0100
Subject: [PATCH 2/9] Fix feature flags

---
 config/initializers/load_balancing.rb                  |  2 +-
 ee/lib/gitlab/database/load_balancing.rb               | 10 +++++++++-
 ee/lib/gitlab/database/load_balancing/load_balancer.rb |  5 ++++-
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb
index 8c1597ddb8122a..7a75f927b58a65 100644
--- a/config/initializers/load_balancing.rb
+++ b/config/initializers/load_balancing.rb
@@ -13,7 +13,7 @@
 
       Gitlab::Database::LoadBalancing.configure_proxy
 
-      if ::Feature.enabled?(:load_balancer_for_sidekiq)
+      if Gitlab::Database::LoadBalancing.load_balancing_for_sidekiq?
         Sidekiq.configure_server do |config|
           config.server_middleware do |chain|
             chain.add(Gitlab::Database::LoadBalancing::SidekiqServerMiddleware)
diff --git a/ee/lib/gitlab/database/load_balancing.rb b/ee/lib/gitlab/database/load_balancing.rb
index 8ba681be3559f6..9ccacea9ffd8ee 100644
--- a/ee/lib/gitlab/database/load_balancing.rb
+++ b/ee/lib/gitlab/database/load_balancing.rb
@@ -91,7 +91,15 @@ def self.enable?
       end
 
       def self.disabled_for_sidekiq?
-         Gitlab::Runtime.sidekiq? && ::Feature.disabled?(:load_balancer_for_sidekiq)
+         Gitlab::Runtime.sidekiq? && !load_balancing_for_sidekiq?
+      end
+
+
+      def self.load_balancing_for_sidekiq?
+        return @load_balancing_for_sidekiq if defined?(@load_balancing_for_sidekiq)
+
+        @load_balancing_for_sidekiq = false
+        @load_balancing_for_sidekiq = ::Feature.enabled?(:load_balancer_for_sidekiq)
       end
 
       # Returns true if load balancing has been configured. Since
diff --git a/ee/lib/gitlab/database/load_balancing/load_balancer.rb b/ee/lib/gitlab/database/load_balancing/load_balancer.rb
index f4c96fe349f209..773a52f3dbeae8 100644
--- a/ee/lib/gitlab/database/load_balancing/load_balancer.rb
+++ b/ee/lib/gitlab/database/load_balancing/load_balancer.rb
@@ -171,7 +171,10 @@ def all_caught_up?(location)
         end
 
         def  logging_for_all_caught_up_method?
-          @logging_for_all_caught_up_method ||= ::Feature.enabled?(:logging_for_all_caught_up_method)
+          return @logging_for_all_caught_up_method if defined?(@logging_for_all_caught_up_method)
+
+          @logging_for_all_caught_up_method = false
+          @logging_for_all_caught_up_method = ::Feature.enabled?(:logging_for_all_caught_up_method)
         end
 
         # Yields a block, retrying it upon error using an exponential backoff.
-- 
GitLab


From 9ba617dd680fed966fb7afc5b2f6105d4c362836 Mon Sep 17 00:00:00 2001
From: nmilojevic1 <nmilojevic@gitlab.com>
Date: Mon, 8 Mar 2021 14:54:45 +0100
Subject: [PATCH 3/9] Fix lint errors

- Improve logging when hosts are not caught up
---
 app/workers/concerns/worker_attributes.rb     | 13 +++--------
 ee/lib/gitlab/database/load_balancing.rb      |  3 +--
 .../database/load_balancing/load_balancer.rb  | 18 +++++++++------
 .../sidekiq_server_middleware.rb              | 23 +++++++++++++++----
 4 files changed, 33 insertions(+), 24 deletions(-)

diff --git a/app/workers/concerns/worker_attributes.rb b/app/workers/concerns/worker_attributes.rb
index 13357c37e07190..18a527e361abcc 100644
--- a/app/workers/concerns/worker_attributes.rb
+++ b/app/workers/concerns/worker_attributes.rb
@@ -11,7 +11,7 @@ module WorkerAttributes
   # Urgencies that workers can declare through the `urgencies` attribute
   VALID_URGENCIES = [:high, :low, :throttled].freeze
 
-  VALID_DATA_CONSISTENCIES = [:always, :sticky ,:delayed].freeze
+  VALID_DATA_CONSISTENCIES = [:always, :sticky, :delayed].freeze
 
   NAMESPACE_WEIGHTS = {
     auto_devops: 2,
@@ -71,23 +71,16 @@ def get_urgency
       class_attributes[:urgency] || :low
     end
 
-    def data_consistency (data_consistency)
+    def data_consistency(data_consistency)
       raise "Invalid data consistency: #{data_consistency}" unless VALID_DATA_CONSISTENCIES.include?(data_consistency)
 
       class_attributes[:data_consistency] = data_consistency
     end
 
-    VALID_DATA_CONSISTENCIES.each do |data_consistency|
-      client_method_name = "#{data_consistency}?".to_sym
-
-      define_singleton_method(client_method_name) do
-        class_attributes[:data_consistency] == data_consistency
-      end
-    end
-
     def get_data_consistency
       class_attributes[:data_consistency] || :always
     end
+
     # Set this attribute on a job when it will call to services outside of the
     # application, such as 3rd party applications, other k8s clusters etc See
     # doc/development/sidekiq_style_guide.md#jobs-with-external-dependencies for
diff --git a/ee/lib/gitlab/database/load_balancing.rb b/ee/lib/gitlab/database/load_balancing.rb
index 9ccacea9ffd8ee..c95c0bd439a6b5 100644
--- a/ee/lib/gitlab/database/load_balancing.rb
+++ b/ee/lib/gitlab/database/load_balancing.rb
@@ -91,10 +91,9 @@ def self.enable?
       end
 
       def self.disabled_for_sidekiq?
-         Gitlab::Runtime.sidekiq? && !load_balancing_for_sidekiq?
+        Gitlab::Runtime.sidekiq? && !load_balancing_for_sidekiq?
       end
 
-
       def self.load_balancing_for_sidekiq?
         return @load_balancing_for_sidekiq if defined?(@load_balancing_for_sidekiq)
 
diff --git a/ee/lib/gitlab/database/load_balancing/load_balancer.rb b/ee/lib/gitlab/database/load_balancing/load_balancer.rb
index 773a52f3dbeae8..76df189ee32970 100644
--- a/ee/lib/gitlab/database/load_balancing/load_balancer.rb
+++ b/ee/lib/gitlab/database/load_balancing/load_balancer.rb
@@ -157,12 +157,16 @@ def all_caught_up?(location)
               end
             end
 
-            LoadBalancing::Logger.warn(
-              event: :hosts_not_caught_up,
-              message: 'Not all Hosts have caught up to the given transaction write location.',
-              caught_up_count: caught_up_count,
-              host_list_length: @host_list.length
-            ) unless all_caught_up
+            unless all_caught_up
+              LoadBalancing::Logger.warn(
+                event: :hosts_not_caught_up,
+                message: 'Not all Hosts have caught up to the given transaction write location.',
+                caught_up_count: caught_up_count,
+                sidekiq: Gitlab::Runtime.sidekiq?,
+                web_server: Gitlab::Runtime.web_server?,
+                host_list_length: @host_list.length
+              )
+            end
 
             all_caught_up
           else
@@ -170,7 +174,7 @@ def all_caught_up?(location)
           end
         end
 
-        def  logging_for_all_caught_up_method?
+        def logging_for_all_caught_up_method?
           return @logging_for_all_caught_up_method if defined?(@logging_for_all_caught_up_method)
 
           @logging_for_all_caught_up_method = false
diff --git a/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb b/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
index 7dc3e611608e7c..145fa382eb3e54 100644
--- a/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
+++ b/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
@@ -10,12 +10,12 @@ class SidekiqServerMiddleware
         def call(worker, job, _queue)
           clear
 
-          if worker.class.always?
+          if worker.class.data_consistency == :always
             Session.current.use_primary!
           else
             retry_count = job[:retry_count].present? ? job[:retry_count].to_i + 1 : 0
             location = job[:primary_write_location]
-            stick_or_delay_if_necessary(worker.class.sticky?, location, retry_count)
+            stick_or_delay_if_necessary(worker.class.data_consistency, location, retry_count)
           end
 
           yield
@@ -30,16 +30,29 @@ def clear
           Session.clear_session
         end
 
-        def stick_or_delay_if_necessary(sticky, location, retry_count)
+        def stick_or_delay_if_necessary(data_consistency, location, retry_count)
           unless all_caught_up?(location)
-            if sticky || retry_count > 1
+            if data_consistency == :sticky || retry_count > 1
               Session.current.use_primary!
+              log_warning(data_consistency, retry_count, "Using primary instead.")
             else
-              raise SecondariesStaleError.new("Replicas haven't caught up to the given transaction")
+              log_warning(data_consistency, retry_count, "Retrying the worker.")
+              raise SecondariesStaleError.new("Not all Hosts have caught up to the given transaction write location.")
             end
           end
         end
 
+        def log_warning(data_consistency, retry_count, message)
+          full_message = "Not all Hosts have caught up to the given transaction write location. #{message}"
+
+          LoadBalancing::Logger.warn(
+            event: :hosts_not_caught_up_for_sidekiq,
+            message: full_message,
+            worker_data_consistency: data_consistency,
+            retry_count: retry_count
+          )
+        end
+
         def load_balancer
           LoadBalancing.proxy.load_balancer
         end
-- 
GitLab


From e96f91343db77715a931684490bfc381c6ef2a52 Mon Sep 17 00:00:00 2001
From: nmilojevic1 <nmilojevic@gitlab.com>
Date: Mon, 8 Mar 2021 17:26:17 +0100
Subject: [PATCH 4/9] Address reviewer comments

- Add specs for sidekiq middleware
- Add specs for worker_attributes
- Add env variable to control load_balancing for sidekiq
- Add support for feature flag controling data_consistency
- Raise exception for idempotent jobs
- Pass worker_data_consistency to the structured log
---
 app/workers/concerns/worker_attributes.rb     |  12 +-
 .../development/load_balancer_for_sidekiq.yml |   8 -
 .../logging_for_all_caught_up_method.yml      |   8 -
 config/initializers/load_balancing.rb         |  16 +-
 ee/lib/gitlab/database/load_balancing.rb      |  14 +-
 .../database/load_balancing/load_balancer.rb  |  34 +----
 .../sidekiq_client_middleware.rb              |  21 ++-
 .../sidekiq_server_middleware.rb              |  48 +++---
 .../sidekiq_client_middleware_spec.rb         | 113 ++++++++++++++
 .../sidekiq_server_middleware_spec.rb         | 141 ++++++++++++++++++
 .../gitlab/database/load_balancing_spec.rb    |  13 +-
 lib/gitlab/instrumentation_helper.rb          |   1 +
 .../lib/gitlab/instrumentation_helper_spec.rb |   1 +
 .../concerns/worker_attributes_spec.rb        |  74 +++++++++
 14 files changed, 397 insertions(+), 107 deletions(-)
 delete mode 100644 config/feature_flags/development/load_balancer_for_sidekiq.yml
 delete mode 100644 config/feature_flags/development/logging_for_all_caught_up_method.yml
 create mode 100644 ee/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
 create mode 100644 ee/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
 create mode 100644 spec/workers/concerns/worker_attributes_spec.rb

diff --git a/app/workers/concerns/worker_attributes.rb b/app/workers/concerns/worker_attributes.rb
index 18a527e361abcc..4d9193026f7c02 100644
--- a/app/workers/concerns/worker_attributes.rb
+++ b/app/workers/concerns/worker_attributes.rb
@@ -71,9 +71,11 @@ def get_urgency
       class_attributes[:urgency] || :low
     end
 
-    def data_consistency(data_consistency)
+    def data_consistency(data_consistency, feature_flag: nil)
       raise "Invalid data consistency: #{data_consistency}" unless VALID_DATA_CONSISTENCIES.include?(data_consistency)
+      raise "Class can't be marked as idempotent if data_consistency is not set to :always" if idempotent? && data_consistency != :always
 
+      class_attributes[:data_consistency_feature_flag] = feature_flag if feature_flag
       class_attributes[:data_consistency] = data_consistency
     end
 
@@ -81,6 +83,12 @@ def get_data_consistency
       class_attributes[:data_consistency] || :always
     end
 
+    def get_data_consistency_feature_flag_enabled?
+      return true unless class_attributes[:data_consistency_feature_flag]
+
+      Feature.enabled?(class_attributes[:data_consistency_feature_flag], default_enabled: :yaml)
+    end
+
     # Set this attribute on a job when it will call to services outside of the
     # application, such as 3rd party applications, other k8s clusters etc See
     # doc/development/sidekiq_style_guide.md#jobs-with-external-dependencies for
@@ -107,6 +115,8 @@ def get_worker_resource_boundary
     end
 
     def idempotent!
+      raise "Class can't be marked as idempotent if data_consistency is not set to :always" unless get_data_consistency == :always
+
       class_attributes[:idempotent] = true
     end
 
diff --git a/config/feature_flags/development/load_balancer_for_sidekiq.yml b/config/feature_flags/development/load_balancer_for_sidekiq.yml
deleted file mode 100644
index 047d76d53df22c..00000000000000
--- a/config/feature_flags/development/load_balancer_for_sidekiq.yml
+++ /dev/null
@@ -1,8 +0,0 @@
----
-name: load_balancer_for_sidekiq
-introduced_by_url:
-rollout_issue_url:
-milestone: '13.10'
-type: development
-group: group::memory
-default_enabled: false
diff --git a/config/feature_flags/development/logging_for_all_caught_up_method.yml b/config/feature_flags/development/logging_for_all_caught_up_method.yml
deleted file mode 100644
index 53a30a3b921b82..00000000000000
--- a/config/feature_flags/development/logging_for_all_caught_up_method.yml
+++ /dev/null
@@ -1,8 +0,0 @@
----
-name: logging_for_all_caught_up_method
-introduced_by_url:
-rollout_issue_url:
-milestone: '13.10'
-type: development
-group: group::memory
-default_enabled: false
diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb
index 7a75f927b58a65..cafefc47a6d96f 100644
--- a/config/initializers/load_balancing.rb
+++ b/config/initializers/load_balancing.rb
@@ -13,17 +13,15 @@
 
       Gitlab::Database::LoadBalancing.configure_proxy
 
-      if Gitlab::Database::LoadBalancing.load_balancing_for_sidekiq?
-        Sidekiq.configure_server do |config|
-          config.server_middleware do |chain|
-            chain.add(Gitlab::Database::LoadBalancing::SidekiqServerMiddleware)
-          end
+      Sidekiq.configure_server do |config|
+        config.server_middleware do |chain|
+          chain.add(Gitlab::Database::LoadBalancing::SidekiqServerMiddleware)
         end
+      end
 
-        Sidekiq.configure_client do |config|
-          config.client_middleware do |chain|
-            chain.add(Gitlab::Database::LoadBalancing::SidekiqClientMiddleware)
-          end
+      Sidekiq.configure_client do |config|
+        config.client_middleware do |chain|
+          chain.add(Gitlab::Database::LoadBalancing::SidekiqClientMiddleware)
         end
       end
 
diff --git a/ee/lib/gitlab/database/load_balancing.rb b/ee/lib/gitlab/database/load_balancing.rb
index c95c0bd439a6b5..7fad432daae27a 100644
--- a/ee/lib/gitlab/database/load_balancing.rb
+++ b/ee/lib/gitlab/database/load_balancing.rb
@@ -84,21 +84,14 @@ def self.pool_size
 
       # Returns true if load balancing is to be enabled.
       def self.enable?
-        return false if program_name == 'rake' || disabled_for_sidekiq?
+        return false if Gitlab::Runtime.rake? || disabled_for_sidekiq
         return false unless self.configured?
 
         true
       end
 
-      def self.disabled_for_sidekiq?
-        Gitlab::Runtime.sidekiq? && !load_balancing_for_sidekiq?
-      end
-
-      def self.load_balancing_for_sidekiq?
-        return @load_balancing_for_sidekiq if defined?(@load_balancing_for_sidekiq)
-
-        @load_balancing_for_sidekiq = false
-        @load_balancing_for_sidekiq = ::Feature.enabled?(:load_balancer_for_sidekiq)
+      def self.disabled_for_sidekiq
+        Gitlab::Runtime.sidekiq? && ENV['ENABLE_LOAD_BALANCING_FOR_SIDEKIQ'] != 'true'
       end
 
       # Returns true if load balancing has been configured. Since
@@ -124,7 +117,6 @@ def self.feature_available?
         #   -> Set @feature_available  to true
         #   -> return true
         # - Second call: return @feature_available right away
-
         return @feature_available if defined?(@feature_available)
 
         @feature_available = false
diff --git a/ee/lib/gitlab/database/load_balancing/load_balancer.rb b/ee/lib/gitlab/database/load_balancing/load_balancer.rb
index 76df189ee32970..85543e46b241ae 100644
--- a/ee/lib/gitlab/database/load_balancing/load_balancer.rb
+++ b/ee/lib/gitlab/database/load_balancing/load_balancer.rb
@@ -146,39 +146,7 @@ def primary_write_location
         # Returns true if all hosts have caught up to the given transaction
         # write location.
         def all_caught_up?(location)
-          if logging_for_all_caught_up_method?
-            all_caught_up = true
-            caught_up_count = 0
-            @host_list.hosts.each do |host|
-              if host.caught_up?(location)
-                caught_up_count += 1
-              else
-                all_caught_up = false
-              end
-            end
-
-            unless all_caught_up
-              LoadBalancing::Logger.warn(
-                event: :hosts_not_caught_up,
-                message: 'Not all Hosts have caught up to the given transaction write location.',
-                caught_up_count: caught_up_count,
-                sidekiq: Gitlab::Runtime.sidekiq?,
-                web_server: Gitlab::Runtime.web_server?,
-                host_list_length: @host_list.length
-              )
-            end
-
-            all_caught_up
-          else
-            @host_list.hosts.all? { |host| host.caught_up?(location) }
-          end
-        end
-
-        def logging_for_all_caught_up_method?
-          return @logging_for_all_caught_up_method if defined?(@logging_for_all_caught_up_method)
-
-          @logging_for_all_caught_up_method = false
-          @logging_for_all_caught_up_method = ::Feature.enabled?(:logging_for_all_caught_up_method)
+          @host_list.hosts.all? { |host| host.caught_up?(location) }
         end
 
         # Yields a block, retrying it upon error using an exponential backoff.
diff --git a/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb b/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
index 1a923675b379f5..0f6709207f8b42 100644
--- a/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
+++ b/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
@@ -5,18 +5,29 @@ module Database
     module LoadBalancing
       class SidekiqClientMiddleware
         def call(worker_class, job, _queue, _redis_pool)
-          mark_primary_write_location(worker_class, job)
+          mark_data_consistency_location(worker_class, job)
 
           yield
         end
 
-        def mark_primary_write_location(worker_class, job)
-          return unless LoadBalancing.enable? || !worker_class.always? || Session.current.performed_write?
+        private
 
-          Sticking.with_primary_write_location do |location|
-            job[:primary_write_location] = location
+        def mark_data_consistency_location(worker_class, job)
+          worker_class = worker_class.to_s.safe_constantize
+          return if worker_class.get_data_consistency == :always
+
+          return unless worker_class.get_data_consistency_feature_flag_enabled?
+
+          if Session.current.performed_write?
+            job['database_write_location'] = load_balancer.primary_write_location
+          else
+            job['database_replica_location'] = true
           end
         end
+
+        def load_balancer
+          LoadBalancing.proxy.load_balancer
+        end
       end
     end
   end
diff --git a/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb b/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
index 145fa382eb3e54..772b06cdb895e5 100644
--- a/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
+++ b/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
@@ -4,18 +4,11 @@ module Gitlab
   module Database
     module LoadBalancing
       class SidekiqServerMiddleware
-
-        SecondariesStaleError = Class.new(StandardError)
+        JobReplicaNotUpToDate = Class.new(StandardError)
 
         def call(worker, job, _queue)
-          clear
-
-          if worker.class.data_consistency == :always
+          if requires_primary?(worker.class, job)
             Session.current.use_primary!
-          else
-            retry_count = job[:retry_count].present? ? job[:retry_count].to_i + 1 : 0
-            location = job[:primary_write_location]
-            stick_or_delay_if_necessary(worker.class.data_consistency, location, retry_count)
           end
 
           yield
@@ -30,35 +23,28 @@ def clear
           Session.clear_session
         end
 
-        def stick_or_delay_if_necessary(data_consistency, location, retry_count)
-          unless all_caught_up?(location)
-            if data_consistency == :sticky || retry_count > 1
-              Session.current.use_primary!
-              log_warning(data_consistency, retry_count, "Using primary instead.")
-            else
-              log_warning(data_consistency, retry_count, "Retrying the worker.")
-              raise SecondariesStaleError.new("Not all Hosts have caught up to the given transaction write location.")
-            end
-          end
-        end
-
-        def log_warning(data_consistency, retry_count, message)
-          full_message = "Not all Hosts have caught up to the given transaction write location. #{message}"
+        def requires_primary?(worker_class, job)
+          return true if worker_class.get_data_consistency == :always
+          return true unless worker_class.get_data_consistency_feature_flag_enabled?
 
-          LoadBalancing::Logger.warn(
-            event: :hosts_not_caught_up_for_sidekiq,
-            message: full_message,
-            worker_data_consistency: data_consistency,
-            retry_count: retry_count
-          )
+          if job['database_replica_location'] || replica_caught_up?(job['database_write_location'] )
+            false
+          elsif worker_class.get_data_consistency == :delayed && job['retry_count'].to_i == 0
+            raise JobReplicaNotUpToDate, "Sidekiq job #{worker_class} JID-#{job['jid']} couldn't use the replica."\
+               "  Replica was not up to date."
+          else
+            true
+          end
         end
 
         def load_balancer
           LoadBalancing.proxy.load_balancer
         end
 
-        def all_caught_up?(location)
-          load_balancer.all_caught_up?(location)
+        def replica_caught_up?(location)
+          return true unless location
+
+          load_balancer.host.caught_up?(location)
         end
       end
     end
diff --git a/ee/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb b/ee/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
new file mode 100644
index 00000000000000..b00cfa61b11cf0
--- /dev/null
+++ b/ee/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
@@ -0,0 +1,113 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
+  let(:middleware) { described_class.new }
+
+  after do
+    Gitlab::Database::LoadBalancing::Session.clear_session
+  end
+
+  describe '#call' do
+    shared_context 'data consistency worker class' do |data_consistency, feature_flag|
+      let(:worker_class) do
+        Class.new do
+          def self.name
+            'TestDataConsistencyWorker'
+          end
+
+          include ApplicationWorker
+
+          data_consistency data_consistency, feature_flag: feature_flag
+
+          def perform(*args)
+          end
+        end
+      end
+
+      before do
+        stub_const('TestDataConsistencyWorker', worker_class)
+      end
+    end
+
+    shared_examples_for 'mark database_replica_location' do
+      it 'passes database_replica_location' do
+        expect(middleware).not_to receive(:load_balancer)
+
+        middleware.call(worker_class, job, double(:queue), redis_pool) { 10 }
+
+        expect(job['database_replica_location']).to be_truthy
+      end
+    end
+
+    shared_examples_for 'does not pass database locations' do
+      it 'does not pass database locations', :aggregate_failures do
+        middleware.call(worker_class, job, double(:queue), redis_pool) { 10 }
+
+        expect(job['database_replica_location']).to be_nil
+        expect(job['database_write_location']).to be_nil
+      end
+    end
+
+    shared_examples_for 'mark data consistency location' do |data_consistency|
+      include_context 'data consistency worker class', data_consistency, :load_balancing_for_test_data_consistency_worker
+
+      let(:location) { '0/D525E3A8' }
+
+      context 'when feature flag load_balancing_for_sidekiq is disabled' do
+        before do
+          stub_feature_flags(load_balancing_for_test_data_consistency_worker: false)
+        end
+
+        include_examples 'does not pass database locations'
+      end
+
+      context 'when write was not performed' do
+        before do
+          allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:performed_write?).and_return(false)
+        end
+
+        include_examples 'mark database_replica_location'
+      end
+
+      context 'when write was performed' do
+        before do
+          allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:performed_write?).and_return(true)
+        end
+
+        it 'passes primary write location', :aggregate_failures do
+          expect(middleware).to receive_message_chain(:load_balancer, :primary_write_location).and_return(location)
+
+          middleware.call(worker_class, job, double(:queue), redis_pool) { 10 }
+
+          expect(job['database_write_location']).to eq(location)
+        end
+      end
+    end
+
+    let(:queue) { 'default' }
+    let(:redis_pool) { Sidekiq.redis_pool }
+    let(:worker_class) { 'TestDataConsistencyWorker' }
+    let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e" } }
+
+    before do
+      skip_feature_flags_yaml_validation
+      skip_default_enabled_yaml_check
+    end
+
+    context 'when worker data consistency is :always' do
+      include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker
+
+      include_examples 'does not pass database locations'
+    end
+
+    context 'when worker data consistency is :delayed' do
+      include_examples  'mark data consistency location', :delayed
+    end
+
+    context 'when worker data consistency is :sticky' do
+      include_examples  'mark data consistency location', :sticky
+    end
+  end
+end
diff --git a/ee/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb b/ee/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
new file mode 100644
index 00000000000000..c929c8f1de7792
--- /dev/null
+++ b/ee/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
@@ -0,0 +1,141 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
+  let(:middleware) { described_class.new }
+
+  after do
+    Gitlab::Database::LoadBalancing::Session.clear_session
+  end
+
+  describe '#call' do
+    shared_context 'data consistency worker class' do |data_consistency, feature_flag|
+      let(:worker_class) do
+        Class.new do
+          def self.name
+            'TestDataConsistencyWorker'
+          end
+
+          include ApplicationWorker
+
+          data_consistency data_consistency, feature_flag: feature_flag
+
+          def perform(*args)
+          end
+        end
+      end
+
+      before do
+        stub_const('TestDataConsistencyWorker', worker_class)
+      end
+    end
+
+    shared_examples_for 'stick to the primary' do
+      it 'sticks to the primary' do
+        middleware.call(worker, job, double(:queue)) do
+          expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).to be_truthy
+        end
+      end
+    end
+
+    shared_examples_for 'sticks based on data consistency' do |data_consistency|
+      include_context 'data consistency worker class', data_consistency, :load_balancing_for_test_data_consistency_worker
+
+      context 'when load_balancing_for_test_data_consistency_worker is disabled' do
+        before do
+          stub_feature_flags(load_balancing_for_test_data_consistency_worker: false)
+        end
+
+        include_examples 'stick to the primary'
+      end
+
+      context 'database replica location is set' do
+        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'database_replica_location' => 'true' } }
+
+        it 'do not stick to the primary' do
+          middleware.call(worker, job, double(:queue)) do
+            expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
+          end
+        end
+      end
+
+      context 'write was not performed' do
+        let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e' } }
+
+        it 'do not stick to the primary' do
+          middleware.call(worker, job, double(:queue)) do
+            expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
+          end
+        end
+      end
+
+      context 'replica is up to date' do
+        before do
+          allow(middleware).to receive(:replica_caught_up?).and_return(true)
+        end
+
+        it 'do not stick to the primary' do
+          middleware.call(worker, job, double(:queue)) do
+            expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
+          end
+        end
+      end
+    end
+
+    let(:queue) { 'default' }
+    let(:redis_pool) { Sidekiq.redis_pool }
+    let(:worker) { TestDataConsistencyWorker.new }
+    let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'primary_write_location' => '0/D525E3A8' } }
+    let(:block) { 10 }
+
+    before do
+      skip_feature_flags_yaml_validation
+      skip_default_enabled_yaml_check
+      allow(middleware).to receive(:clear)
+      allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:performed_write?).and_return(true)
+    end
+
+    context 'when worker data consistency is :always' do
+      include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker
+
+      include_examples 'stick to the primary'
+    end
+
+    context 'when worker data consistency is :delayed' do
+      include_examples 'sticks based on data consistency', :delayed
+
+      context 'when replica is not up to date' do
+        before do
+          allow(middleware).to receive(:replica_caught_up?).and_return(false)
+        end
+
+        context 'when job is retried once' do
+          it 'raise an error and retries' do
+            expect { middleware.call(worker, job, double(:queue)) { block } }.to raise_error(Gitlab::Database::LoadBalancing::SidekiqServerMiddleware::JobReplicaNotUpToDate)
+          end
+        end
+
+        context 'when job is retried more then once' do
+          before do
+            job['retry_count'] = 1
+          end
+
+          include_examples 'stick to the primary'
+        end
+      end
+    end
+
+    context 'when worker data consistency is :sticky' do
+      include_examples 'sticks based on data consistency', :sticky
+
+      context 'when replica is not up to date' do
+        before do
+          allow(middleware).to receive(:replica_caught_up?).and_return(false)
+        end
+
+        include_examples 'stick to the primary'
+      end
+    end
+  end
+end
diff --git a/ee/spec/lib/gitlab/database/load_balancing_spec.rb b/ee/spec/lib/gitlab/database/load_balancing_spec.rb
index 94d91e92e00660..e86bb88580e77c 100644
--- a/ee/spec/lib/gitlab/database/load_balancing_spec.rb
+++ b/ee/spec/lib/gitlab/database/load_balancing_spec.rb
@@ -142,7 +142,6 @@
     end
 
     it 'returns false when Sidekiq is being used' do
-      allow(described_class).to receive(:hosts).and_return(%w(foo))
       allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
 
       expect(described_class.enable?).to eq(false)
@@ -171,6 +170,18 @@
       expect(described_class.enable?).to eq(true)
     end
 
+    context 'when ENABLE_LOAD_BALANCING_FOR_SIDEKIQ environment variable is set' do
+      before do
+        stub_env('ENABLE_LOAD_BALANCING_FOR_SIDEKIQ', 'true')
+      end
+
+      it 'returns true when Sidekiq is being used' do
+        allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
+
+        expect(described_class.enable?).to eq(true)
+      end
+    end
+
     context 'without a license' do
       before do
         License.destroy_all # rubocop: disable Cop/DestroyAll
diff --git a/lib/gitlab/instrumentation_helper.rb b/lib/gitlab/instrumentation_helper.rb
index 61de6b02453e9e..420803a17e7862 100644
--- a/lib/gitlab/instrumentation_helper.rb
+++ b/lib/gitlab/instrumentation_helper.rb
@@ -16,6 +16,7 @@ def keys
         :elasticsearch_calls,
         :elasticsearch_duration_s,
         :elasticsearch_timed_out_count,
+        :worker_data_consistency,
         *::Gitlab::Memory::Instrumentation::KEY_MAPPING.values,
         *::Gitlab::Instrumentation::Redis.known_payload_keys,
         *::Gitlab::Metrics::Subscribers::ActiveRecord::DB_COUNTERS,
diff --git a/spec/lib/gitlab/instrumentation_helper_spec.rb b/spec/lib/gitlab/instrumentation_helper_spec.rb
index a5c9cde4c37e11..bbfdb8f4c2ca70 100644
--- a/spec/lib/gitlab/instrumentation_helper_spec.rb
+++ b/spec/lib/gitlab/instrumentation_helper_spec.rb
@@ -17,6 +17,7 @@
         :elasticsearch_calls,
         :elasticsearch_duration_s,
         :elasticsearch_timed_out_count,
+        :worker_data_consistency,
         :mem_objects,
         :mem_bytes,
         :mem_mallocs,
diff --git a/spec/workers/concerns/worker_attributes_spec.rb b/spec/workers/concerns/worker_attributes_spec.rb
new file mode 100644
index 00000000000000..a654ecbd3e2fc1
--- /dev/null
+++ b/spec/workers/concerns/worker_attributes_spec.rb
@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe WorkerAttributes do
+  let(:worker) do
+    Class.new do
+      def self.name
+        "TestWorker"
+      end
+
+      include ApplicationWorker
+    end
+  end
+
+  describe '.data_consistency' do
+    context 'with valid data_consistency' do
+      it 'returns correct data_consistency' do
+        worker.data_consistency(:sticky)
+
+        expect(worker.get_data_consistency).to eq(:sticky)
+      end
+    end
+
+    context 'when data_consistency is not provided' do
+      it 'defaults to :always' do
+        expect(worker.get_data_consistency).to eq(:always)
+      end
+    end
+
+    context 'with invalid data_consistency' do
+      it 'raise exception' do
+        expect { worker.data_consistency(:invalid) }
+          .to raise_error('Invalid data consistency: invalid')
+      end
+    end
+
+    context 'when job is idempotent' do
+      context 'when data_consistency is not :always' do
+        it 'raise exception' do
+          worker.idempotent!
+
+          expect { worker.data_consistency(:sticky) }
+            .to raise_error("Class can't be marked as idempotent if data_consistency is not set to :always")
+        end
+      end
+
+      context 'when feature_flag is provided' do
+        before do
+          stub_feature_flags(test_feature_flag: false)
+          skip_feature_flags_yaml_validation
+          skip_default_enabled_yaml_check
+        end
+
+        it 'returns correct feature flag value' do
+          worker.data_consistency(:sticky, feature_flag: :test_feature_flag)
+
+          expect(worker.get_data_consistency_feature_flag_enabled?).not_to be_truthy
+        end
+      end
+    end
+  end
+
+  describe '.idempotent!' do
+    context 'when data consistency is not :always' do
+      it 'raise exception' do
+        worker.data_consistency(:sticky)
+
+        expect { worker.idempotent! }
+          .to raise_error("Class can't be marked as idempotent if data_consistency is not set to :always")
+      end
+    end
+  end
+end
-- 
GitLab


From f7cf43755de3b13272dc27f0bd9a580883720646 Mon Sep 17 00:00:00 2001
From: nmilojevic1 <nmilojevic@gitlab.com>
Date: Thu, 18 Mar 2021 14:58:21 +0100
Subject: [PATCH 5/9] Add comment for deduplication

---
 app/workers/concerns/worker_attributes.rb | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/app/workers/concerns/worker_attributes.rb b/app/workers/concerns/worker_attributes.rb
index 4d9193026f7c02..d2cb4eef910afe 100644
--- a/app/workers/concerns/worker_attributes.rb
+++ b/app/workers/concerns/worker_attributes.rb
@@ -73,6 +73,9 @@ def get_urgency
 
     def data_consistency(data_consistency, feature_flag: nil)
       raise "Invalid data consistency: #{data_consistency}" unless VALID_DATA_CONSISTENCIES.include?(data_consistency)
+      # Since the deduplication should always take into account the latest binary replication pointer into account,
+      # not the first one, the deduplication will not work with sticky or delayed.
+      # Follow up issue to improve this: https://gitlab.com/gitlab-org/gitlab/-/issues/325291
       raise "Class can't be marked as idempotent if data_consistency is not set to :always" if idempotent? && data_consistency != :always
 
       class_attributes[:data_consistency_feature_flag] = feature_flag if feature_flag
@@ -115,6 +118,9 @@ def get_worker_resource_boundary
     end
 
     def idempotent!
+      # Since the deduplication should always take into account the latest binary replication pointer into account,
+      # not the first one, the deduplication will not work with sticky or delayed.
+      # Follow up issue to improve this: https://gitlab.com/gitlab-org/gitlab/-/issues/325291
       raise "Class can't be marked as idempotent if data_consistency is not set to :always" unless get_data_consistency == :always
 
       class_attributes[:idempotent] = true
-- 
GitLab


From cea30ec82d53040754da2b9615d984aee0bf39d7 Mon Sep 17 00:00:00 2001
From: nmilojevic1 <nmilojevic@gitlab.com>
Date: Wed, 24 Mar 2021 01:07:04 +0100
Subject: [PATCH 6/9] Addressed maintainer comments

- Extract ee sidekiq_middleware
- Raise exception if data_consistency is already set
- Made code more clear
---
 app/workers/concerns/worker_attributes.rb     |  25 ++--
 config/initializers/load_balancing.rb         |  12 --
 ee/lib/ee/gitlab/sidekiq_middleware.rb        |  35 +++++
 ee/lib/gitlab/database/load_balancing.rb      |   7 +-
 .../sidekiq_client_middleware.rb              |   4 +
 .../sidekiq_server_middleware.rb              |   2 +
 .../lib/ee/gitlab/sidekiq_middleware_spec.rb  | 122 ++++++++++++++++++
 lib/gitlab/sidekiq_middleware.rb              |   2 +
 8 files changed, 182 insertions(+), 27 deletions(-)
 create mode 100644 ee/lib/ee/gitlab/sidekiq_middleware.rb
 create mode 100644 ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb

diff --git a/app/workers/concerns/worker_attributes.rb b/app/workers/concerns/worker_attributes.rb
index d2cb4eef910afe..6f99fd089aca4d 100644
--- a/app/workers/concerns/worker_attributes.rb
+++ b/app/workers/concerns/worker_attributes.rb
@@ -72,14 +72,22 @@ def get_urgency
     end
 
     def data_consistency(data_consistency, feature_flag: nil)
-      raise "Invalid data consistency: #{data_consistency}" unless VALID_DATA_CONSISTENCIES.include?(data_consistency)
-      # Since the deduplication should always take into account the latest binary replication pointer into account,
-      # not the first one, the deduplication will not work with sticky or delayed.
-      # Follow up issue to improve this: https://gitlab.com/gitlab-org/gitlab/-/issues/325291
-      raise "Class can't be marked as idempotent if data_consistency is not set to :always" if idempotent? && data_consistency != :always
+      raise ArgumentError, "Invalid data consistency: #{data_consistency}" unless VALID_DATA_CONSISTENCIES.include?(data_consistency)
+      raise ArgumentError, 'Data consistency is already set' if class_attributes[:data_consistency]
 
       class_attributes[:data_consistency_feature_flag] = feature_flag if feature_flag
       class_attributes[:data_consistency] = data_consistency
+
+      validate_worker_attributes!
+    end
+
+    def validate_worker_attributes!
+      # Since the deduplication should always take into account the latest binary replication pointer into account,
+      # not the first one, the deduplication will not work with sticky or delayed.
+      # Follow up issue to improve this: https://gitlab.com/gitlab-org/gitlab/-/issues/325291
+      if idempotent? && get_data_consistency != :always
+        raise ArgumentError, "Class can't be marked as idempotent if data_consistency is not set to :always"
+      end
     end
 
     def get_data_consistency
@@ -118,12 +126,9 @@ def get_worker_resource_boundary
     end
 
     def idempotent!
-      # Since the deduplication should always take into account the latest binary replication pointer into account,
-      # not the first one, the deduplication will not work with sticky or delayed.
-      # Follow up issue to improve this: https://gitlab.com/gitlab-org/gitlab/-/issues/325291
-      raise "Class can't be marked as idempotent if data_consistency is not set to :always" unless get_data_consistency == :always
-
       class_attributes[:idempotent] = true
+
+      validate_worker_attributes!
     end
 
     def idempotent?
diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb
index cafefc47a6d96f..7502a6299aeae9 100644
--- a/config/initializers/load_balancing.rb
+++ b/config/initializers/load_balancing.rb
@@ -13,18 +13,6 @@
 
       Gitlab::Database::LoadBalancing.configure_proxy
 
-      Sidekiq.configure_server do |config|
-        config.server_middleware do |chain|
-          chain.add(Gitlab::Database::LoadBalancing::SidekiqServerMiddleware)
-        end
-      end
-
-      Sidekiq.configure_client do |config|
-        config.client_middleware do |chain|
-          chain.add(Gitlab::Database::LoadBalancing::SidekiqClientMiddleware)
-        end
-      end
-
       # This needs to be executed after fork of clustered processes
       Gitlab::Cluster::LifecycleEvents.on_worker_start do
         # Service discovery must be started after configuring the proxy, as service
diff --git a/ee/lib/ee/gitlab/sidekiq_middleware.rb b/ee/lib/ee/gitlab/sidekiq_middleware.rb
new file mode 100644
index 00000000000000..61c3fc0c2b1471
--- /dev/null
+++ b/ee/lib/ee/gitlab/sidekiq_middleware.rb
@@ -0,0 +1,35 @@
+# 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)
+
+          chain.add ::Gitlab::Database::LoadBalancing::SidekiqServerMiddleware if load_balancing_enabled?
+        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
diff --git a/ee/lib/gitlab/database/load_balancing.rb b/ee/lib/gitlab/database/load_balancing.rb
index 7fad432daae27a..a5b18ac09c5ab5 100644
--- a/ee/lib/gitlab/database/load_balancing.rb
+++ b/ee/lib/gitlab/database/load_balancing.rb
@@ -84,16 +84,13 @@ def self.pool_size
 
       # Returns true if load balancing is to be enabled.
       def self.enable?
-        return false if Gitlab::Runtime.rake? || disabled_for_sidekiq
+        return false if Gitlab::Runtime.rake?
+        return false if Gitlab::Runtime.sidekiq? && !Gitlab::Utils.to_boolean(ENV['ENABLE_LOAD_BALANCING_FOR_SIDEKIQ'], default: false)
         return false unless self.configured?
 
         true
       end
 
-      def self.disabled_for_sidekiq
-        Gitlab::Runtime.sidekiq? && ENV['ENABLE_LOAD_BALANCING_FOR_SIDEKIQ'] != 'true'
-      end
-
       # Returns true if load balancing has been configured. Since
       # Sidekiq does not currently use load balancing, we
       # may want Web application servers to detect replication lag by
diff --git a/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb b/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
index 0f6709207f8b42..7511435b5c97b6 100644
--- a/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
+++ b/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
@@ -21,6 +21,10 @@ def mark_data_consistency_location(worker_class, job)
           if Session.current.performed_write?
             job['database_write_location'] = load_balancer.primary_write_location
           else
+            # It is possible that the current replica has a different write-ahead
+            # replication log location from the sidekiq server replica.
+            # In the follow-up issue https://gitlab.com/gitlab-org/gitlab/-/issues/325519,
+            # we want to pass database replica location as well
             job['database_replica_location'] = true
           end
         end
diff --git a/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb b/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
index 772b06cdb895e5..7ce4cbcfe723a0 100644
--- a/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
+++ b/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
@@ -24,6 +24,8 @@ def clear
         end
 
         def requires_primary?(worker_class, job)
+          job[:worker_data_consistency] = worker_class.get_data_consistency
+
           return true if worker_class.get_data_consistency == :always
           return true unless worker_class.get_data_consistency_feature_flag_enabled?
 
diff --git a/ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb b/ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb
new file mode 100644
index 00000000000000..a82f828f03dd24
--- /dev/null
+++ b/ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb
@@ -0,0 +1,122 @@
+# 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_instance_of(middleware) do |middleware_instance|
+          expect(middleware_instance).to receive(:call).with(*middleware_expected_args).ordered.once.and_call_original
+        end
+      end
+
+      expect { |b| chain.invoke(*worker_args, &b) }.to yield_control.once
+    end
+  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::Database::LoadBalancing::SidekiqServerMiddleware
+      ]
+    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
+    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
+    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
+    end
+
+    context "when load balancing is enabled" do
+      it_behaves_like "a middleware chain", true
+    end
+  end
+end
diff --git a/lib/gitlab/sidekiq_middleware.rb b/lib/gitlab/sidekiq_middleware.rb
index a2696e1707849e..563a105484de06 100644
--- a/lib/gitlab/sidekiq_middleware.rb
+++ b/lib/gitlab/sidekiq_middleware.rb
@@ -43,3 +43,5 @@ def self.client_configurator
     end
   end
 end
+
+Gitlab::SidekiqMiddleware.singleton_class.prepend_if_ee('EE::Gitlab::SidekiqMiddleware')
-- 
GitLab


From 1af21006512463f3a50bde924b9cd2cff9c884b0 Mon Sep 17 00:00:00 2001
From: nmilojevic1 <nmilojevic@gitlab.com>
Date: Thu, 25 Mar 2021 12:35:04 +0100
Subject: [PATCH 7/9] Fix chain ordering for sidekiq

- Fix spec to reflect new ordering
- Fix sidekiq server_configuration chain order
- Add ordering support for next_instance_of
---
 ee/lib/ee/gitlab/sidekiq_middleware.rb        |  5 ++-
 .../lib/ee/gitlab/sidekiq_middleware_spec.rb  | 36 +++++++++----------
 spec/support/helpers/next_instance_of.rb      | 11 +++---
 3 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/ee/lib/ee/gitlab/sidekiq_middleware.rb b/ee/lib/ee/gitlab/sidekiq_middleware.rb
index 61c3fc0c2b1471..dff07d51be4b0e 100644
--- a/ee/lib/ee/gitlab/sidekiq_middleware.rb
+++ b/ee/lib/ee/gitlab/sidekiq_middleware.rb
@@ -12,7 +12,10 @@ def server_configurator(metrics: true, arguments_logger: true, memory_killer: tr
         lambda do |chain|
           super.call(chain)
 
-          chain.add ::Gitlab::Database::LoadBalancing::SidekiqServerMiddleware if load_balancing_enabled?
+          if load_balancing_enabled?
+            chain.insert_after(::Labkit::Middleware::Sidekiq::Server,
+                               ::Gitlab::Database::LoadBalancing::SidekiqServerMiddleware)
+          end
         end
       end
 
diff --git a/ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb b/ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb
index a82f828f03dd24..7cb30be9ae1851 100644
--- a/ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb
+++ b/ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb
@@ -33,9 +33,8 @@ def perform(*args)
 
     it "passes through the right middlewares", :aggregate_failures 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).ordered.once.and_call_original
+        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
 
@@ -43,28 +42,27 @@ def perform(*args)
     end
   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::Database::LoadBalancing::SidekiqServerMiddleware
+        ::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::Database::LoadBalancing::SidekiqServerMiddleware,
+        ::Gitlab::SidekiqMiddleware::InstrumentationLogger,
+        ::Gitlab::SidekiqMiddleware::AdminMode::Server,
+        ::Gitlab::SidekiqVersioning::Middleware,
+        ::Gitlab::SidekiqStatus::ServerMiddleware,
+        ::Gitlab::SidekiqMiddleware::WorkerContext::Server,
+        ::Gitlab::SidekiqMiddleware::DuplicateJobs::Server
       ]
     end
 
diff --git a/spec/support/helpers/next_instance_of.rb b/spec/support/helpers/next_instance_of.rb
index a8e9ab2bafea3a..fdbc643c002c2c 100644
--- a/spec/support/helpers/next_instance_of.rb
+++ b/spec/support/helpers/next_instance_of.rb
@@ -5,22 +5,23 @@ def expect_next_instance_of(klass, *new_args, &blk)
     stub_new(expect(klass), nil, *new_args, &blk)
   end
 
-  def expect_next_instances_of(klass, number, *new_args, &blk)
-    stub_new(expect(klass), number, *new_args, &blk)
+  def expect_next_instances_of(klass, number, ordered = false, *new_args, &blk)
+    stub_new(expect(klass), number, ordered, *new_args, &blk)
   end
 
   def allow_next_instance_of(klass, *new_args, &blk)
     stub_new(allow(klass), nil, *new_args, &blk)
   end
 
-  def allow_next_instances_of(klass, number, *new_args, &blk)
-    stub_new(allow(klass), number, *new_args, &blk)
+  def allow_next_instances_of(klass, number, ordered = false, *new_args, &blk)
+    stub_new(allow(klass), number, ordered, *new_args, &blk)
   end
 
   private
 
-  def stub_new(target, number, *new_args, &blk)
+  def stub_new(target, number, ordered = false, *new_args, &blk)
     receive_new = receive(:new)
+    receive_new.ordered if ordered
     receive_new.exactly(number).times if number
     receive_new.with(*new_args) if new_args.any?
 
-- 
GitLab


From 39c9d79fc06a46d6e1eb893b1358908cdc7def67 Mon Sep 17 00:00:00 2001
From: nmilojevic1 <nmilojevic@gitlab.com>
Date: Thu, 25 Mar 2021 16:47:16 +0100
Subject: [PATCH 8/9] Add guards for sidekiq middleware

- Fix middleware specs
- Fix case when worker not include ApplicationWorker
---
 .../load_balancing/sidekiq_client_middleware.rb      |  9 ++++++---
 .../load_balancing/sidekiq_server_middleware.rb      |  2 ++
 ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb     | 10 ++++++++++
 .../load_balancing/sidekiq_client_middleware_spec.rb | 12 ++++++++++++
 .../load_balancing/sidekiq_server_middleware_spec.rb |  8 +++++++-
 5 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb b/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
index 7511435b5c97b6..3044585285f6a4 100644
--- a/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
+++ b/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
@@ -5,6 +5,8 @@ module Database
     module LoadBalancing
       class SidekiqClientMiddleware
         def call(worker_class, job, _queue, _redis_pool)
+          worker_class = worker_class.to_s.safe_constantize
+
           mark_data_consistency_location(worker_class, job)
 
           yield
@@ -13,10 +15,11 @@ def call(worker_class, job, _queue, _redis_pool)
         private
 
         def mark_data_consistency_location(worker_class, job)
-          worker_class = worker_class.to_s.safe_constantize
-          return if worker_class.get_data_consistency == :always
-
+          # Mailers can't be constantized
+          return unless worker_class
+          return unless worker_class.include?(::ApplicationWorker)
           return unless worker_class.get_data_consistency_feature_flag_enabled?
+          return if worker_class.get_data_consistency == :always
 
           if Session.current.performed_write?
             job['database_write_location'] = load_balancer.primary_write_location
diff --git a/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb b/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
index 7ce4cbcfe723a0..12578292748018 100644
--- a/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
+++ b/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
@@ -24,6 +24,8 @@ def clear
         end
 
         def requires_primary?(worker_class, job)
+          return true unless worker_class.include?(::ApplicationWorker)
+
           job[:worker_data_consistency] = worker_class.get_data_consistency
 
           return true if worker_class.get_data_consistency == :always
diff --git a/ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb b/ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb
index 7cb30be9ae1851..a6b71b5236e3d1 100644
--- a/ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb
+++ b/ee/spec/lib/ee/gitlab/sidekiq_middleware_spec.rb
@@ -42,6 +42,12 @@ def perform(*args)
     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] }
@@ -72,6 +78,7 @@ def perform(*args)
       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
@@ -82,6 +89,7 @@ def perform(*args)
       end
 
       it_behaves_like "a middleware chain", false
+      it_behaves_like "a middleware chain for mailer", false
     end
   end
 
@@ -111,10 +119,12 @@ def perform(*args)
       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
diff --git a/ee/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb b/ee/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
index b00cfa61b11cf0..468b34695504f4 100644
--- a/ee/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
+++ b/ee/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
@@ -96,6 +96,18 @@ def perform(*args)
       skip_default_enabled_yaml_check
     end
 
+    context 'when worker cannot be constantized' do
+      let(:worker_class) { 'ActionMailer::MailDeliveryJob' }
+
+      include_examples 'does not pass database locations'
+    end
+
+    context 'when worker class does not include ApplicationWorker' do
+      let(:worker_class) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper }
+
+      include_examples 'does not pass database locations'
+    end
+
     context 'when worker data consistency is :always' do
       include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker
 
diff --git a/ee/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb b/ee/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
index c929c8f1de7792..aa3a4bf28571d1 100644
--- a/ee/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
+++ b/ee/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
@@ -85,7 +85,7 @@ def perform(*args)
 
     let(:queue) { 'default' }
     let(:redis_pool) { Sidekiq.redis_pool }
-    let(:worker) { TestDataConsistencyWorker.new }
+    let(:worker) { worker_class.new }
     let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'primary_write_location' => '0/D525E3A8' } }
     let(:block) { 10 }
 
@@ -96,6 +96,12 @@ def perform(*args)
       allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:performed_write?).and_return(true)
     end
 
+    context 'when worker class does not include ApplicationWorker' do
+      let(:worker) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper.new }
+
+      include_examples 'stick to the primary'
+    end
+
     context 'when worker data consistency is :always' do
       include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker
 
-- 
GitLab


From 5ef1914142cf16dc9f6395174f13559e0df9c27a Mon Sep 17 00:00:00 2001
From: nmilojevic1 <nmilojevic@gitlab.com>
Date: Fri, 26 Mar 2021 01:23:26 +0100
Subject: [PATCH 9/9] Fix next_instance_of wrong params

---
 spec/support/helpers/next_instance_of.rb | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/spec/support/helpers/next_instance_of.rb b/spec/support/helpers/next_instance_of.rb
index fdbc643c002c2c..95d8936588c3e4 100644
--- a/spec/support/helpers/next_instance_of.rb
+++ b/spec/support/helpers/next_instance_of.rb
@@ -2,7 +2,7 @@
 
 module NextInstanceOf
   def expect_next_instance_of(klass, *new_args, &blk)
-    stub_new(expect(klass), nil, *new_args, &blk)
+    stub_new(expect(klass), nil, false, *new_args, &blk)
   end
 
   def expect_next_instances_of(klass, number, ordered = false, *new_args, &blk)
@@ -10,7 +10,7 @@ def expect_next_instances_of(klass, number, ordered = false, *new_args, &blk)
   end
 
   def allow_next_instance_of(klass, *new_args, &blk)
-    stub_new(allow(klass), nil, *new_args, &blk)
+    stub_new(allow(klass), nil, false, *new_args, &blk)
   end
 
   def allow_next_instances_of(klass, number, ordered = false, *new_args, &blk)
-- 
GitLab