From 62f60e86c281ae1cd0356a7630dbb76069e2ffde Mon Sep 17 00:00:00 2001
From: Eugen Rochko <eugen@zeonfederated.com>
Date: Wed, 2 Oct 2019 04:59:37 +0200
Subject: [PATCH] Fix account counters being overwritten by parallel writes
 (#12045)

---
 app/models/account_stat.rb                    | 22 ++++++++
 ...13028_add_lock_version_to_account_stats.rb | 15 ++++++
 db/schema.rb                                  |  3 +-
 spec/models/account_stat_spec.rb              | 53 +++++++++++++++++++
 4 files changed, 92 insertions(+), 1 deletion(-)
 create mode 100644 db/migrate/20191001213028_add_lock_version_to_account_stats.rb

diff --git a/app/models/account_stat.rb b/app/models/account_stat.rb
index 1351f7d8a4d..c84e4217c8c 100644
--- a/app/models/account_stat.rb
+++ b/app/models/account_stat.rb
@@ -11,6 +11,7 @@
 #  created_at      :datetime         not null
 #  updated_at      :datetime         not null
 #  last_status_at  :datetime
+#  lock_version    :integer          default(0), not null
 #
 
 class AccountStat < ApplicationRecord
@@ -20,10 +21,26 @@ class AccountStat < ApplicationRecord
 
   def increment_count!(key)
     update(attributes_for_increment(key))
+  rescue ActiveRecord::StaleObjectError
+    begin
+      reload_with_id
+    rescue ActiveRecord::RecordNotFound
+      # Nothing to do
+    else
+      retry
+    end
   end
 
   def decrement_count!(key)
     update(key => [public_send(key) - 1, 0].max)
+  rescue ActiveRecord::StaleObjectError
+    begin
+      reload_with_id
+    rescue ActiveRecord::RecordNotFound
+      # Nothing to do
+    else
+      retry
+    end
   end
 
   private
@@ -33,4 +50,9 @@ class AccountStat < ApplicationRecord
     attrs[:last_status_at] = Time.now.utc if key == :statuses_count
     attrs
   end
+
+  def reload_with_id
+    self.id = find_by!(account: account).id if new_record?
+    reload
+  end
 end
diff --git a/db/migrate/20191001213028_add_lock_version_to_account_stats.rb b/db/migrate/20191001213028_add_lock_version_to_account_stats.rb
new file mode 100644
index 00000000000..47f37cca29d
--- /dev/null
+++ b/db/migrate/20191001213028_add_lock_version_to_account_stats.rb
@@ -0,0 +1,15 @@
+require Rails.root.join('lib', 'mastodon', 'migration_helpers')
+
+class AddLockVersionToAccountStats < ActiveRecord::Migration[5.2]
+  include Mastodon::MigrationHelpers
+
+  disable_ddl_transaction!
+
+  def up
+    safety_assured { add_column_with_default :account_stats, :lock_version, :integer, allow_null: false, default: 0 }
+  end
+
+  def down
+    remove_column :account_stats, :lock_version
+  end
+end
diff --git a/db/schema.rb b/db/schema.rb
index 557b777e01f..2d516965c5f 100644
--- a/db/schema.rb
+++ b/db/schema.rb
@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema.define(version: 2019_09_27_232842) do
+ActiveRecord::Schema.define(version: 2019_10_01_213028) do
 
   # These are extensions that must be enabled in order to support this database
   enable_extension "plpgsql"
@@ -97,6 +97,7 @@ ActiveRecord::Schema.define(version: 2019_09_27_232842) do
     t.datetime "created_at", null: false
     t.datetime "updated_at", null: false
     t.datetime "last_status_at"
+    t.integer "lock_version", default: 0, null: false
     t.index ["account_id"], name: "index_account_stats_on_account_id", unique: true
   end
 
diff --git a/spec/models/account_stat_spec.rb b/spec/models/account_stat_spec.rb
index a94185109c7..8adc0d1d638 100644
--- a/spec/models/account_stat_spec.rb
+++ b/spec/models/account_stat_spec.rb
@@ -1,4 +1,57 @@
 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