Add metric to track querying write location on Request level

ee/changelogs/unreleased/326225-add-load-balancing-metrics-on-request-level.yml

+---
+title: Add metric to track querying write ahead log on Request level
+merge_request: 58673
+author:
+type: other

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

@@ -9,11 +9,13 @@ module EE
           extend ::Gitlab::Utils::Override
 
           DB_LOAD_BALANCING_COUNTERS = %i{
-            db_replica_count db_replica_cached_count
-            db_primary_count db_primary_cached_count
+            db_replica_count db_replica_cached_count db_replica_wal_count
+            db_primary_count db_primary_cached_count db_primary_wal_count
           }.freeze
           DB_LOAD_BALANCING_DURATIONS = %i{db_primary_duration_s db_replica_duration_s}.freeze
 
+          SQL_WAL_LOCATION_REGEX = /(pg_current_wal_insert_lsn\(\)::text|pg_last_wal_replay_lsn\(\)::text)/.freeze
+
           class_methods do
             extend ::Gitlab::Utils::Override
 
@@ -55,9 +57,14 @@ module EE
 
           private
 
+          def wal_command?(payload)
+            payload[:sql].match(SQL_WAL_LOCATION_REGEX)
+          end
+
           def increment_db_role_counters(db_role, payload)
             increment("db_#{db_role}_count".to_sym)
             increment("db_#{db_role}_cached_count".to_sym) if cached_query?(payload)
+            increment("db_#{db_role}_wal_count".to_sym) if !cached_query?(payload) && wal_command?(payload)
           end
 
           def observe_db_role_duration(db_role, event)

ee/spec/initializers/lograge_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'lograge', type: :request do
+  context 'with a log subscriber' do
+    include_context 'parsed logs'
+    include_context 'clear DB Load Balancing configuration'
+
+    let(:subscriber) { Lograge::LogSubscribers::ActionController.new }
+
+    let(:event) do
+      ActiveSupport::Notifications::Event.new(
+        'process_action.action_controller',
+        Time.now,
+        Time.now,
+        2,
+        status: 200,
+        controller: 'HomeController',
+        action: 'index',
+        format: 'application/json',
+        method: 'GET',
+        path: '/home?foo=bar',
+        params: {},
+        db_runtime: 0.02,
+        view_runtime: 0.01
+      )
+    end
+
+    let(:logging_keys) do
+      %w[db_primary_wal_count
+         db_replica_wal_count
+         db_replica_count
+         db_replica_cached_count
+         db_primary_count
+         db_primary_cached_count
+         db_primary_duration_s
+         db_replica_duration_s]
+    end
+
+    context 'when load balancing is enabled' do
+      before do
+        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+      end
+
+      context 'with db payload' do
+        context 'when RequestStore is enabled', :request_store do
+          it 'includes db counters' do
+            subscriber.process_action(event)
+            expect(log_data).to include(*logging_keys)
+          end
+        end
+
+        context 'when RequestStore is disabled' do
+          it 'does not include db counters' do
+            subscriber.process_action(event)
+
+            expect(log_data).not_to include(*logging_keys)
+          end
+        end
+      end
+    end
+
+    context 'when load balancing is disabled' do
+      before do
+        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
+      end
+
+      it 'does not include db counters' do
+        subscriber.process_action(event)
+
+        expect(log_data).not_to include(*logging_keys)
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/instrumentation_helper_spec.rb

@@ -21,7 +21,9 @@ RSpec.describe Gitlab::InstrumentationHelper do
         expect(payload).to include(db_replica_count: 0,
                                    db_replica_cached_count: 0,
                                    db_primary_count: 0,
-                                   db_primary_cached_count: 0)
+                                   db_primary_cached_count: 0,
+                                   db_primary_wal_count: 0,
+                                   db_replica_wal_count: 0)
       end
     end
 
@@ -30,13 +32,15 @@ RSpec.describe Gitlab::InstrumentationHelper do
         allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
       end
 
-      it 'includes DB counts' do
+      it 'does not include DB counts' do
         subject
 
         expect(payload).not_to include(db_replica_count: 0,
                                    db_replica_cached_count: 0,
                                    db_primary_count: 0,
-                                   db_primary_cached_count: 0)
+                                   db_primary_cached_count: 0,
+                                   db_primary_wal_count: 0,
+                                   db_replica_wal_count: 0)
       end
     end
 

ee/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb

@@ -29,18 +29,20 @@ RSpec.describe ::Gitlab::Metrics::Subscribers::ActiveRecord do
   end
 
   shared_examples 'track sql events for each role' do
-    where(:name, :sql_query, :record_query, :record_write_query, :record_cached_query) do
-      'SQL' | 'SELECT * FROM users WHERE id = 10' | true | false | false
-      'SQL' | 'WITH active_milestones AS (SELECT COUNT(*), state FROM milestones GROUP BY state) SELECT * FROM active_milestones' | true | false | false
-      'SQL' | 'SELECT * FROM users WHERE id = 10 FOR UPDATE' | true | true | false
-      'SQL' | 'WITH archived_rows AS (SELECT * FROM users WHERE archived = true) INSERT INTO products_log SELECT * FROM archived_rows' | true | true | false
-      'SQL' | 'DELETE FROM users where id = 10' | true | true | false
-      'SQL' | 'INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects' | true | true | false
-      'SQL' | 'UPDATE users SET admin = true WHERE id = 10' | true | true | false
-      'CACHE' | 'SELECT * FROM users WHERE id = 10' | true | false | true
-      'SCHEMA' | "SELECT attr.attname FROM pg_attribute attr INNER JOIN pg_constraint cons ON attr.attrelid = cons.conrelid AND attr.attnum = any(cons.conkey) WHERE cons.contype = 'p' AND cons.conrelid = '\"projects\"'::regclass" | false | false | false
-      nil | 'BEGIN' | false | false | false
-      nil | 'COMMIT' | false | false | false
+    where(:name, :sql_query, :record_query, :record_write_query, :record_cached_query, :record_wal_query) do
+      'SQL' | 'SELECT * FROM users WHERE id = 10' | true | false | false | false
+      'SQL' | 'WITH active_milestones AS (SELECT COUNT(*), state FROM milestones GROUP BY state) SELECT * FROM active_milestones' | true | false | false | false
+      'SQL' | 'SELECT * FROM users WHERE id = 10 FOR UPDATE' | true | true | false | false
+      'SQL' | 'WITH archived_rows AS (SELECT * FROM users WHERE archived = true) INSERT INTO products_log SELECT * FROM archived_rows' | true | true | false | false
+      'SQL' | 'DELETE FROM users where id = 10' | true | true | false | false
+      'SQL' | 'INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects' | true | true | false | false
+      'SQL' | 'UPDATE users SET admin = true WHERE id = 10' | true | true | false | false
+      'SQL' | 'SELECT pg_current_wal_insert_lsn()::text AS location' | true | false | false | true
+      'SQL' | 'SELECT pg_last_wal_replay_lsn()::text AS location' | true | false | false | true
+      'CACHE' | 'SELECT * FROM users WHERE id = 10' | true | false | true | false
+      'SCHEMA' | "SELECT attr.attname FROM pg_attribute attr INNER JOIN pg_constraint cons ON attr.attrelid = cons.conrelid AND attr.attnum = any(cons.conkey) WHERE cons.contype = 'p' AND cons.conrelid = '\"projects\"'::regclass" | false | false | false | false
+      nil | 'BEGIN' | false | false | false | false
+      nil | 'COMMIT' | false | false | false | false
     end
 
     with_them do

ee/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb

@@ -76,7 +76,9 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
             'db_replica_count' => 0,
             'db_replica_cached_count' => 0,
             'db_primary_count' => a_value >= 1,
-            'db_primary_cached_count' => 0
+            'db_primary_cached_count' => 0,
+            'db_primary_wal_count' => 0,
+            'db_replica_wal_count' => 0
           )
         end
 
@@ -94,7 +96,9 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
             'db_replica_count' => 0,
             'db_replica_cached_count' => 0,
             'db_primary_count' => 0,
-            'db_primary_cached_count' => 0
+            'db_primary_cached_count' => 0,
+            'db_primary_wal_count' => 0,
+            'db_replica_wal_count' => 0
           )
         end
 

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

@@ -9,7 +9,9 @@ RSpec.describe Gitlab::SidekiqMiddleware::InstrumentationLogger do
         :db_replica_count,
         :db_replica_cached_count,
         :db_primary_count,
-        :db_primary_cached_count
+        :db_primary_cached_count,
+        :db_primary_wal_count,
+        :db_replica_wal_count
       ]
 
       expect(described_class.keys).to include(*expected_keys)

spec/lib/gitlab/metrics/subscribers/active_record_spec.rb

@@ -124,6 +124,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
 
       with_them do
         let(:payload) { { name: name, sql: sql(sql_query, comments: comments), connection: connection } }
+        let(:record_wal_query) { false }
 
         it 'marks the current thread as using the database' do
           # since it would already have been toggled by other specs

spec/support/shared_examples/metrics/active_record_subscriber_shared_examples.rb

@@ -16,7 +16,9 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
             db_primary_duration_s:  record_query ? 0.002 : 0,
             db_replica_cached_count:  0,
             db_replica_count:  0,
-            db_replica_duration_s:  0.0
+            db_replica_duration_s:  0.0,
+            db_primary_wal_count: record_wal_query ? 1 : 0,
+            db_replica_wal_count: 0
           )
         elsif db_role == :replica
           expect(described_class.db_counter_payload).to eq(
@@ -28,7 +30,9 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
             db_primary_duration_s:  0.0,
             db_replica_cached_count:  record_cached_query ? 1 : 0,
             db_replica_count:  record_query ? 1 : 0,
-            db_replica_duration_s:  record_query ? 0.002 : 0
+            db_replica_duration_s:  record_query ? 0.002 : 0,
+            db_replica_wal_count: record_wal_query ? 1 : 0,
+            db_primary_wal_count: 0
           )
         else
           expect(described_class.db_counter_payload).to eq(
@@ -66,6 +70,12 @@ RSpec.shared_examples 'record ActiveRecord metrics in a metrics transaction' do
       expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_#{db_role}_cached_count_total".to_sym, 1) if db_role
     end
 
+    if record_wal_query
+      expect(transaction).to receive(:increment).with("gitlab_transaction_db_#{db_role}_wal_count_total".to_sym, 1) if db_role
+    else
+      expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_#{db_role}_wal_count_total".to_sym, 1) if db_role
+    end
+
     subscriber.sql(event)
   end