Merge branch 'sh-add-gitaly-and-rugged-data-sidekiq' into 'master'

Add Gitaly and Rugged call timing in Sidekiq logs

See merge request gitlab-org/gitlab-ce!31651

9 files changed, 87 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..b17135922f6 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..979a3fce7e6
--- /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 middleware 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..b118d3f7017 100644
--- a/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb
+++ b/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb
@@ -153,5 +153,29 @@ 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
+
+      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