Merge branch 'sidekiq-metadata-log' into 'master'

Log extra metadata sidekiq done for ElasticIndexBulkCronWorker

See merge request gitlab-org/gitlab!31297

app/workers/concerns/application_worker.rb

@@ -11,6 +11,8 @@ module ApplicationWorker
   include WorkerAttributes
   include WorkerContext
 
+  LOGGING_EXTRA_KEY = 'extra'
+
   included do
     set_queue
 
@@ -24,6 +26,21 @@ module ApplicationWorker
 
       payload.stringify_keys.merge(context)
     end
+
+    def log_extra_metadata_on_done(key, value)
+      @done_log_extra_metadata ||= {}
+      @done_log_extra_metadata[key] = value
+    end
+
+    def logging_extras
+      return {} unless @done_log_extra_metadata
+
+      # Prefix keys with class name to avoid conflicts in Elasticsearch types.
+      # Also prefix with "extra." so that we know to log these new fields.
+      @done_log_extra_metadata.transform_keys do |k|
+        "#{LOGGING_EXTRA_KEY}.#{self.class.name.gsub("::", "_").underscore}.#{k}"
+      end
+    end
   end
 
   class_methods do

ee/app/services/elastic/process_bookkeeping_service.rb

@@ -57,7 +57,7 @@ module Elastic
       start_time = Time.current
 
       specs = redis.zrangebyscore(REDIS_SET_KEY, '-inf', '+inf', limit: [0, LIMIT], with_scores: true)
-      return if specs.empty?
+      return 0 if specs.empty?
 
       first_score = specs.first.last
       last_score = specs.last.last
@@ -79,14 +79,18 @@ module Elastic
       # Remove all the successes
       redis.zremrangebyscore(REDIS_SET_KEY, first_score, last_score)
 
+      records_count = specs.count
+
       logger.info(
         message: 'bulk_indexing_end',
-        records_count: specs.count,
+        records_count: records_count,
         failures_count: failures.count,
         first_score: first_score,
         last_score: last_score,
         bulk_execution_duration_s: Time.current - start_time
       )
+
+      records_count
     end
 
     def deserialize_all(specs)

ee/app/workers/elastic_index_bulk_cron_worker.rb

@@ -14,7 +14,8 @@ class ElasticIndexBulkCronWorker
 
   def perform
     in_lock(self.class.name.underscore, ttl: 10.minutes, retries: 10, sleep_sec: 1) do
-      Elastic::ProcessBookkeepingService.new.execute
+      records_count = Elastic::ProcessBookkeepingService.new.execute
+      log_extra_metadata_on_done(:records_count, records_count)
     end
   rescue Gitlab::ExclusiveLeaseHelpers::FailedToObtainLockError
     # We're scheduled on a cronjob, so nothing to do here

ee/spec/services/elastic/process_bookkeeping_service_spec.rb

@@ -93,6 +93,20 @@ describe Elastic::ProcessBookkeepingService, :clean_gitlab_redis_shared_state do
       expect { described_class.new.execute }.to change(described_class, :queue_size).by(-limit)
     end
 
+    it 'returns the number of documents processed' do
+      described_class.track!(*fake_refs)
+
+      expect_processing(*fake_refs[0...limit])
+
+      expect(described_class.new.execute).to eq(limit)
+    end
+
+    it 'returns 0 without writing to the index when there are no documents' do
+      expect(::Gitlab::Elastic::BulkIndexer).not_to receive(:new)
+
+      expect(described_class.new.execute).to eq(0)
+    end
+
     it 'retries failed documents' do
       described_class.track!(*fake_refs)
       failed = fake_refs[0]

ee/spec/workers/elastic_index_bulk_cron_worker_spec.rb

@@ -14,5 +14,18 @@ describe ElasticIndexBulkCronWorker do
 
       described_class.new.perform
     end
+
+    it 'adds the elastic_bulk_count to the done log' do
+      expect_next_instance_of(::Elastic::ProcessBookkeepingService) do |service|
+        expect(service).to receive(:execute).and_return(15)
+      end
+      worker = described_class.new
+
+      worker.perform
+
+      expect(worker.logging_extras).to eq(
+        "#{ApplicationWorker::LOGGING_EXTRA_KEY}.elastic_index_bulk_cron_worker.records_count" => 15
+      )
+    end
   end
 end

lib/gitlab/sidekiq_logging/structured_logger.rb

@@ -30,6 +30,12 @@ module Gitlab
         output_payload.merge!(job.slice(*::Gitlab::InstrumentationHelper::KEYS))
       end
 
+      def add_logging_extras!(job, output_payload)
+        output_payload.merge!(
+          job.select { |key, _| key.to_s.start_with?("#{ApplicationWorker::LOGGING_EXTRA_KEY}.") }
+        )
+      end
+
       def log_job_start(payload)
         payload['message'] = "#{base_message(payload)}: start"
         payload['job_status'] = 'start'
@@ -43,6 +49,7 @@ module Gitlab
       def log_job_done(job, started_time, payload, job_exception = nil)
         payload = payload.dup
         add_instrumentation_keys!(job, payload)
+        add_logging_extras!(job, payload)
 
         elapsed_time = elapsed(started_time)
         add_time_keys!(elapsed_time, payload)

lib/gitlab/sidekiq_middleware.rb

@@ -14,6 +14,7 @@ module Gitlab
         chain.add ::Gitlab::SidekiqMiddleware::ArgumentsLogger if arguments_logger
         chain.add ::Gitlab::SidekiqMiddleware::MemoryKiller if memory_killer
         chain.add ::Gitlab::SidekiqMiddleware::RequestStoreMiddleware
+        chain.add ::Gitlab::SidekiqMiddleware::ExtraDoneLogMetadata
         chain.add ::Gitlab::SidekiqMiddleware::BatchLoader
         chain.add ::Labkit::Middleware::Sidekiq::Server
         chain.add ::Gitlab::SidekiqMiddleware::InstrumentationLogger

lib/gitlab/sidekiq_middleware/extra_done_log_metadata.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module SidekiqMiddleware
+    class ExtraDoneLogMetadata
+      def call(worker, job, queue)
+        yield
+
+        # We needed a way to pass state from a worker in to the
+        # Gitlab::SidekiqLogging::StructuredLogger . Unfortunately the
+        # StructuredLogger itself is not a middleware so cannot access the
+        # worker object. We also tried to use SafeRequestStore but to pass the
+        # data up but that doesn't work either because this is reset in
+        # Gitlab::SidekiqMiddleware::RequestStoreMiddleware inside yield for
+        # the StructuredLogger so it's cleared before we get to logging the
+        # done statement. As such the only way to do this is to pass the data
+        # up in the `job` object. Since `job` is just a Hash we can add this
+        # extra metadata there.
+        if worker.respond_to?(:logging_extras)
+          job.merge!(worker.logging_extras)
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb

@@ -218,6 +218,27 @@ describe Gitlab::SidekiqLogging::StructuredLogger do
         subject.call(job, 'test_queue') { }
       end
     end
+
+    context 'when there is extra metadata set for the done log' do
+      let(:expected_start_payload) { start_payload.except('args') }
+
+      let(:expected_end_payload) do
+        end_payload.except('args').merge("#{ApplicationWorker::LOGGING_EXTRA_KEY}.key1" => 15, "#{ApplicationWorker::LOGGING_EXTRA_KEY}.key2" => 16)
+      end
+
+      it 'logs it in the done log' do
+        Timecop.freeze(timestamp) do
+          expect(logger).to receive(:info).with(expected_start_payload).ordered
+          expect(logger).to receive(:info).with(expected_end_payload).ordered
+
+          subject.call(job, 'test_queue') do
+            job["#{ApplicationWorker::LOGGING_EXTRA_KEY}.key1"] = 15
+            job["#{ApplicationWorker::LOGGING_EXTRA_KEY}.key2"] = 16
+            job['key that will be ignored because it does not start with extra.'] = 17
+          end
+        end
+      end
+    end
   end
 
   describe '#add_time_keys!' do

spec/lib/gitlab/sidekiq_middleware/extra_done_log_metadata_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::SidekiqMiddleware::ExtraDoneLogMetadata do
+  # Cannot use Class.new for this as ApplicationWorker will need the class to
+  # have a name during `included do`.
+  let(:worker) { AdminEmailWorker.new }
+
+  let(:worker_without_application_worker) do
+    Class.new do
+    end.new
+  end
+
+  subject { described_class.new }
+
+  let(:job) { { 'jid' => 123 } }
+  let(:queue) { 'test_queue' }
+
+  describe '#call' do
+    it 'merges Application#logging_extras in to job' do
+      worker.log_extra_metadata_on_done(:key1, 15)
+      worker.log_extra_metadata_on_done(:key2, 16)
+      expect { |b| subject.call(worker, job, queue, &b) }.to yield_control
+
+      expect(job).to eq({ 'jid' => 123, 'extra.admin_email_worker.key1' => 15, 'extra.admin_email_worker.key2' => 16 })
+    end
+
+    it 'does not raise when the worker does not respond to #done_log_extra_metadata' do
+      expect { |b| subject.call(worker_without_application_worker, job, queue, &b) }.to yield_control
+
+      expect(job).to eq({ 'jid' => 123 })
+    end
+  end
+end

spec/lib/gitlab/sidekiq_middleware_spec.rb

@@ -9,6 +9,7 @@ describe Gitlab::SidekiqMiddleware do
 
     TestWorker.class_eval do
       include Sidekiq::Worker
+      include ApplicationWorker
 
       def perform(_arg)
         Gitlab::SafeRequestStore['gitaly_call_actual'] = 1
@@ -55,6 +56,7 @@ describe Gitlab::SidekiqMiddleware do
        Gitlab::SidekiqMiddleware::ArgumentsLogger,
        Gitlab::SidekiqMiddleware::MemoryKiller,
        Gitlab::SidekiqMiddleware::RequestStoreMiddleware,
+       Gitlab::SidekiqMiddleware::ExtraDoneLogMetadata,
        Gitlab::SidekiqMiddleware::WorkerContext::Server,
        Gitlab::SidekiqMiddleware::AdminMode::Server,
        Gitlab::SidekiqMiddleware::DuplicateJobs::Server

spec/workers/concerns/application_worker_spec.rb

@@ -21,6 +21,21 @@ describe ApplicationWorker do
     end
   end
 
+  describe '#logging_extras' do
+    it 'returns extra data to be logged that was set from #log_extra_metadata_on_done' do
+      instance.log_extra_metadata_on_done(:key1, "value1")
+      instance.log_extra_metadata_on_done(:key2, "value2")
+
+      expect(instance.logging_extras).to eq({ 'extra.gitlab_foo_bar_dummy_worker.key1' => "value1", 'extra.gitlab_foo_bar_dummy_worker.key2' => "value2" })
+    end
+
+    context 'when nothing is set' do
+      it 'returns {}' do
+        expect(instance.logging_extras).to eq({})
+      end
+    end
+  end
+
   describe '#structured_payload' do
     let(:payload) { {} }