Improve account counters handling (mastodon#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

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

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

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

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