Allow setting Sidekiq status only when opted in

This adds the `opt_in_sidekiq_status` feature flag. When this is
disabled, everything works as normal: we always set an entry in Redis
for a Sidekiq job's status.

When the flag is enabled, we will only perform the Redis SET command if
the job itself has opted in to tracking its status. It can do this on
the job class level:

    class SomeWorker
      sidekiq_options status_expiration: 1.hour
    end

Or on an individual call:

    SomeOtherWorker.with_status.perform_async

The aim here is to heavily reduce the number of SET commands we send to
the Sidekiq Redis, as most are unnecessary - most workers don't need to
have their status tracked.

config/feature_flags/development/opt_in_sidekiq_status.yml

+---
+name: opt_in_sidekiq_status
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/77349
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/343964
+milestone: '14.7'
+type: development
+group: group::scalability
+default_enabled: false

lib/gitlab/sidekiq_status.rb

@@ -34,6 +34,8 @@ module Gitlab
     # jid - The Sidekiq job ID
     # expire - The expiration time of the Redis key.
     def self.set(jid, expire = DEFAULT_EXPIRATION)
+      return unless expire
+
       Sidekiq.redis do |redis|
         redis.set(key_for(jid), 1, ex: expire)
       end

lib/gitlab/sidekiq_status/client_middleware.rb

@@ -4,10 +4,14 @@ module Gitlab
   module SidekiqStatus
     class ClientMiddleware
       def call(_, job, _, _)
-        status_expiration = job['status_expiration'] || Gitlab::SidekiqStatus::DEFAULT_EXPIRATION
-        value = job['status_expiration'] ? 2 : Gitlab::SidekiqStatus::DEFAULT_VALUE
+        status_expiration = job['status_expiration']
+
+        unless ::Feature.enabled?(:opt_in_sidekiq_status, default_enabled: :yaml)
+          status_expiration ||= Gitlab::SidekiqStatus::DEFAULT_EXPIRATION
+        end
+
+        Gitlab::SidekiqStatus.set(job['jid'], status_expiration)
 
-        Gitlab::SidekiqStatus.set(job['jid'], status_expiration, value: value)
         yield
       end
     end

spec/lib/gitlab/sidekiq_status/client_middleware_spec.rb

 # frozen_string_literal: true
 
-require 'fast_spec_helper'
+# This can use fast_spec_helper when the feature flag stubbing is removed.
+require 'spec_helper'
 
-RSpec.describe Gitlab::SidekiqStatus::ClientMiddleware do
+RSpec.describe Gitlab::SidekiqStatus::ClientMiddleware, :clean_gitlab_redis_queues do
   describe '#call' do
-    context 'when the job has status_expiration set' do
-      it 'tracks the job in Redis' do
-        expect(Gitlab::SidekiqStatus).to receive(:set).with('123', 1.hour.to_i)
+    context 'when opt_in_sidekiq_status is disabled' do
+      before do
+        stub_feature_flags(opt_in_sidekiq_status: false)
+      end
+
+      context 'when the job has status_expiration set' do
+        it 'tracks the job in Redis' do
+          expect(Gitlab::SidekiqStatus).to receive(:set).with('123', 1.hour.to_i).and_call_original
+
+          described_class.new
+            .call('Foo', { 'jid' => '123', 'status_expiration' => 1.hour.to_i }, double(:queue), double(:pool)) { nil }
+
+          expect(Gitlab::SidekiqStatus.num_running(['123'])).to eq(1)
+        end
+      end
+
+      context 'when the job does not have status_expiration set' do
+        it 'tracks the job in Redis' do
+          expect(Gitlab::SidekiqStatus).to receive(:set).with('123', 30.minutes.to_i).and_call_original
+
+          described_class.new
+            .call('Foo', { 'jid' => '123' }, double(:queue), double(:pool)) { nil }
 
-        described_class.new
-          .call('Foo', { 'jid' => '123', 'status_expiration' => 1.hour.to_i }, double(:queue), double(:pool)) { nil }
+          expect(Gitlab::SidekiqStatus.num_running(['123'])).to eq(1)
+        end
       end
     end
 
-    context 'when the job does not have status_expiration set' do
-      it 'tracks the job in Redis' do
-        expect(Gitlab::SidekiqStatus).to receive(:set).with('123', Gitlab::SidekiqStatus::DEFAULT_EXPIRATION)
+    context 'when opt_in_sidekiq_status is enabled' do
+      before do
+        stub_feature_flags(opt_in_sidekiq_status: true)
+      end
+
+      context 'when the job has status_expiration set' do
+        it 'tracks the job in Redis' do
+          expect(Gitlab::SidekiqStatus).to receive(:set).with('123', 1.hour.to_i).and_call_original
+
+          described_class.new
+            .call('Foo', { 'jid' => '123', 'status_expiration' => 1.hour.to_i }, double(:queue), double(:pool)) { nil }
+
+          expect(Gitlab::SidekiqStatus.num_running(['123'])).to eq(1)
+        end
+      end
+
+      context 'when the job does not have status_expiration set' do
+        it 'does not track the job in Redis' do
+          described_class.new
+            .call('Foo', { 'jid' => '123' }, double(:queue), double(:pool)) { nil }
 
-        described_class.new
-          .call('Foo', { 'jid' => '123' }, double(:queue), double(:pool)) { nil }
+          expect(Gitlab::SidekiqStatus.num_running(['123'])).to be_zero
+        end
       end
     end
   end

spec/lib/gitlab/sidekiq_status_spec.rb

@@ -27,6 +27,16 @@ RSpec.describe Gitlab::SidekiqStatus, :clean_gitlab_redis_queues, :clean_gitlab_
         expect(redis.get(key)).to eq('1')
       end
     end
+
+    it 'does not store anything with a nil expiry' do
+      described_class.set('123', nil)
+
+      key = described_class.key_for('123')
+
+      Sidekiq.redis do |redis|
+        expect(redis.exists(key)).to eq(false)
+      end
+    end
   end
 
   describe '.unset' do

spec/requests/api/merge_requests_spec.rb

@@ -3278,9 +3278,10 @@ RSpec.describe API::MergeRequests do
 
       context 'when skip_ci parameter is set' do
         it 'enqueues a rebase of the merge request with skip_ci flag set' do
-          allow(RebaseWorker).to receive(:with_status).and_return(RebaseWorker)
+          with_status = RebaseWorker.with_status
 
-          expect(RebaseWorker).to receive(:perform_async).with(merge_request.id, user.id, true).and_call_original
+          expect(RebaseWorker).to receive(:with_status).and_return(with_status)
+          expect(with_status).to receive(:perform_async).with(merge_request.id, user.id, true).and_call_original
 
           Sidekiq::Testing.fake! do
             expect do