From 741d0952b174740e70a09fe6db6862624dfe1e44 Mon Sep 17 00:00:00 2001 From: Claire Date: Fri, 19 Mar 2021 13:14:57 +0100 Subject: [PATCH] Improve account counters handling (#15913) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 --- app/models/account_stat.rb | 42 ------------- app/models/concerns/account_counters.rb | 60 ++++++++++++++++++- lib/tasks/db.rake | 8 ++- spec/models/account_stat_spec.rb | 57 ------------------ spec/models/concerns/account_counters_spec.rb | 60 +++++++++++++++++++ 5 files changed, 125 insertions(+), 102 deletions(-) delete mode 100644 spec/models/account_stat_spec.rb create mode 100644 spec/models/concerns/account_counters_spec.rb diff --git a/app/models/account_stat.rb b/app/models/account_stat.rb index e70b54d79..a826a9af3 100644 --- a/app/models/account_stat.rb +++ b/app/models/account_stat.rb @@ -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 diff --git a/app/models/concerns/account_counters.rb b/app/models/concerns/account_counters.rb index 6e25e1905..fd3f161ad 100644 --- a/app/models/concerns/account_counters.rb +++ b/app/models/concerns/account_counters.rb @@ -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 diff --git a/lib/tasks/db.rake b/lib/tasks/db.rake index 552a02b3f..7e6c1c8fc 100644 --- a/lib/tasks/db.rake +++ b/lib/tasks/db.rake @@ -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 diff --git a/spec/models/account_stat_spec.rb b/spec/models/account_stat_spec.rb deleted file mode 100644 index 8adc0d1d6..000000000 --- a/spec/models/account_stat_spec.rb +++ /dev/null @@ -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 diff --git a/spec/models/concerns/account_counters_spec.rb b/spec/models/concerns/account_counters_spec.rb new file mode 100644 index 000000000..4350496e7 --- /dev/null +++ b/spec/models/concerns/account_counters_spec.rb @@ -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