Improve account counters handling (#15913)

* Improve account counters handling

* Use ActiveRecord::Base::sanitize_sql to pass values instead of interpolating them

Keep using string interpolation for `key` as it is safe and using
“ActiveRecord::Base::sanitize_sql_hash_for_assignment” would require stitching
bits of SQL in a way that is not more easily checked for safety.

* Add migration hook to catch PostgreSQL versions earlier than 9.5
This commit is contained in:
Claire 2021-03-19 13:14:57 +01:00 committed by GitHub
parent c31c95ffe4
commit 741d0952b1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 125 additions and 102 deletions

View File

@ -18,46 +18,4 @@ class AccountStat < ApplicationRecord
belongs_to :account, inverse_of: :account_stat belongs_to :account, inverse_of: :account_stat
update_index('accounts#account', :account) update_index('accounts#account', :account)
def increment_count!(key)
update(attributes_for_increment(key))
rescue ActiveRecord::StaleObjectError, ActiveRecord::RecordNotUnique
begin
reload_with_id
rescue ActiveRecord::RecordNotFound
return
end
retry
end
def decrement_count!(key)
update(attributes_for_decrement(key))
rescue ActiveRecord::StaleObjectError, ActiveRecord::RecordNotUnique
begin
reload_with_id
rescue ActiveRecord::RecordNotFound
return
end
retry
end
private
def attributes_for_increment(key)
attrs = { key => public_send(key) + 1 }
attrs[:last_status_at] = Time.now.utc if key == :statuses_count
attrs
end
def attributes_for_decrement(key)
attrs = { key => [public_send(key) - 1, 0].max }
attrs
end
def reload_with_id
self.id = self.class.find_by!(account: account).id if new_record?
reload
end
end end

View File

@ -3,6 +3,8 @@
module AccountCounters module AccountCounters
extend ActiveSupport::Concern extend ActiveSupport::Concern
ALLOWED_COUNTER_KEYS = %i(statuses_count following_count followers_count).freeze
included do included do
has_one :account_stat, inverse_of: :account has_one :account_stat, inverse_of: :account
after_save :save_account_stat after_save :save_account_stat
@ -14,11 +16,65 @@ module AccountCounters
:following_count=, :following_count=,
:followers_count, :followers_count,
:followers_count=, :followers_count=,
:increment_count!,
:decrement_count!,
:last_status_at, :last_status_at,
to: :account_stat to: :account_stat
# @param [Symbol] key
def increment_count!(key)
update_count!(key, 1)
end
# @param [Symbol] key
def decrement_count!(key)
update_count!(key, -1)
end
# @param [Symbol] key
# @param [Integer] value
def update_count!(key, value)
raise ArgumentError, "Invalid key #{key}" unless ALLOWED_COUNTER_KEYS.include?(key)
raise ArgumentError, 'Do not call update_count! on dirty objects' if association(:account_stat).loaded? && account_stat&.changed? && account_stat.changed_attribute_names_to_save == %w(id)
value = value.to_i
default_value = value.positive? ? value : 0
# We do an upsert using manually written SQL, as Rails' upsert method does
# not seem to support writing expressions in the UPDATE clause, but only
# re-insert the provided values instead.
# Even ARel seem to be missing proper handling of upserts.
sql = if value.positive? && key == :statuses_count
<<-SQL.squish
INSERT INTO account_stats(account_id, #{key}, created_at, updated_at, last_status_at)
VALUES (:account_id, :default_value, now(), now(), now())
ON CONFLICT (account_id) DO UPDATE
SET #{key} = account_stats.#{key} + :value,
last_status_at = now(),
lock_version = account_stats.lock_version + 1,
updated_at = now()
RETURNING id;
SQL
else
<<-SQL.squish
INSERT INTO account_stats(account_id, #{key}, created_at, updated_at)
VALUES (:account_id, :default_value, now(), now())
ON CONFLICT (account_id) DO UPDATE
SET #{key} = account_stats.#{key} + :value,
lock_version = account_stats.lock_version + 1,
updated_at = now()
RETURNING id;
SQL
end
sql = AccountStat.sanitize_sql([sql, account_id: id, default_value: default_value, value: value])
account_stat_id = AccountStat.connection.exec_query(sql)[0]['id']
# Reload account_stat if it was loaded, taking into account newly-created unsaved records
if association(:account_stat).loaded?
account_stat.id = account_stat_id if account_stat.new_record?
account_stat.reload
end
end
def account_stat def account_stat
super || build_account_stat super || build_account_stat
end end

View File

@ -19,7 +19,7 @@ namespace :db do
task :post_migration_hook do task :post_migration_hook do
at_exit do at_exit do
unless %w(C POSIX).include?(ActiveRecord::Base.connection.execute('SELECT datcollate FROM pg_database WHERE datname = current_database();').first['datcollate']) unless %w(C POSIX).include?(ActiveRecord::Base.connection.select_one('SELECT datcollate FROM pg_database WHERE datname = current_database();')['datcollate'])
warn <<~WARNING warn <<~WARNING
Your database collation is susceptible to index corruption. Your database collation is susceptible to index corruption.
(This warning does not indicate that index corruption has occured and can be ignored) (This warning does not indicate that index corruption has occured and can be ignored)
@ -29,5 +29,11 @@ namespace :db do
end end
end end
task :pre_migration_check do
version = ActiveRecord::Base.connection.select_one("SELECT current_setting('server_version_num') AS v")['v'].to_i
abort 'ERROR: This version of Mastodon requires PostgreSQL 9.5 or newer. Please update PostgreSQL before updating Mastodon.' if version < 90_500
end
Rake::Task['db:migrate'].enhance(['db:pre_migration_check'])
Rake::Task['db:migrate'].enhance(['db:post_migration_hook']) Rake::Task['db:migrate'].enhance(['db:post_migration_hook'])
end end

View File

@ -1,57 +0,0 @@
require 'rails_helper'
RSpec.describe AccountStat, type: :model do
describe '#increment_count!' do
it 'increments the count' do
account_stat = AccountStat.create(account: Fabricate(:account))
expect(account_stat.followers_count).to eq 0
account_stat.increment_count!(:followers_count)
expect(account_stat.followers_count).to eq 1
end
it 'increments the count in multi-threaded an environment' do
account_stat = AccountStat.create(account: Fabricate(:account), statuses_count: 0)
increment_by = 15
wait_for_start = true
threads = Array.new(increment_by) do
Thread.new do
true while wait_for_start
AccountStat.find(account_stat.id).increment_count!(:statuses_count)
end
end
wait_for_start = false
threads.each(&:join)
expect(account_stat.reload.statuses_count).to eq increment_by
end
end
describe '#decrement_count!' do
it 'decrements the count' do
account_stat = AccountStat.create(account: Fabricate(:account), followers_count: 15)
expect(account_stat.followers_count).to eq 15
account_stat.decrement_count!(:followers_count)
expect(account_stat.followers_count).to eq 14
end
it 'decrements the count in multi-threaded an environment' do
account_stat = AccountStat.create(account: Fabricate(:account), statuses_count: 15)
decrement_by = 10
wait_for_start = true
threads = Array.new(decrement_by) do
Thread.new do
true while wait_for_start
AccountStat.find(account_stat.id).decrement_count!(:statuses_count)
end
end
wait_for_start = false
threads.each(&:join)
expect(account_stat.reload.statuses_count).to eq 5
end
end
end

View File

@ -0,0 +1,60 @@
require 'rails_helper'
describe AccountCounters do
let!(:account) { Fabricate(:account) }
describe '#increment_count!' do
it 'increments the count' do
expect(account.followers_count).to eq 0
account.increment_count!(:followers_count)
expect(account.followers_count).to eq 1
end
it 'increments the count in multi-threaded an environment' do
increment_by = 15
wait_for_start = true
threads = Array.new(increment_by) do
Thread.new do
true while wait_for_start
account.increment_count!(:statuses_count)
end
end
wait_for_start = false
threads.each(&:join)
expect(account.statuses_count).to eq increment_by
end
end
describe '#decrement_count!' do
it 'decrements the count' do
account.followers_count = 15
account.save!
expect(account.followers_count).to eq 15
account.decrement_count!(:followers_count)
expect(account.followers_count).to eq 14
end
it 'decrements the count in multi-threaded an environment' do
decrement_by = 10
wait_for_start = true
account.statuses_count = 15
account.save!
threads = Array.new(decrement_by) do
Thread.new do
true while wait_for_start
account.decrement_count!(:statuses_count)
end
end
wait_for_start = false
threads.each(&:join)
expect(account.statuses_count).to eq 5
end
end
end