1 files changed, 30 insertions, 17 deletions
diff --git a/lib/gitlab/sidekiq_middleware/server_metrics.rb b/lib/gitlab/sidekiq_middleware/server_metrics.rb
index cf768811ffd..f5fee8050ac 100644
--- a/lib/gitlab/sidekiq_middleware/server_metrics.rb
+++ b/lib/gitlab/sidekiq_middleware/server_metrics.rb
@@ -21,6 +21,16 @@ module Gitlab
         Thread.current.name ||= Gitlab::Metrics::Samplers::ThreadsSampler::SIDEKIQ_WORKER_THREAD_NAME
 
         labels = create_labels(worker.class, queue, job)
+        instrument(job, labels) do
+          yield
+        end
+      end
+
+      protected
+
+      attr_reader :metrics
+
+      def instrument(job, labels)
         queue_duration = ::Gitlab::InstrumentationHelper.queue_duration_for_job(job)
 
         @metrics[:sidekiq_jobs_queue_duration_seconds].observe(labels, queue_duration) if queue_duration
@@ -50,19 +60,18 @@ module Gitlab
 
           # job_status: done, fail match the job_status attribute in structured logging
           labels[:job_status] = job_succeeded ? "done" : "fail"
+          instrumentation = job[:instrumentation] || {}
           @metrics[:sidekiq_jobs_cpu_seconds].observe(labels, job_thread_cputime)
           @metrics[:sidekiq_jobs_completion_seconds].observe(labels, monotonic_time)
           @metrics[:sidekiq_jobs_db_seconds].observe(labels, ActiveRecord::LogSubscriber.runtime / 1000)
-          @metrics[:sidekiq_jobs_gitaly_seconds].observe(labels, get_gitaly_time(job))
-          @metrics[:sidekiq_redis_requests_total].increment(labels, get_redis_calls(job))
-          @metrics[:sidekiq_redis_requests_duration_seconds].observe(labels, get_redis_time(job))
-          @metrics[:sidekiq_elasticsearch_requests_total].increment(labels, get_elasticsearch_calls(job))
-          @metrics[:sidekiq_elasticsearch_requests_duration_seconds].observe(labels, get_elasticsearch_time(job))
+          @metrics[:sidekiq_jobs_gitaly_seconds].observe(labels, get_gitaly_time(instrumentation))
+          @metrics[:sidekiq_redis_requests_total].increment(labels, get_redis_calls(instrumentation))
+          @metrics[:sidekiq_redis_requests_duration_seconds].observe(labels, get_redis_time(instrumentation))
+          @metrics[:sidekiq_elasticsearch_requests_total].increment(labels, get_elasticsearch_calls(instrumentation))
+          @metrics[:sidekiq_elasticsearch_requests_duration_seconds].observe(labels, get_elasticsearch_time(instrumentation))
         end
       end
 
-      private
-
       def init_metrics
         {
           sidekiq_jobs_cpu_seconds:                ::Gitlab::Metrics.histogram(:sidekiq_jobs_cpu_seconds, 'Seconds of cpu time to run Sidekiq job', {}, SIDEKIQ_LATENCY_BUCKETS),
@@ -81,29 +90,33 @@ module Gitlab
         }
       end
 
+      private
+
       def get_thread_cputime
         defined?(Process::CLOCK_THREAD_CPUTIME_ID) ? Process.clock_gettime(Process::CLOCK_THREAD_CPUTIME_ID) : 0
       end
 
-      def get_redis_time(job)
-        job.fetch(:redis_duration_s, 0)
+      def get_redis_time(payload)
+        payload.fetch(:redis_duration_s, 0)
       end
 
-      def get_redis_calls(job)
-        job.fetch(:redis_calls, 0)
+      def get_redis_calls(payload)
+        payload.fetch(:redis_calls, 0)
       end
 
-      def get_elasticsearch_time(job)
-        job.fetch(:elasticsearch_duration_s, 0)
+      def get_elasticsearch_time(payload)
+        payload.fetch(:elasticsearch_duration_s, 0)
       end
 
-      def get_elasticsearch_calls(job)
-        job.fetch(:elasticsearch_calls, 0)
+      def get_elasticsearch_calls(payload)
+        payload.fetch(:elasticsearch_calls, 0)
       end
 
-      def get_gitaly_time(job)
-        job.fetch(:gitaly_duration_s, 0)
+      def get_gitaly_time(payload)
+        payload.fetch(:gitaly_duration_s, 0)
       end
     end
   end
 end
+
+Gitlab::SidekiqMiddleware::ServerMetrics.prepend_if_ee('EE::Gitlab::SidekiqMiddleware::ServerMetrics')