diff --git a/ee/lib/ee/gitlab/metrics/subscribers/active_record.rb b/ee/lib/ee/gitlab/metrics/subscribers/active_record.rb deleted file mode 100644 index 2258d5af0bc60f48036486a7ef131dc5de9b6362..0000000000000000000000000000000000000000 --- a/ee/lib/ee/gitlab/metrics/subscribers/active_record.rb +++ /dev/null @@ -1,78 +0,0 @@ -# frozen_string_literal: true - -module EE - module Gitlab - module Metrics - module Subscribers - module ActiveRecord - extend ActiveSupport::Concern - extend ::Gitlab::Utils::Override - - DB_LOAD_BALANCING_COUNTERS = %i{ - 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 - - override :db_counter_payload - def db_counter_payload - super.tap do |payload| - if ::Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable? - DB_LOAD_BALANCING_COUNTERS.each do |counter| - payload[counter] = ::Gitlab::SafeRequestStore[counter].to_i - end - DB_LOAD_BALANCING_DURATIONS.each do |duration| - payload[duration] = ::Gitlab::SafeRequestStore[duration].to_f.round(3) - end - end - end - end - end - - override :sql - def sql(event) - super - - return unless ::Gitlab::Database::LoadBalancing.enable? - - payload = event.payload - return if ignored_query?(payload) - - db_role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(payload[:connection]) - return if db_role.blank? - - increment_db_role_counters(db_role, payload) - observe_db_role_duration(db_role, event) - end - - 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) - observe("gitlab_sql_#{db_role}_duration_seconds".to_sym, event) do - buckets ::Gitlab::Metrics::Subscribers::ActiveRecord::SQL_DURATION_BUCKET - end - - duration = event.duration / 1000.0 - duration_key = "db_#{db_role}_duration_s".to_sym - ::Gitlab::SafeRequestStore[duration_key] = (::Gitlab::SafeRequestStore[duration_key].presence || 0) + duration - end - end - end - end - end -end diff --git a/ee/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb b/ee/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb deleted file mode 100644 index e74f56aad00f81cce2d3d3cdf14f1564b4b2a065..0000000000000000000000000000000000000000 --- a/ee/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb +++ /dev/null @@ -1,144 +0,0 @@ -# frozen_string_literal: true - -require 'spec_helper' - -RSpec.describe ::Gitlab::Metrics::Subscribers::ActiveRecord do - using RSpec::Parameterized::TableSyntax - - let(:env) { {} } - let(:subscriber) { described_class.new } - let(:connection) { double(:connection) } - let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10', connection: connection } } - - let(:event) do - double( - :event, - name: 'sql.active_record', - duration: 2, - payload: payload - ) - end - - # Emulate Marginalia pre-pending comments - def sql(query, comments: true) - if comments && !%w[BEGIN COMMIT].include?(query) - "/*application:web,controller:badges,action:pipeline,correlation_id:01EYN39K9VMJC56Z7808N7RSRH*/ #{query}" - else - query - end - end - - shared_examples 'track sql events for each role' do - 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 - let(:payload) { { name: name, sql: sql(sql_query, comments: comments), connection: connection } } - - before do - allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true) - end - - context 'query using a connection to a replica' do - before do - allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(:replica) - end - - it 'queries connection db role' do - subscriber.sql(event) - - if record_query - expect(Gitlab::Database::LoadBalancing).to have_received(:db_role_for_connection).with(connection) - end - end - - it_behaves_like 'record ActiveRecord metrics', :replica - it_behaves_like 'store ActiveRecord info in RequestStore', :replica - end - - context 'query using a connection to a primary' do - before do - allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(:primary) - end - - it 'queries connection db role' do - subscriber.sql(event) - - if record_query - expect(Gitlab::Database::LoadBalancing).to have_received(:db_role_for_connection).with(connection) - end - end - - it_behaves_like 'record ActiveRecord metrics', :primary - it_behaves_like 'store ActiveRecord info in RequestStore', :primary - end - - context 'query using a connection to an unknown source' do - let(:transaction) { double('Gitlab::Metrics::WebTransaction') } - - before do - allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(nil) - - allow(::Gitlab::Metrics::WebTransaction).to receive(:current).and_return(transaction) - allow(::Gitlab::Metrics::BackgroundTransaction).to receive(:current).and_return(nil) - - allow(transaction).to receive(:increment) - allow(transaction).to receive(:observe) - end - - it 'does not record DB role metrics' do - expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_primary_count_total".to_sym, any_args) - expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_replica_count_total".to_sym, any_args) - - expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_primary_cached_count_total".to_sym, any_args) - expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_replica_cached_count_total".to_sym, any_args) - - expect(transaction).not_to receive(:observe).with("gitlab_sql_primary_duration_seconds".to_sym, any_args) - expect(transaction).not_to receive(:observe).with("gitlab_sql_replica_duration_seconds".to_sym, any_args) - - subscriber.sql(event) - end - - it 'does not store DB roles into into RequestStore' do - Gitlab::WithRequestStore.with_request_store do - subscriber.sql(event) - - expect(described_class.db_counter_payload).to include( - db_primary_cached_count: 0, - db_primary_count: 0, - db_primary_duration_s: 0, - db_replica_cached_count: 0, - db_replica_count: 0, - db_replica_duration_s: 0 - ) - end - end - end - end - end - - context 'without Marginalia comments' do - let(:comments) { false } - - it_behaves_like 'track sql events for each role' - end - - context 'with Marginalia comments' do - let(:comments) { true } - - it_behaves_like 'track sql events for each role' - end -end diff --git a/lib/gitlab/metrics/subscribers/active_record.rb b/lib/gitlab/metrics/subscribers/active_record.rb index 24006d75c4d6d9bedd13fae7b66935e8998d91b2..9f7884e13648e734619f881b1c0663bb6b960038 100644 --- a/lib/gitlab/metrics/subscribers/active_record.rb +++ b/lib/gitlab/metrics/subscribers/active_record.rb @@ -14,6 +14,14 @@ module Gitlab SQL_DURATION_BUCKET = [0.05, 0.1, 0.25].freeze TRANSACTION_DURATION_BUCKET = [0.1, 0.25, 1].freeze + DB_LOAD_BALANCING_COUNTERS = %i{ + 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 + # This event is published from ActiveRecordBaseTransactionMetrics and # used to record a database transaction duration when calling # ActiveRecord::Base.transaction {} block. @@ -39,20 +47,57 @@ module Gitlab observe(:gitlab_sql_duration_seconds, event) do buckets SQL_DURATION_BUCKET end + + if ::Gitlab::Database::LoadBalancing.enable? + db_role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(payload[:connection]) + return if db_role.blank? + + increment_db_role_counters(db_role, payload) + observe_db_role_duration(db_role, event) + end end def self.db_counter_payload return {} unless Gitlab::SafeRequestStore.active? - payload = {} - DB_COUNTERS.each do |counter| - payload[counter] = Gitlab::SafeRequestStore[counter].to_i + {}.tap do |payload| + DB_COUNTERS.each do |counter| + payload[counter] = Gitlab::SafeRequestStore[counter].to_i + end + + if ::Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable? + DB_LOAD_BALANCING_COUNTERS.each do |counter| + payload[counter] = ::Gitlab::SafeRequestStore[counter].to_i + end + DB_LOAD_BALANCING_DURATIONS.each do |duration| + payload[duration] = ::Gitlab::SafeRequestStore[duration].to_f.round(3) + end + end end - payload end 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) + observe("gitlab_sql_#{db_role}_duration_seconds".to_sym, event) do + buckets ::Gitlab::Metrics::Subscribers::ActiveRecord::SQL_DURATION_BUCKET + end + + duration = event.duration / 1000.0 + duration_key = "db_#{db_role}_duration_s".to_sym + ::Gitlab::SafeRequestStore[duration_key] = (::Gitlab::SafeRequestStore[duration_key].presence || 0) + duration + end + def ignored_query?(payload) payload[:name] == 'SCHEMA' || IGNORABLE_SQL.include?(payload[:sql]) end @@ -82,5 +127,3 @@ module Gitlab end end end - -Gitlab::Metrics::Subscribers::ActiveRecord.prepend_mod_with('Gitlab::Metrics::Subscribers::ActiveRecord') diff --git a/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb index 6bfcfa2128956563400ab35b23905bb6c840d196..cffa62c3a5290d5b0e94328a14a928cee19f7996 100644 --- a/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb +++ b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb @@ -150,4 +150,140 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do it_behaves_like 'track generic sql events' end end + + context 'Database Load Balancing enabled' do + let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10', connection: connection } } + + let(:event) do + double( + :event, + name: 'sql.active_record', + duration: 2, + payload: payload + ) + end + + # Emulate Marginalia pre-pending comments + def sql(query, comments: true) + if comments && !%w[BEGIN COMMIT].include?(query) + "/*application:web,controller:badges,action:pipeline,correlation_id:01EYN39K9VMJC56Z7808N7RSRH*/ #{query}" + else + query + end + end + + shared_examples 'track sql events for each role' do + 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 + let(:payload) { { name: name, sql: sql(sql_query, comments: comments), connection: connection } } + + before do + allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true) + end + + context 'query using a connection to a replica' do + before do + allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(:replica) + end + + it 'queries connection db role' do + subscriber.sql(event) + + if record_query + expect(Gitlab::Database::LoadBalancing).to have_received(:db_role_for_connection).with(connection) + end + end + + it_behaves_like 'record ActiveRecord metrics', :replica + it_behaves_like 'store ActiveRecord info in RequestStore', :replica + end + + context 'query using a connection to a primary' do + before do + allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(:primary) + end + + it 'queries connection db role' do + subscriber.sql(event) + + if record_query + expect(Gitlab::Database::LoadBalancing).to have_received(:db_role_for_connection).with(connection) + end + end + + it_behaves_like 'record ActiveRecord metrics', :primary + it_behaves_like 'store ActiveRecord info in RequestStore', :primary + end + + context 'query using a connection to an unknown source' do + let(:transaction) { double('Gitlab::Metrics::WebTransaction') } + + before do + allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(nil) + + allow(::Gitlab::Metrics::WebTransaction).to receive(:current).and_return(transaction) + allow(::Gitlab::Metrics::BackgroundTransaction).to receive(:current).and_return(nil) + + allow(transaction).to receive(:increment) + allow(transaction).to receive(:observe) + end + + it 'does not record DB role metrics' do + expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_primary_count_total".to_sym, any_args) + expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_replica_count_total".to_sym, any_args) + + expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_primary_cached_count_total".to_sym, any_args) + expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_replica_cached_count_total".to_sym, any_args) + + expect(transaction).not_to receive(:observe).with("gitlab_sql_primary_duration_seconds".to_sym, any_args) + expect(transaction).not_to receive(:observe).with("gitlab_sql_replica_duration_seconds".to_sym, any_args) + + subscriber.sql(event) + end + + it 'does not store DB roles into into RequestStore' do + Gitlab::WithRequestStore.with_request_store do + subscriber.sql(event) + + expect(described_class.db_counter_payload).to include( + db_primary_cached_count: 0, + db_primary_count: 0, + db_primary_duration_s: 0, + db_replica_cached_count: 0, + db_replica_count: 0, + db_replica_duration_s: 0 + ) + end + end + end + end + end + + context 'without Marginalia comments' do + let(:comments) { false } + + it_behaves_like 'track sql events for each role' + end + + context 'with Marginalia comments' do + let(:comments) { true } + + it_behaves_like 'track sql events for each role' + end + end end