Add Gitaly and Rugged call timing in Sidekiq logssh-add-gitaly-and-rugged-data-sidekiq

This will help identify Sidekiq jobs that invoke excessive number of
filesystem access.

The timing data is stored in `RequestStore`, but this is only active
within the middleware and is not directly accessible to the Sidekiq
logger. However, it is possible for the middleware to modify the job
hash to pass this data along to the logger.

9 files changed, 88 insertions, 17 deletions

diff --git a/changelogs/unreleased/sh-add-gitaly-and-rugged-data-sidekiq.yml b/changelogs/unreleased/sh-add-gitaly-and-rugged-data-sidekiq.yml
new file mode 100644
index 00000000000..d2143e83045
--- /dev/null
+++ b/changelogs/unreleased/sh-add-gitaly-and-rugged-data-sidekiq.yml
@@ -0,0 +1,5 @@
+---
+title: Add Gitaly and Rugged call timing in Sidekiq logs
+merge_request: 31651
+author:
+type: other
diff --git a/config/initializers/lograge.rb b/config/initializers/lograge.rb
index 3d84b4e44ce..346725e4080 100644
--- a/config/initializers/lograge.rb
+++ b/config/initializers/lograge.rb
@@ -27,19 +27,7 @@ unless Sidekiq.server?
         queue_duration: event.payload[:queue_duration]
       }
 
-      gitaly_calls = Gitlab::GitalyClient.get_request_count
-
-      if gitaly_calls > 0
-        payload[:gitaly_calls] = gitaly_calls
-        payload[:gitaly_duration] = Gitlab::GitalyClient.query_time_ms
-      end
-
-      rugged_calls = Gitlab::RuggedInstrumentation.query_count
-
-      if rugged_calls > 0
-        payload[:rugged_calls] = rugged_calls
-        payload[:rugged_duration_ms] = Gitlab::RuggedInstrumentation.query_time_ms
-      end
+      ::Gitlab::InstrumentationHelper.add_instrumentation_data(payload)
 
       payload[:response] = event.payload[:response] if event.payload[:response]
       payload[Labkit::Correlation::CorrelationId::LOG_KEY] = Labkit::Correlation::CorrelationId.current_id
diff --git a/config/initializers/sidekiq.rb b/config/initializers/sidekiq.rb
index 9f88fb9895b..7217f098fd9 100644
--- a/config/initializers/sidekiq.rb
+++ b/config/initializers/sidekiq.rb
@@ -39,6 +39,7 @@ Sidekiq.configure_server do |config|
     chain.add Gitlab::SidekiqMiddleware::RequestStoreMiddleware unless ENV['SIDEKIQ_REQUEST_STORE'] == '0'
     chain.add Gitlab::SidekiqMiddleware::BatchLoader
     chain.add Gitlab::SidekiqMiddleware::CorrelationLogger
+    chain.add Gitlab::SidekiqMiddleware::InstrumentationLogger
     chain.add Gitlab::SidekiqStatus::ServerMiddleware
   end
 
diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index 4783832961d..e6cbfb00f60 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -240,7 +240,7 @@ module Gitlab
 
     # Ensures that Gitaly is not being abuse through n+1 misuse etc
     def self.enforce_gitaly_request_limits(call_site)
-      # Only count limits in request-response environments (not sidekiq for example)
+      # Only count limits in request-response environments
       return unless Gitlab::SafeRequestStore.active?
 
       # This is this actual number of times this call was made. Used for information purposes only
diff --git a/lib/gitlab/instrumentation_helper.rb b/lib/gitlab/instrumentation_helper.rb
new file mode 100644
index 00000000000..e6a5facb2a5
--- /dev/null
+++ b/lib/gitlab/instrumentation_helper.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module InstrumentationHelper
+    extend self
+
+    KEYS = %i(gitaly_calls gitaly_duration rugged_calls rugged_duration_ms).freeze
+
+    def add_instrumentation_data(payload)
+      gitaly_calls = Gitlab::GitalyClient.get_request_count
+
+      if gitaly_calls > 0
+        payload[:gitaly_calls] = gitaly_calls
+        payload[:gitaly_duration] = Gitlab::GitalyClient.query_time_ms
+      end
+
+      rugged_calls = Gitlab::RuggedInstrumentation.query_count
+
+      if rugged_calls > 0
+        payload[:rugged_calls] = rugged_calls
+        payload[:rugged_duration_ms] = Gitlab::RuggedInstrumentation.query_time_ms
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/sidekiq_logging/structured_logger.rb b/lib/gitlab/sidekiq_logging/structured_logger.rb
index d556d5ef129..e1241392c4d 100644
--- a/lib/gitlab/sidekiq_logging/structured_logger.rb
+++ b/lib/gitlab/sidekiq_logging/structured_logger.rb
@@ -15,9 +15,9 @@ module Gitlab
 
         yield
 
-        Sidekiq.logger.info log_job_done(started_at, base_payload)
+        Sidekiq.logger.info log_job_done(job, started_at, base_payload)
       rescue => job_exception
-        Sidekiq.logger.warn log_job_done(started_at, base_payload, job_exception)
+        Sidekiq.logger.warn log_job_done(job, started_at, base_payload, job_exception)
 
         raise
       end
@@ -28,6 +28,10 @@ module Gitlab
         "#{payload['class']} JID-#{payload['jid']}"
       end
 
+      def add_instrumentation_keys(job, output_payload)
+        output_payload.merge!(job.slice(*::Gitlab::InstrumentationHelper::KEYS))
+      end
+
       def log_job_start(started_at, payload)
         payload['message'] = "#{base_message(payload)}: start"
         payload['job_status'] = 'start'
@@ -41,8 +45,9 @@ module Gitlab
         payload
       end
 
-      def log_job_done(started_at, payload, job_exception = nil)
+      def log_job_done(job, started_at, payload, job_exception = nil)
         payload = payload.dup
+        add_instrumentation_keys(job, payload)
         payload['duration'] = elapsed(started_at)
         payload['completed_at'] = Time.now.utc
 
diff --git a/lib/gitlab/sidekiq_middleware/instrumentation_logger.rb b/lib/gitlab/sidekiq_middleware/instrumentation_logger.rb
new file mode 100644
index 00000000000..8d60e630045
--- /dev/null
+++ b/lib/gitlab/sidekiq_middleware/instrumentation_logger.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module SidekiqMiddleware
+    class InstrumentationLogger
+      def call(worker, job, queue)
+        yield
+
+        # The Sidekiq logger is called outside the middleware block, so
+        # we need to modify the job hash to pass along this information
+        # since RequestStore is only active in the Sidekiq middleware.
+        #
+        # Modifying the job hash in a middlware is permitted by Sidekiq
+        # because Sidekiq keeps a pristine copy of the original hash before
+        # sending it to the middleware:
+        # https://github.com/mperham/sidekiq/blob/53bd529a0c3f901879925b8390353129c465b1f2/lib/sidekiq/processor.rb#L115-L118
+        ::Gitlab::InstrumentationHelper.add_instrumentation_data(job)
+      end
+    end
+  end
+end
diff --git a/scripts/lint-rugged b/scripts/lint-rugged
index d862571c1c5..1b3fb54f70b 100755
--- a/scripts/lint-rugged
+++ b/scripts/lint-rugged
@@ -19,6 +19,7 @@ ALLOWED = [
   'config/initializers/peek.rb',
   'config/initializers/lograge.rb',
   'lib/gitlab/grape_logging/loggers/perf_logger.rb',
+  'lib/gitlab/instrumentation_helper.rb',
   'lib/gitlab/rugged_instrumentation.rb',
   'lib/peek/views/rugged.rb'
 ].freeze
diff --git a/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb b/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb
index 98286eb432d..65de44c89d0 100644
--- a/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb
+++ b/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb
@@ -153,5 +153,30 @@ describe Gitlab::SidekiqLogging::StructuredLogger do
         end
       end
     end
+
+    context 'with Gitaly and Rugged calls' do
+      let(:timing_data) do
+        {
+          gitaly_calls: 10,
+          gitaly_duration: 10000,
+          rugged_calls: 1,
+          rugged_duration_ms: 5000
+        }
+      end
+
+      before do
+        job.merge!(timing_data)
+        end_payload.merge!(timing_data)
+      end
+
+      it 'logs with Gitaly and Rugged timing data' do
+        Timecop.freeze(timestamp) do
+          expect(logger).to receive(:info).with(start_payload.except('args')).ordered
+          expect(logger).to receive(:info).with(end_payload.except('args')).ordered
+
+          subject.call(job, 'test_queue') { }
+        end
+      end
+    end
   end
 end