13 files changed, 70 insertions, 4 deletions

diff --git a/app/workers/build_finished_worker.rb b/app/workers/build_finished_worker.rb
index 8e2a18a8fd8..2f09d4e14e8 100644
--- a/app/workers/build_finished_worker.rb
+++ b/app/workers/build_finished_worker.rb
@@ -6,6 +6,10 @@ class BuildFinishedWorker
 
   queue_namespace :pipeline_processing
 
+  def self.sidekiq_status_enabled?
+    false
+  end
+
   # rubocop: disable CodeReuse/ActiveRecord
   def perform(build_id)
     Ci::Build.find_by(id: build_id).try do |build|
diff --git a/app/workers/build_queue_worker.rb b/app/workers/build_queue_worker.rb
index 67d5b0f5f5b..95a2f8fcfb7 100644
--- a/app/workers/build_queue_worker.rb
+++ b/app/workers/build_queue_worker.rb
@@ -6,6 +6,10 @@ class BuildQueueWorker
 
   queue_namespace :pipeline_processing
 
+  def self.sidekiq_status_enabled?
+    false
+  end
+
   # rubocop: disable CodeReuse/ActiveRecord
   def perform(build_id)
     Ci::Build.find_by(id: build_id).try do |build|
diff --git a/app/workers/concerns/application_worker.rb b/app/workers/concerns/application_worker.rb
index 2b36ccb8304..ba902d2c469 100644
--- a/app/workers/concerns/application_worker.rb
+++ b/app/workers/concerns/application_worker.rb
@@ -64,5 +64,9 @@ module ApplicationWorker
 
       Sidekiq::Client.push_bulk('class' => self, 'args' => args_list, 'at' => schedule)
     end
+
+    def sidekiq_status_enabled?
+      true
+    end
   end
 end
diff --git a/app/workers/expire_pipeline_cache_worker.rb b/app/workers/expire_pipeline_cache_worker.rb
index 78e68d7bf46..f135c4cccb9 100644
--- a/app/workers/expire_pipeline_cache_worker.rb
+++ b/app/workers/expire_pipeline_cache_worker.rb
@@ -6,6 +6,10 @@ class ExpirePipelineCacheWorker
 
   queue_namespace :pipeline_cache
 
+  def self.sidekiq_status_enabled?
+    false
+  end
+
   # rubocop: disable CodeReuse/ActiveRecord
   def perform(pipeline_id)
     pipeline = Ci::Pipeline.find_by(id: pipeline_id)
diff --git a/app/workers/pipeline_process_worker.rb b/app/workers/pipeline_process_worker.rb
index f2aa17acb51..388a26cca74 100644
--- a/app/workers/pipeline_process_worker.rb
+++ b/app/workers/pipeline_process_worker.rb
@@ -6,6 +6,10 @@ class PipelineProcessWorker
 
   queue_namespace :pipeline_processing
 
+  def self.sidekiq_status_enabled?
+    false
+  end
+
   # rubocop: disable CodeReuse/ActiveRecord
   def perform(pipeline_id)
     Ci::Pipeline.find_by(id: pipeline_id)
diff --git a/app/workers/stage_update_worker.rb b/app/workers/stage_update_worker.rb
index ea587789d03..c3f7429e155 100644
--- a/app/workers/stage_update_worker.rb
+++ b/app/workers/stage_update_worker.rb
@@ -6,6 +6,10 @@ class StageUpdateWorker
 
   queue_namespace :pipeline_processing
 
+  def self.sidekiq_status_enabled?
+    false
+  end
+
   # rubocop: disable CodeReuse/ActiveRecord
   def perform(stage_id)
     Ci::Stage.find_by(id: stage_id).try do |stage|
diff --git a/app/workers/update_project_statistics_worker.rb b/app/workers/update_project_statistics_worker.rb
index 9a29cc12707..ef5a37707ac 100644
--- a/app/workers/update_project_statistics_worker.rb
+++ b/app/workers/update_project_statistics_worker.rb
@@ -5,6 +5,10 @@
 class UpdateProjectStatisticsWorker
   include ApplicationWorker
 
+  def self.sidekiq_status_enabled?
+    false
+  end
+
   # project_id - The ID of the project for which to flush the cache.
   # statistics - An Array containing columns from ProjectStatistics to
   #              refresh, if empty all columns will be refreshed
diff --git a/app/workers/web_hook_worker.rb b/app/workers/web_hook_worker.rb
index 09219a24a16..1dfcb45a8ea 100644
--- a/app/workers/web_hook_worker.rb
+++ b/app/workers/web_hook_worker.rb
@@ -5,6 +5,10 @@ class WebHookWorker
 
   sidekiq_options retry: 4, dead: false
 
+  def self.sidekiq_status_enabled?
+    false
+  end
+
   def perform(hook_id, data, hook_name)
     hook = WebHook.find(hook_id)
     data = data.with_indifferent_access
diff --git a/lib/gitlab/sidekiq_status/base_middleware.rb b/lib/gitlab/sidekiq_status/base_middleware.rb
new file mode 100644
index 00000000000..3682122e3d5
--- /dev/null
+++ b/lib/gitlab/sidekiq_status/base_middleware.rb
@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module SidekiqStatus
+    class BaseMiddleware
+      # @param [Hash] job the full job payload
+      def status_enabled?(job)
+        worker = job['class']&.constantize
+
+        byebug
+        return worker.sidekiq_status_enabled? if worker&.respond_to?(:sidekiq_status_enabled?)
+
+        true
+      rescue NameError
+        true
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/sidekiq_status/client_middleware.rb b/lib/gitlab/sidekiq_status/client_middleware.rb
index bfd5038557d..b69d84531f0 100644
--- a/lib/gitlab/sidekiq_status/client_middleware.rb
+++ b/lib/gitlab/sidekiq_status/client_middleware.rb
@@ -2,11 +2,12 @@
 
 module Gitlab
   module SidekiqStatus
-    class ClientMiddleware
+    class ClientMiddleware < BaseMiddleware
       def call(_, job, _, _)
         status_expiration = job['status_expiration'] || Gitlab::SidekiqStatus::DEFAULT_EXPIRATION
 
-        Gitlab::SidekiqStatus.set(job['jid'], status_expiration)
+        Gitlab::SidekiqStatus.set(job['jid'], status_expiration) if status_enabled?(job)
+
         yield
       end
     end
diff --git a/lib/gitlab/sidekiq_status/server_middleware.rb b/lib/gitlab/sidekiq_status/server_middleware.rb
index 01bc58fd2be..7a2ae99d155 100644
--- a/lib/gitlab/sidekiq_status/server_middleware.rb
+++ b/lib/gitlab/sidekiq_status/server_middleware.rb
@@ -2,11 +2,11 @@
 
 module Gitlab
   module SidekiqStatus
-    class ServerMiddleware
+    class ServerMiddleware < BaseMiddleware
       def call(worker, job, queue)
         ret = yield
 
-        Gitlab::SidekiqStatus.unset(job['jid'])
+        Gitlab::SidekiqStatus.unset(job['jid']) if status_enabled?(job)
 
         ret
       end
diff --git a/spec/lib/gitlab/sidekiq_status/client_middleware_spec.rb b/spec/lib/gitlab/sidekiq_status/client_middleware_spec.rb
index 37d9e1d3e6b..bc2f0eddcb6 100644
--- a/spec/lib/gitlab/sidekiq_status/client_middleware_spec.rb
+++ b/spec/lib/gitlab/sidekiq_status/client_middleware_spec.rb
@@ -8,5 +8,12 @@ describe Gitlab::SidekiqStatus::ClientMiddleware do
       described_class.new
         .call('Foo', { 'jid' => '123' }, double(:queue), double(:pool)) { nil }
     end
+
+    it 'does not track the job in Redis for excluded jobs' do
+      expect(Gitlab::SidekiqStatus).not_to receive(:set)
+
+      described_class.new
+        .call('BuildQueueWorker', { 'class' => 'BuildQueueWorker', 'jid' => '123' }, double(:queue), double(:pool)) { nil }
+    end
   end
 end
diff --git a/spec/lib/gitlab/sidekiq_status/server_middleware_spec.rb b/spec/lib/gitlab/sidekiq_status/server_middleware_spec.rb
index 04e09d3dec8..379b39958df 100644
--- a/spec/lib/gitlab/sidekiq_status/server_middleware_spec.rb
+++ b/spec/lib/gitlab/sidekiq_status/server_middleware_spec.rb
@@ -10,5 +10,12 @@ describe Gitlab::SidekiqStatus::ServerMiddleware do
 
       expect(ret).to eq(10)
     end
+
+    it 'does not track the job in Redis for excluded jobs' do
+      expect(Gitlab::SidekiqStatus).not_to receive(:unset)
+
+      described_class.new
+        .call(double(:worker), { 'class' => 'BuildQueueWorker', 'jid' => '123' }, double(:queue)) { 10 }
+    end
   end
 end