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

5 files changed, 125 insertions, 102 deletions

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