Commit fe93987c authored by Robert Speicher's avatar Robert Speicher

Merge branch 'move_lb_calls_ar_metrics' into 'master'

Measure DB Load balancing for queries in Core

See merge request gitlab-org/gitlab!62751
parents f28ae91a 83277779
# 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
# 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
......@@ -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 = {}
{}.tap do |payload|
DB_COUNTERS.each do |counter|
payload[counter] = Gitlab::SafeRequestStore[counter].to_i
end
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
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')
......@@ -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
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment