Fix Sidekiq feature category logging

The intent of our feature category logging is:

1. If the worker is owned by a feature category, we log that category.
2. If the worker is not owned, we use the category from the surrounding
   context.

We removed the client middleware step previously because it interfered
with case 2 for metrics, because we could end up putting `not_owned`
into the context. With this change, we add that back, but only when the
worker is owned.

lib/gitlab/sidekiq_config/dummy_worker.rb

@@ -32,6 +32,10 @@ module Gitlab
         Gitlab::ApplicationContext.current_context_attribute('meta.feature_category') || :not_owned
       end
 
+      def feature_category_not_owned?
+        true
+      end
+
       def get_worker_context
         nil
       end

lib/gitlab/sidekiq_middleware/worker_context/client.rb

@@ -15,7 +15,19 @@ module Gitlab
 
           context_for_args = worker_class.context_for_arguments(job['args'])
 
-          wrap_in_optional_context(context_for_args, &block)
+          wrap_in_optional_context(context_for_args) do
+            # This should be inside the context for the arguments so
+            # that we don't override the feature category on the worker
+            # with the one from the caller.
+            #
+            # We do not want to set anything explicitly in the context
+            # when the feature category is 'not_owned'.
+            if worker_class.feature_category_not_owned?
+              yield
+            else
+              Gitlab::ApplicationContext.with_context(feature_category: worker_class.get_feature_category.to_s, &block)
+            end
+          end
         end
       end
     end

spec/lib/gitlab/sidekiq_middleware/worker_context/client_spec.rb

@@ -11,6 +11,8 @@ RSpec.describe Gitlab::SidekiqMiddleware::WorkerContext::Client do
 
       include ApplicationWorker
 
+      feature_category :issue_tracking
+
       def self.job_for_args(args)
         jobs.find { |job| job['args'] == args }
       end
@@ -20,8 +22,31 @@ RSpec.describe Gitlab::SidekiqMiddleware::WorkerContext::Client do
     end
   end
 
+  let(:not_owned_worker_class) do
+    Class.new(worker_class) do
+      def self.name
+        'TestNotOwnedWithContextWorker'
+      end
+
+      feature_category_not_owned!
+    end
+  end
+
+  let(:mailer_class) do
+    Class.new(ApplicationMailer) do
+      def self.name
+        'TestMailer'
+      end
+
+      def test_mail
+      end
+    end
+  end
+
   before do
     stub_const(worker_class.name, worker_class)
+    stub_const(not_owned_worker_class.name, not_owned_worker_class)
+    stub_const(mailer_class.name, mailer_class)
   end
 
   describe "#call" do
@@ -41,5 +66,75 @@ RSpec.describe Gitlab::SidekiqMiddleware::WorkerContext::Client do
       expect(job1['meta.user']).to eq(user_per_job['job1'].username)
       expect(job2['meta.user']).to eq(user_per_job['job2'].username)
     end
+
+    context 'when the feature category is set in the context_proc' do
+      it 'takes the feature category from the worker, not the caller' do
+        TestWithContextWorker.bulk_perform_async_with_contexts(
+          %w(job1 job2),
+          arguments_proc: -> (name) { [name, 1, 2, 3] },
+          context_proc: -> (_) { { feature_category: 'code_review' } }
+        )
+
+        job1 = TestWithContextWorker.job_for_args(['job1', 1, 2, 3])
+        job2 = TestWithContextWorker.job_for_args(['job2', 1, 2, 3])
+
+        expect(job1['meta.feature_category']).to eq('issue_tracking')
+        expect(job2['meta.feature_category']).to eq('issue_tracking')
+      end
+
+      it 'takes the feature category from the caller if the worker is not owned' do
+        TestNotOwnedWithContextWorker.bulk_perform_async_with_contexts(
+          %w(job1 job2),
+          arguments_proc: -> (name) { [name, 1, 2, 3] },
+          context_proc: -> (_) { { feature_category: 'code_review' } }
+        )
+
+        job1 = TestNotOwnedWithContextWorker.job_for_args(['job1', 1, 2, 3])
+        job2 = TestNotOwnedWithContextWorker.job_for_args(['job2', 1, 2, 3])
+
+        expect(job1['meta.feature_category']).to eq('code_review')
+        expect(job2['meta.feature_category']).to eq('code_review')
+      end
+
+      it 'does not set any explicit feature category for mailers', :sidekiq_mailers do
+        expect(Gitlab::ApplicationContext).not_to receive(:with_context)
+
+        TestMailer.test_mail.deliver_later
+      end
+    end
+
+    context 'when the feature category is already set in the surrounding block' do
+      it 'takes the feature category from the worker, not the caller' do
+        Gitlab::ApplicationContext.with_context(feature_category: 'authentication_and_authorization') do
+          TestWithContextWorker.bulk_perform_async_with_contexts(
+            %w(job1 job2),
+            arguments_proc: -> (name) { [name, 1, 2, 3] },
+            context_proc: -> (_) { {} }
+          )
+        end
+
+        job1 = TestWithContextWorker.job_for_args(['job1', 1, 2, 3])
+        job2 = TestWithContextWorker.job_for_args(['job2', 1, 2, 3])
+
+        expect(job1['meta.feature_category']).to eq('issue_tracking')
+        expect(job2['meta.feature_category']).to eq('issue_tracking')
+      end
+
+      it 'takes the feature category from the caller if the worker is not owned' do
+        Gitlab::ApplicationContext.with_context(feature_category: 'authentication_and_authorization') do
+          TestNotOwnedWithContextWorker.bulk_perform_async_with_contexts(
+            %w(job1 job2),
+            arguments_proc: -> (name) { [name, 1, 2, 3] },
+            context_proc: -> (_) { {} }
+          )
+        end
+
+        job1 = TestNotOwnedWithContextWorker.job_for_args(['job1', 1, 2, 3])
+        job2 = TestNotOwnedWithContextWorker.job_for_args(['job2', 1, 2, 3])
+
+        expect(job1['meta.feature_category']).to eq('authentication_and_authorization')
+        expect(job2['meta.feature_category']).to eq('authentication_and_authorization')
+      end
+    end
   end
 end