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
GitHub
parent
c31c95ffe4
commit
741d0952b1
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
60
spec/models/concerns/account_counters_spec.rb
Normal file
60
spec/models/concerns/account_counters_spec.rb
Normal 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
|
||||
Loading…
Reference in New Issue
Block a user