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
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

View file

@ -3,6 +3,8 @@
module AccountCounters
extend ActiveSupport::Concern
ALLOWED_COUNTER_KEYS = %i(statuses_count following_count followers_count).freeze
included do
has_one :account_stat, inverse_of: :account
after_save :save_account_stat
@ -14,11 +16,65 @@ module AccountCounters
:following_count=,
:followers_count,
:followers_count=,
:increment_count!,
:decrement_count!,
:last_status_at,
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
super || build_account_stat
end

View file

@ -19,7 +19,7 @@ namespace :db do
task :post_migration_hook 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
Your database collation is susceptible to index corruption.
(This warning does not indicate that index corruption has occured and can be ignored)
@ -29,5 +29,11 @@ namespace :db do
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'])
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