Add load balancing counters to structured_log

- Move keys to instrumentation_logger middleware
- Slice instrumentation values in logger middleware
- Add missing specs for instrumentation_logger

ee/lib/ee/gitlab/metrics/subscribers/active_record.rb

@@ -17,6 +17,11 @@ module EE
           class_methods do
             extend ::Gitlab::Utils::Override
 
+            override :known_payload_keys
+            def known_payload_keys
+              super + DB_LOAD_BALANCING_COUNTERS
+            end
+
             override :db_counter_payload
             def db_counter_payload
               super.tap do |payload|

ee/spec/lib/gitlab/instrumentation_helper_spec.rb

@@ -8,6 +8,15 @@ RSpec.describe Gitlab::InstrumentationHelper do
 
     subject { described_class.add_instrumentation_data(payload) }
 
+    it 'includes DB counts' do
+      subject
+
+      expect(payload).to include(db_replica_count: 0,
+                                 db_replica_cached_count: 0,
+                                 db_primary_count: 0,
+                                 db_primary_cached_count: 0)
+    end
+
     # We don't want to interact with Elasticsearch in GitLab FOSS so we test
     # this in ee/ only. The code exists in FOSS and won't do anything.
     context 'when Elasticsearch calls are made', :elastic do

ee/spec/lib/gitlab/sidekiq_middleware/instrumentation_logger_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::SidekiqMiddleware::InstrumentationLogger do
+
+  describe '.keys' do
+    it 'contains load balancer keys' do
+      expected_keys = [
+        :db_replica_count,
+        :db_replica_cached_count,
+        :db_primary_count,
+        :db_primary_cached_count
+      ]
+
+      expect(described_class.keys).to include(*expected_keys)
+    end
+  end
+end

lib/gitlab/instrumentation_helper.rb

@@ -6,25 +6,6 @@ module Gitlab
 
     DURATION_PRECISION = 6 # microseconds
 
-    def keys
-      @keys ||= [
-        :cpu_s,
-        :gitaly_calls,
-        :gitaly_duration_s,
-        :rugged_calls,
-        :rugged_duration_s,
-        :elasticsearch_calls,
-        :elasticsearch_duration_s,
-        :elasticsearch_timed_out_count,
-        :worker_data_consistency,
-        *::Gitlab::Memory::Instrumentation::KEY_MAPPING.values,
-        *::Gitlab::Instrumentation::Redis.known_payload_keys,
-        *::Gitlab::Metrics::Subscribers::ActiveRecord::DB_COUNTERS,
-        *::Gitlab::Metrics::Subscribers::ExternalHttp::KNOWN_PAYLOAD_KEYS,
-        *::Gitlab::Metrics::Subscribers::RackAttack::PAYLOAD_KEYS
-      ]
-    end
-
     def init_instrumentation_data(request_ip: nil)
       # Set `request_start_time` only if this is request
       # This is done, as `request_start_time` imply `request_deadline`

lib/gitlab/metrics/subscribers/active_record.rb

@@ -51,6 +51,10 @@ module Gitlab
           payload
         end
 
+        def self.known_payload_keys
+          DB_COUNTERS
+        end
+
         private
 
         def ignored_query?(payload)

lib/gitlab/sidekiq_logging/structured_logger.rb

@@ -39,9 +39,7 @@ module Gitlab
       private
 
       def add_instrumentation_keys!(job, output_payload)
-        instrumentation_values = job.slice(*::Gitlab::InstrumentationHelper.keys).stringify_keys
-
-        output_payload.merge!(instrumentation_values)
+        output_payload.merge!(job[:instrumentation])
       end
 
       def add_logging_extras!(job, output_payload)

lib/gitlab/sidekiq_middleware/instrumentation_logger.rb

@@ -3,6 +3,25 @@
 module Gitlab
   module SidekiqMiddleware
     class InstrumentationLogger
+      def self.keys
+        @keys ||= [
+          :cpu_s,
+          :gitaly_calls,
+          :gitaly_duration_s,
+          :rugged_calls,
+          :rugged_duration_s,
+          :elasticsearch_calls,
+          :elasticsearch_duration_s,
+          :elasticsearch_timed_out_count,
+          :worker_data_consistency,
+          *::Gitlab::Memory::Instrumentation::KEY_MAPPING.values,
+          *::Gitlab::Instrumentation::Redis.known_payload_keys,
+          *::Gitlab::Metrics::Subscribers::ActiveRecord.known_payload_keys,
+          *::Gitlab::Metrics::Subscribers::ExternalHttp::KNOWN_PAYLOAD_KEYS,
+          *::Gitlab::Metrics::Subscribers::RackAttack::PAYLOAD_KEYS
+        ]
+      end
+
       def call(worker, job, queue)
         ::Gitlab::InstrumentationHelper.init_instrumentation_data
 
@@ -17,7 +36,11 @@ module Gitlab
         # 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)
+        job[:instrumentation] = {}.tap do |instrumentation_values|
+          ::Gitlab::InstrumentationHelper.add_instrumentation_data(instrumentation_values)
+          instrumentation_values.slice!(*self.class.keys)
+          instrumentation_values.stringify_keys!
+        end
       end
     end
   end

spec/lib/gitlab/instrumentation_helper_spec.rb

@@ -6,54 +6,6 @@ require 'rspec-parameterized'
 RSpec.describe Gitlab::InstrumentationHelper do
   using RSpec::Parameterized::TableSyntax
 
-  describe '.keys' do
-    it 'returns all available payload keys' do
-      expected_keys = [
-        :cpu_s,
-        :gitaly_calls,
-        :gitaly_duration_s,
-        :rugged_calls,
-        :rugged_duration_s,
-        :elasticsearch_calls,
-        :elasticsearch_duration_s,
-        :elasticsearch_timed_out_count,
-        :worker_data_consistency,
-        :mem_objects,
-        :mem_bytes,
-        :mem_mallocs,
-        :redis_calls,
-        :redis_duration_s,
-        :redis_read_bytes,
-        :redis_write_bytes,
-        :redis_action_cable_calls,
-        :redis_action_cable_duration_s,
-        :redis_action_cable_read_bytes,
-        :redis_action_cable_write_bytes,
-        :redis_cache_calls,
-        :redis_cache_duration_s,
-        :redis_cache_read_bytes,
-        :redis_cache_write_bytes,
-        :redis_queues_calls,
-        :redis_queues_duration_s,
-        :redis_queues_read_bytes,
-        :redis_queues_write_bytes,
-        :redis_shared_state_calls,
-        :redis_shared_state_duration_s,
-        :redis_shared_state_read_bytes,
-        :redis_shared_state_write_bytes,
-        :db_count,
-        :db_write_count,
-        :db_cached_count,
-        :external_http_count,
-        :external_http_duration_s,
-        :rack_attack_redis_count,
-        :rack_attack_redis_duration_s
-      ]
-
-      expect(described_class.keys).to eq(expected_keys)
-    end
-  end
-
   describe '.add_instrumentation_data', :request_store do
     let(:payload) { {} }
 

spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb

@@ -283,14 +283,19 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
         end_payload.merge(timing_data.stringify_keys)
       end
 
-      it 'logs with Gitaly and Rugged timing data' do
+      before do
+        allow(::Gitlab::InstrumentationHelper).to receive(:add_instrumentation_data).and_wrap_original do |method, values|
+          method.call(values)
+          values.merge!(timing_data)
+        end
+      end
+
+      it 'logs with Gitaly and Rugged timing data', :aggregate_failures do
         Timecop.freeze(timestamp) do
           expect(logger).to receive(:info).with(start_payload).ordered
           expect(logger).to receive(:info).with(expected_end_payload).ordered
 
-          call_subject(job, 'test_queue') do
-            job.merge!(timing_data)
-          end
+          call_subject(job, 'test_queue') { }
         end
       end
     end

spec/lib/gitlab/sidekiq_middleware/instrumentation_logger_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::SidekiqMiddleware::InstrumentationLogger do
+  let(:job) { { 'jid' => 123 } }
+  let(:queue) { 'test_queue' }
+  let(:worker) do
+    Class.new do
+      def self.name
+        'TestDWorker'
+      end
+
+      include ApplicationWorker
+
+      def perform(*args)
+      end
+    end
+  end
+
+  subject { described_class.new }
+
+  before do
+    stub_const('TestWorker', worker)
+  end
+
+  describe '.keys' do
+    it 'returns all available payload keys' do
+      expected_keys = [
+        :cpu_s,
+        :gitaly_calls,
+        :gitaly_duration_s,
+        :rugged_calls,
+        :rugged_duration_s,
+        :elasticsearch_calls,
+        :elasticsearch_duration_s,
+        :elasticsearch_timed_out_count,
+        :worker_data_consistency,
+        :mem_objects,
+        :mem_bytes,
+        :mem_mallocs,
+        :redis_calls,
+        :redis_duration_s,
+        :redis_read_bytes,
+        :redis_write_bytes,
+        :redis_action_cable_calls,
+        :redis_action_cable_duration_s,
+        :redis_action_cable_read_bytes,
+        :redis_action_cable_write_bytes,
+        :redis_cache_calls,
+        :redis_cache_duration_s,
+        :redis_cache_read_bytes,
+        :redis_cache_write_bytes,
+        :redis_queues_calls,
+        :redis_queues_duration_s,
+        :redis_queues_read_bytes,
+        :redis_queues_write_bytes,
+        :redis_shared_state_calls,
+        :redis_shared_state_duration_s,
+        :redis_shared_state_read_bytes,
+        :redis_shared_state_write_bytes,
+        :db_count,
+        :db_write_count,
+        :db_cached_count,
+        :external_http_count,
+        :external_http_duration_s,
+        :rack_attack_redis_count,
+        :rack_attack_redis_duration_s
+      ]
+
+      expect(described_class.keys).to include(*expected_keys)
+    end
+  end
+
+  describe '#call', :request_store do
+    let(:instrumentation_values) do
+      {
+        cpu_s: 10,
+        unknown_attribute: 123,
+        db_count: 0,
+        db_cached_count: 0,
+        db_write_count: 0,
+        gitaly_calls: 0,
+        redis_calls: 0
+      }
+    end
+
+    before do
+      allow(::Gitlab::InstrumentationHelper).to receive(:add_instrumentation_data) do |values|
+        values.merge!(instrumentation_values)
+      end
+    end
+
+    it 'merges correct instrumentation data in the job' do
+      expect { |b| subject.call(worker, job, queue, &b) }.to yield_control
+
+      expected_values = instrumentation_values.except(:unknown_attribute).stringify_keys
+
+      expect(job[:instrumentation]).to eq(expected_values)
+    end
+  end
+end